Skip to content

Commit de56b35

Browse files
kvchandrew-farries
andauthored
New option for columns: generated (#605)
This PR adds support for generated columns in `create_table` and `add_column` operations. The column definition is extended with the setting named `generated`. Define a stored column: ```json ... "generated": { "expression": "upper(name)" } ``` Define a generated identity column: ```json ... "generated": { "identity": { "user_specified_values": "ALWAYS", "sequence_options": "start 5" } ``` Identity columns are marked not null by default, so setting `nullable` is unnecessary. I changed the existing `fruits` table to use a generated identity column as a primary key. ```json { "name": "08_create_fruits_table", "operations": [ { "create_table": { "name": "fruits", "columns": [ { "name": "id", "type": "bigint", "pk": true, "generated": { "identity": { "user_specified_values": "BY DEFAULT" } } }, { "name": "name", "type": "varchar(255)", "unique": true }, { "name": "price", "type": "decimal(10,2)" } ] } } ] } ``` The PR also contains the transformation required by `sql2pgroll` project. Requires xataio/pg_query_go#7 --------- Co-authored-by: Andrew Farries <[email protected]>
1 parent fd074b5 commit de56b35

15 files changed

+348
-17
lines changed

docs/operations/create_table.mdx

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,13 @@ Each `column` is defined as:
3030
"name": "name of check constraint",
3131
"constraint": "constraint expression"
3232
},
33+
"generated": {
34+
"expression": "expression for stored column",
35+
"identity": {
36+
"user_specified_values": "user specified values can be used, can be ALWAYS and BY DEFAULT. Default is ALWAYS",
37+
"sequence_options": "sequence options for identity columns"
38+
}
39+
},
3340
"references": {
3441
"name": "name of foreign key constraint",
3542
"table": "name of referenced table",
@@ -41,6 +48,8 @@ Each `column` is defined as:
4148

4249
Default values are subject to the usual rules for quoting SQL expressions. In particular, string literals should be surrounded with single quotes.
4350

51+
Generated columns can either be stored or identity columns. Options `generated.expression` and `generated.identity` cannot be set at the same time.
52+
4453
Each `constraint` is defined as:
4554

4655
```json
@@ -82,7 +91,7 @@ Create one table:
8291

8392
### Create one table (2)
8493

85-
Create one table:
94+
Create one table with generated identity column:
8695

8796
<ExampleSnippet example="08_create_fruits_table.json" language="json" />
8897

examples/08_create_fruits_table.json

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,13 @@
77
"columns": [
88
{
99
"name": "id",
10-
"type": "serial",
11-
"pk": true
10+
"type": "bigint",
11+
"pk": true,
12+
"generated": {
13+
"identity": {
14+
"user_specified_values": "BY DEFAULT"
15+
}
16+
}
1217
},
1318
{
1419
"name": "name",

go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ require (
1414
github.com/stretchr/testify v1.10.0
1515
github.com/testcontainers/testcontainers-go v0.35.0
1616
github.com/testcontainers/testcontainers-go/modules/postgres v0.35.0
17-
github.com/xataio/pg_query_go/v6 v6.0.0-20241223083246-5033a992750b
17+
github.com/xataio/pg_query_go/v6 v6.0.0-20250122133641-54118c062181
1818
golang.org/x/tools v0.29.0
1919
)
2020

go.sum

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -221,8 +221,8 @@ github.com/tklauser/go-sysconf v0.3.14 h1:g5vzr9iPFFz24v2KZXs/pvpvh8/V9Fw6vQK5ZZ
221221
github.com/tklauser/go-sysconf v0.3.14/go.mod h1:1ym4lWMLUOhuBOPGtRcJm7tEGX4SCYNEEEtghGG/8uY=
222222
github.com/tklauser/numcpus v0.8.0 h1:Mx4Wwe/FjZLeQsK/6kt2EOepwwSl7SmJrK5bV/dXYgY=
223223
github.com/tklauser/numcpus v0.8.0/go.mod h1:ZJZlAY+dmR4eut8epnzf0u/VwodKmryxR8txiloSqBE=
224-
github.com/xataio/pg_query_go/v6 v6.0.0-20241223083246-5033a992750b h1:rZyagmQlpbFFKUMW7k/hQVklPGKAZ9Nzv0xqPPd7SQU=
225-
github.com/xataio/pg_query_go/v6 v6.0.0-20241223083246-5033a992750b/go.mod h1:GK6bpfAhPtZb7wG/IccqvnH+cz3cmvvRTkC+MosESGo=
224+
github.com/xataio/pg_query_go/v6 v6.0.0-20250122133641-54118c062181 h1:iLOHgul20WFUhO4eFpJ/lZRkHzZICF2ghzncxtOcD0E=
225+
github.com/xataio/pg_query_go/v6 v6.0.0-20250122133641-54118c062181/go.mod h1:GK6bpfAhPtZb7wG/IccqvnH+cz3cmvvRTkC+MosESGo=
226226
github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778/go.mod h1:2MuV+tbUrU1zIOPMxZ5EncGwgmMJsa+9ucAQZXxsObs=
227227
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
228228
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=

pkg/migrations/errors.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,3 +254,12 @@ type PrimaryKeysAreAlreadySetError struct {
254254
func (e PrimaryKeysAreAlreadySetError) Error() string {
255255
return fmt.Sprintf("table %q already has a primary key configuration in columns list", e.Table)
256256
}
257+
258+
type InvalidGeneratedColumnError struct {
259+
Table string
260+
Column string
261+
}
262+
263+
func (e InvalidGeneratedColumnError) Error() string {
264+
return fmt.Sprintf("column %q on table %q is invalid: only one of generated.expression and generated.identity may be set", e.Column, e.Table)
265+
}

pkg/migrations/op_add_column.go

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,11 @@ func (o *OpAddColumn) Validate(ctx context.Context, s *schema.Schema) error {
184184
}
185185
}
186186

187-
if !o.Column.IsNullable() && o.Column.Default == nil && o.Up == "" && !o.Column.HasImplicitDefault() {
187+
if o.Column.Generated != nil && o.Column.Generated.Expression != "" && o.Column.Generated.Identity != nil {
188+
return InvalidGeneratedColumnError{Table: o.Table, Column: o.Column.Name}
189+
}
190+
191+
if !o.Column.IsNullable() && o.Column.Default == nil && o.Up == "" && !o.Column.HasImplicitDefault() && o.Column.Generated == nil {
188192
return FieldRequiredError{Name: "up"}
189193
}
190194

@@ -213,6 +217,11 @@ func addColumn(ctx context.Context, conn db.DB, o OpAddColumn, t *schema.Table,
213217
o.Column.Nullable = true
214218
}
215219

220+
// Generated identity columns are marked not null automatically by PostgreSQL.
221+
if o.Column.Generated != nil && o.Column.Generated.Identity != nil && !o.Column.IsNullable() {
222+
o.Column.Nullable = true
223+
}
224+
216225
// Don't add a column with a CHECK constraint directly.
217226
// They are handled by:
218227
// - adding the column without the constraint
@@ -292,6 +301,18 @@ func (w ColumnSQLWriter) Write(col Column) (string, error) {
292301
}
293302
sql += fmt.Sprintf(" DEFAULT %s", d)
294303
}
304+
305+
if col.Generated != nil {
306+
if col.Generated.Expression != "" {
307+
sql += fmt.Sprintf(" GENERATED ALWAYS AS (%s) STORED", col.Generated.Expression)
308+
} else if col.Generated.Identity != nil {
309+
sql += fmt.Sprintf(" GENERATED %s AS IDENTITY", col.Generated.Identity.UserSpecifiedValues)
310+
if col.Generated.Identity.SequenceOptions != "" {
311+
sql += fmt.Sprintf(" (%s)", col.Generated.Identity.SequenceOptions)
312+
}
313+
}
314+
}
315+
295316
if col.References != nil {
296317
onDelete := string(ForeignKeyReferenceOnDeleteNOACTION)
297318
if col.References.OnDelete != "" {

pkg/migrations/op_add_column_test.go

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,101 @@ func TestAddColumn(t *testing.T) {
254254
}, res)
255255
},
256256
},
257+
{
258+
name: "add generated identity column and a regular stored column",
259+
migrations: []migrations.Migration{
260+
{
261+
Name: "01_add_table",
262+
Operations: migrations.Operations{
263+
&migrations.OpCreateTable{
264+
Name: "users",
265+
Columns: []migrations.Column{
266+
{
267+
Name: "id",
268+
Type: "bigint",
269+
Pk: true,
270+
Generated: &migrations.ColumnGenerated{
271+
Identity: &migrations.ColumnGeneratedIdentity{
272+
UserSpecifiedValues: "ALWAYS",
273+
},
274+
},
275+
},
276+
{
277+
Name: "name",
278+
Type: "varchar(255)",
279+
},
280+
},
281+
},
282+
},
283+
},
284+
{
285+
Name: "02_add_column",
286+
Operations: migrations.Operations{
287+
&migrations.OpAddColumn{
288+
Table: "users",
289+
Column: migrations.Column{
290+
Name: "generated_upper",
291+
Type: "varchar(255)",
292+
Generated: &migrations.ColumnGenerated{
293+
Expression: "upper(name)",
294+
},
295+
Nullable: false,
296+
},
297+
},
298+
},
299+
},
300+
},
301+
afterStart: func(t *testing.T, db *sql.DB, schema string) {
302+
// old and new views of the table should exist
303+
ViewMustExist(t, db, schema, "01_add_table", "users")
304+
ViewMustExist(t, db, schema, "02_add_column", "users")
305+
306+
// inserting via both the old and the new views works
307+
MustInsert(t, db, schema, "01_add_table", "users", map[string]string{
308+
"name": "Alice",
309+
})
310+
MustInsert(t, db, schema, "02_add_column", "users", map[string]string{
311+
"name": "Bob",
312+
})
313+
314+
// selecting from both the old and the new views works
315+
resOld := MustSelect(t, db, schema, "01_add_table", "users")
316+
assert.Equal(t, []map[string]any{
317+
{"id": 1, "name": "Alice"},
318+
{"id": 2, "name": "Bob"},
319+
}, resOld)
320+
resNew := MustSelect(t, db, schema, "02_add_column", "users")
321+
assert.Equal(t, []map[string]any{
322+
{"id": 1, "name": "Alice", "generated_upper": "ALICE"},
323+
{"id": 2, "name": "Bob", "generated_upper": "BOB"},
324+
}, resNew)
325+
},
326+
afterRollback: func(t *testing.T, db *sql.DB, schema string) {
327+
// The new column has been dropped from the underlying table
328+
columnName := migrations.TemporaryName("generated_upper")
329+
ColumnMustNotExist(t, db, schema, "users", columnName)
330+
331+
// The table's column count reflects the drop of the new column
332+
TableMustHaveColumnCount(t, db, schema, "users", 2)
333+
},
334+
afterComplete: func(t *testing.T, db *sql.DB, schema string) {
335+
// The new view still exists
336+
ViewMustExist(t, db, schema, "02_add_column", "users")
337+
338+
// Inserting into the new view still works
339+
MustInsert(t, db, schema, "02_add_column", "users", map[string]string{
340+
"name": "Carl",
341+
})
342+
343+
// Selecting from the new view still works
344+
res := MustSelect(t, db, schema, "02_add_column", "users")
345+
assert.Equal(t, []map[string]any{
346+
{"id": 1, "name": "Alice", "generated_upper": "ALICE"},
347+
{"id": 2, "name": "Bob", "generated_upper": "BOB"},
348+
{"id": 3, "name": "Carl", "generated_upper": "CARL"},
349+
}, res)
350+
},
351+
},
257352
})
258353
}
259354

@@ -1424,6 +1519,31 @@ func TestAddColumnValidation(t *testing.T) {
14241519
},
14251520
wantStartErr: nil,
14261521
},
1522+
{
1523+
name: "generated column is either stored or identity",
1524+
migrations: []migrations.Migration{
1525+
addTableMigrationNoPKNullable,
1526+
{
1527+
Name: "02_add_column",
1528+
Operations: migrations.Operations{
1529+
&migrations.OpAddColumn{
1530+
Table: "users",
1531+
Column: migrations.Column{
1532+
Name: "name_upper",
1533+
Type: "text",
1534+
Generated: &migrations.ColumnGenerated{
1535+
Expression: "upper(name)",
1536+
Identity: &migrations.ColumnGeneratedIdentity{
1537+
SequenceOptions: "start 2",
1538+
},
1539+
},
1540+
},
1541+
},
1542+
},
1543+
},
1544+
},
1545+
wantStartErr: migrations.InvalidGeneratedColumnError{Table: "users", Column: "name_upper"},
1546+
},
14271547
})
14281548
}
14291549

pkg/migrations/op_create_table.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,11 @@ func columnsToSQL(cols []Column, tr SQLTransformer) (string, error) {
227227
if i > 0 {
228228
sql += ", "
229229
}
230+
231+
// Generated identity columns are marked not null automatically by PostgreSQL.
232+
if col.Generated != nil && col.Generated.Identity != nil && !col.IsNullable() {
233+
col.Nullable = true
234+
}
230235
colSQL, err := columnWriter.Write(col)
231236
if err != nil {
232237
return "", err

pkg/migrations/types.go

Lines changed: 26 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pkg/sql2pgroll/alter_table_test.go

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,14 @@ func TestConvertAlterTableStatements(t *testing.T) {
222222
sql: "ALTER TABLE foo ADD COLUMN bar int CONSTRAINT fk_baz REFERENCES baz (bar) ON DELETE CASCADE",
223223
expectedOp: expect.AddColumnOp8WithOnDeleteAction(migrations.ForeignKeyReferenceOnDeleteCASCADE),
224224
},
225+
{
226+
sql: "ALTER TABLE foo ADD COLUMN bar int GENERATED BY DEFAULT AS IDENTITY ",
227+
expectedOp: expect.AddColumnOp9,
228+
},
229+
{
230+
sql: "ALTER TABLE foo ADD COLUMN bar int GENERATED ALWAYS AS ( 123 ) STORED",
231+
expectedOp: expect.AddColumnOp10,
232+
},
225233
}
226234

227235
for _, tc := range tests {
@@ -282,8 +290,6 @@ func TestUnconvertableAlterTableStatements(t *testing.T) {
282290
"ALTER TABLE foo ADD COLUMN IF NOT EXISTS bar int",
283291
"ALTER TABLE foo ADD COLUMN bar int UNIQUE DEFERRABLE",
284292
"ALTER TABLE foo ADD COLUMN bar int UNIQUE INITIALLY DEFERRED",
285-
"ALTER TABLE foo ADD COLUMN bar int GENERATED BY DEFAULT AS IDENTITY ",
286-
"ALTER TABLE foo ADD COLUMN bar int GENERATED ALWAYS AS ( 123 ) STORED",
287293
"ALTER TABLE foo ADD COLUMN bar int COLLATE en_US",
288294
"ALTER TABLE foo ADD COLUMN bar int COMPRESSION pglz",
289295
}

0 commit comments

Comments
 (0)