xataio · kvch · Jan 22, 2025 · Jan 17, 2025 · Jan 17, 2025 · Jan 17, 2025
diff --git a/docs/operations/create_table.mdx b/docs/operations/create_table.mdx
@@ -30,6 +30,13 @@ Each `column` is defined as:
     "name": "name of check constraint",
     "constraint": "constraint expression"
   },
+  "generated": {
+    "expression": "expression for stored column",
+    "identity": {
+      "user_specified_values": "user specified values can be used, can be ALWAYS and BY DEFAULT. Default is ALWAYS",
+      "sequence_options": "sequence options for identity columns"
+    }
+  },
   "references": {
     "name": "name of foreign key constraint",
     "table": "name of referenced table",
@@ -41,6 +48,8 @@ Each `column` is defined as:
 
 Default values are subject to the usual rules for quoting SQL expressions. In particular, string literals should be surrounded with single quotes.
 
+Generated columns can either be stored or identity columns. Options `generated.expression` and `generated.identity` cannot be set at the same time.
+
 Each `constraint` is defined as:
 
 ```json
@@ -82,7 +91,7 @@ Create one table:
 
 ### Create one table (2)
 
-Create one table:
+Create one table with generated identity column:
 
 <ExampleSnippet example="08_create_fruits_table.json" language="json" />
 

diff --git a/examples/08_create_fruits_table.json b/examples/08_create_fruits_table.json
@@ -7,8 +7,13 @@
         "columns": [
           {
             "name": "id",
-            "type": "serial",
-            "pk": true
+            "type": "bigint",
+            "pk": true,
+            "generated": {
+              "identity": {
+                "user_specified_values": "BY DEFAULT"
+              }
+            }
           },
           {
             "name": "name",

diff --git a/go.mod b/go.mod
@@ -14,7 +14,7 @@ require (
 	github.com/stretchr/testify v1.10.0
 	github.com/testcontainers/testcontainers-go v0.35.0
 	github.com/testcontainers/testcontainers-go/modules/postgres v0.35.0
-	github.com/xataio/pg_query_go/v6 v6.0.0-20241223083246-5033a992750b
+	github.com/xataio/pg_query_go/v6 v6.0.0-20250122133641-54118c062181
 	golang.org/x/tools v0.29.0
 )
 

diff --git a/go.sum b/go.sum
@@ -221,8 +221,8 @@ github.com/tklauser/go-sysconf v0.3.14 h1:g5vzr9iPFFz24v2KZXs/pvpvh8/V9Fw6vQK5ZZ
 github.com/tklauser/go-sysconf v0.3.14/go.mod h1:1ym4lWMLUOhuBOPGtRcJm7tEGX4SCYNEEEtghGG/8uY=
 github.com/tklauser/numcpus v0.8.0 h1:Mx4Wwe/FjZLeQsK/6kt2EOepwwSl7SmJrK5bV/dXYgY=
 github.com/tklauser/numcpus v0.8.0/go.mod h1:ZJZlAY+dmR4eut8epnzf0u/VwodKmryxR8txiloSqBE=
-github.com/xataio/pg_query_go/v6 v6.0.0-20241223083246-5033a992750b h1:rZyagmQlpbFFKUMW7k/hQVklPGKAZ9Nzv0xqPPd7SQU=
-github.com/xataio/pg_query_go/v6 v6.0.0-20241223083246-5033a992750b/go.mod h1:GK6bpfAhPtZb7wG/IccqvnH+cz3cmvvRTkC+MosESGo=
+github.com/xataio/pg_query_go/v6 v6.0.0-20250122133641-54118c062181 h1:iLOHgul20WFUhO4eFpJ/lZRkHzZICF2ghzncxtOcD0E=
+github.com/xataio/pg_query_go/v6 v6.0.0-20250122133641-54118c062181/go.mod h1:GK6bpfAhPtZb7wG/IccqvnH+cz3cmvvRTkC+MosESGo=
 github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778/go.mod h1:2MuV+tbUrU1zIOPMxZ5EncGwgmMJsa+9ucAQZXxsObs=
 github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
 github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=

diff --git a/pkg/migrations/errors.go b/pkg/migrations/errors.go
@@ -254,3 +254,12 @@ type PrimaryKeysAreAlreadySetError struct {
 func (e PrimaryKeysAreAlreadySetError) Error() string {
 	return fmt.Sprintf("table %q already has a primary key configuration in columns list", e.Table)
 }
+
+type InvalidGeneratedColumnError struct {
+	Table  string
+	Column string
+}
+
+func (e InvalidGeneratedColumnError) Error() string {
+	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)
+}
diff --git a/pkg/migrations/op_add_column.go b/pkg/migrations/op_add_column.go
@@ -184,7 +184,11 @@ func (o *OpAddColumn) Validate(ctx context.Context, s *schema.Schema) error {
 		}
 	}
 
-	if !o.Column.IsNullable() && o.Column.Default == nil && o.Up == "" && !o.Column.HasImplicitDefault() {
+	if o.Column.Generated != nil && o.Column.Generated.Expression != "" && o.Column.Generated.Identity != nil {
+		return InvalidGeneratedColumnError{Table: o.Table, Column: o.Column.Name}
+	}
+
+	if !o.Column.IsNullable() && o.Column.Default == nil && o.Up == "" && !o.Column.HasImplicitDefault() && o.Column.Generated == nil {
 		return FieldRequiredError{Name: "up"}
 	}
 
@@ -213,6 +217,11 @@ func addColumn(ctx context.Context, conn db.DB, o OpAddColumn, t *schema.Table,
 		o.Column.Nullable = true
 	}
 
+	// Generated identity columns are marked not null automatically by PostgreSQL.
+	if o.Column.Generated != nil && o.Column.Generated.Identity != nil && !o.Column.IsNullable() {
+		o.Column.Nullable = true
+	}
+
 	// Don't add a column with a CHECK constraint directly.
 	// They are handled by:
 	// - adding the column without the constraint
@@ -292,6 +301,18 @@ func (w ColumnSQLWriter) Write(col Column) (string, error) {
 		}
 		sql += fmt.Sprintf(" DEFAULT %s", d)
 	}
+
+	if col.Generated != nil {
+		if col.Generated.Expression != "" {
+			sql += fmt.Sprintf(" GENERATED ALWAYS AS (%s) STORED", col.Generated.Expression)
+		} else if col.Generated.Identity != nil {
+			sql += fmt.Sprintf(" GENERATED %s AS IDENTITY", col.Generated.Identity.UserSpecifiedValues)
+			if col.Generated.Identity.SequenceOptions != "" {
+				sql += fmt.Sprintf(" (%s)", col.Generated.Identity.SequenceOptions)
+			}
+		}
+	}
+
 	if col.References != nil {
 		onDelete := string(ForeignKeyReferenceOnDeleteNOACTION)
 		if col.References.OnDelete != "" {

diff --git a/pkg/migrations/op_add_column_test.go b/pkg/migrations/op_add_column_test.go
@@ -254,6 +254,101 @@ func TestAddColumn(t *testing.T) {
 				}, res)
 			},
 		},
+		{
+			name: "add generated identity column and a regular stored column",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_add_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "users",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "bigint",
+									Pk:   true,
+									Generated: &migrations.ColumnGenerated{
+										Identity: &migrations.ColumnGeneratedIdentity{
+											UserSpecifiedValues: "ALWAYS",
+										},
+									},
+								},
+								{
+									Name: "name",
+									Type: "varchar(255)",
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_add_column",
+					Operations: migrations.Operations{
+						&migrations.OpAddColumn{
+							Table: "users",
+							Column: migrations.Column{
+								Name: "generated_upper",
+								Type: "varchar(255)",
+								Generated: &migrations.ColumnGenerated{
+									Expression: "upper(name)",
+								},
+								Nullable: false,
+							},
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				// old and new views of the table should exist
+				ViewMustExist(t, db, schema, "01_add_table", "users")
+				ViewMustExist(t, db, schema, "02_add_column", "users")
+
+				// inserting via both the old and the new views works
+				MustInsert(t, db, schema, "01_add_table", "users", map[string]string{
+					"name": "Alice",
+				})
+				MustInsert(t, db, schema, "02_add_column", "users", map[string]string{
+					"name": "Bob",
+				})
+
+				// selecting from both the old and the new views works
+				resOld := MustSelect(t, db, schema, "01_add_table", "users")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "name": "Alice"},
+					{"id": 2, "name": "Bob"},
+				}, resOld)
+				resNew := MustSelect(t, db, schema, "02_add_column", "users")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "name": "Alice", "generated_upper": "ALICE"},
+					{"id": 2, "name": "Bob", "generated_upper": "BOB"},
+				}, resNew)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+				// The new column has been dropped from the underlying table
+				columnName := migrations.TemporaryName("generated_upper")
+				ColumnMustNotExist(t, db, schema, "users", columnName)
+
+				// The table's column count reflects the drop of the new column
+				TableMustHaveColumnCount(t, db, schema, "users", 2)
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				// The new view still exists
+				ViewMustExist(t, db, schema, "02_add_column", "users")
+
+				// Inserting into the new view still works
+				MustInsert(t, db, schema, "02_add_column", "users", map[string]string{
+					"name": "Carl",
+				})
+
+				// Selecting from the new view still works
+				res := MustSelect(t, db, schema, "02_add_column", "users")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "name": "Alice", "generated_upper": "ALICE"},
+					{"id": 2, "name": "Bob", "generated_upper": "BOB"},
+					{"id": 3, "name": "Carl", "generated_upper": "CARL"},
+				}, res)
+			},
+		},
 	})
 }
 
@@ -1424,6 +1519,31 @@ func TestAddColumnValidation(t *testing.T) {
 			},
 			wantStartErr: nil,
 		},
+		{
+			name: "generated column is either stored or identity",
+			migrations: []migrations.Migration{
+				addTableMigrationNoPKNullable,
+				{
+					Name: "02_add_column",
+					Operations: migrations.Operations{
+						&migrations.OpAddColumn{
+							Table: "users",
+							Column: migrations.Column{
+								Name: "name_upper",
+								Type: "text",
+								Generated: &migrations.ColumnGenerated{
+									Expression: "upper(name)",
+									Identity: &migrations.ColumnGeneratedIdentity{
+										SequenceOptions: "start 2",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			wantStartErr: migrations.InvalidGeneratedColumnError{Table: "users", Column: "name_upper"},
+		},
 	})
 }
 

diff --git a/pkg/migrations/op_create_table.go b/pkg/migrations/op_create_table.go
@@ -227,6 +227,11 @@ func columnsToSQL(cols []Column, tr SQLTransformer) (string, error) {
 		if i > 0 {
 			sql += ", "
 		}
+
+		// Generated identity columns are marked not null automatically by PostgreSQL.
+		if col.Generated != nil && col.Generated.Identity != nil && !col.IsNullable() {
+			col.Nullable = true
+		}
 		colSQL, err := columnWriter.Write(col)
 		if err != nil {
 			return "", err

diff --git a/pkg/migrations/types.go b/pkg/migrations/types.go
diff --git a/pkg/sql2pgroll/alter_table_test.go b/pkg/sql2pgroll/alter_table_test.go
@@ -222,6 +222,14 @@ func TestConvertAlterTableStatements(t *testing.T) {
 			sql:        "ALTER TABLE foo ADD COLUMN bar int CONSTRAINT fk_baz REFERENCES baz (bar) ON DELETE CASCADE",
 			expectedOp: expect.AddColumnOp8WithOnDeleteAction(migrations.ForeignKeyReferenceOnDeleteCASCADE),
 		},
+		{
+			sql:        "ALTER TABLE foo ADD COLUMN bar int GENERATED BY DEFAULT AS IDENTITY ",
+			expectedOp: expect.AddColumnOp9,
+		},
+		{
+			sql:        "ALTER TABLE foo ADD COLUMN bar int GENERATED ALWAYS AS ( 123 ) STORED",
+			expectedOp: expect.AddColumnOp10,
+		},
 	}
 
 	for _, tc := range tests {
@@ -282,8 +290,6 @@ func TestUnconvertableAlterTableStatements(t *testing.T) {
 		"ALTER TABLE foo ADD COLUMN IF NOT EXISTS bar int",
 		"ALTER TABLE foo ADD COLUMN bar int UNIQUE DEFERRABLE",
 		"ALTER TABLE foo ADD COLUMN bar int UNIQUE INITIALLY DEFERRED",
-		"ALTER TABLE foo ADD COLUMN bar int GENERATED BY DEFAULT AS IDENTITY ",
-		"ALTER TABLE foo ADD COLUMN bar int GENERATED ALWAYS AS ( 123 ) STORED",
 		"ALTER TABLE foo ADD COLUMN bar int COLLATE en_US",
 		"ALTER TABLE foo ADD COLUMN bar int COMPRESSION pglz",
 	}

diff --git a/pkg/sql2pgroll/create_table.go b/pkg/sql2pgroll/create_table.go
@@ -102,6 +102,7 @@ func convertColumnDef(tableName string, col *pgq.ColumnDef) (*migrations.Column,
 	var check *migrations.CheckConstraint
 	var foreignKey *migrations.ForeignKeyReference
 	var defaultValue *string
+	var generated *migrations.ColumnGenerated
 	for _, c := range col.GetConstraints() {
 		switch c.GetConstraint().GetContype() {
 		case pgq.ConstrType_CONSTR_NULL:
@@ -171,11 +172,43 @@ func convertColumnDef(tableName string, col *pgq.ColumnDef) (*migrations.Column,
 			// are supported, but no extra annotation is needed
 			continue
 		case pgq.ConstrType_CONSTR_GENERATED:
-			// Generated columns are not supported
-			return nil, nil
+			if c.GetConstraint().GetRawExpr() != nil {
+				generatorExpr, err := pgq.DeparseExpr(c.GetConstraint().GetRawExpr())
+				if err != nil {
+					return nil, fmt.Errorf("deparsing generated expression: %w", err)
+				}
+				generated = &migrations.ColumnGenerated{
+					Expression: generatorExpr,
+				}
+			} else {
+				return nil, nil
+			}
+			notNull = true
 		case pgq.ConstrType_CONSTR_IDENTITY:
-			// Identity columns are not supported
-			return nil, nil
+			var when migrations.ColumnGeneratedIdentityUserSpecifiedValues
+			switch c.GetConstraint().GeneratedWhen {
+			case "a":
+				when = migrations.ColumnGeneratedIdentityUserSpecifiedValuesALWAYS
+			case "d":
+				when = migrations.ColumnGeneratedIdentityUserSpecifiedValuesBYDEFAULT
+			default:
+				return nil, nil
+			}
+			sequenceOptions := ""
+			if c.GetConstraint().GetOptions() != nil {
+				sequenceOptions, err = pgq.DeparseParenthesizedSeqOptList(c.GetConstraint().GetOptions())
+				if err != nil {
+					return nil, fmt.Errorf("parsing sequence options: %w", err)
+				}
+				sequenceOptions = sequenceOptions[1 : len(sequenceOptions)-1]
+			}
+			generated = &migrations.ColumnGenerated{
+				Identity: &migrations.ColumnGeneratedIdentity{
+					UserSpecifiedValues: when,
+					SequenceOptions:     sequenceOptions,
+				},
+			}
+			notNull = true
 		case pgq.ConstrType_CONSTR_ATTR_DEFERRABLE:
 			// Deferrable constraints are not supported
 			return nil, nil
@@ -197,6 +230,7 @@ func convertColumnDef(tableName string, col *pgq.ColumnDef) (*migrations.Column,
 		References: foreignKey,
 		Default:    defaultValue,
 		Unique:     unique,
+		Generated:  generated,
 	}, nil
 }