Repository: ariga/atlas
Pull request #583 (2022-02-14 16:49:47)

sql/postgres: add support for inspecting and applying identity attributes


Patch:
diff --git a/schema/schemaspec/extension.go b/schema/schemaspec/extension.go
index f63b8904..56597d27 100644
--- a/schema/schemaspec/extension.go
+++ b/schema/schemaspec/extension.go
@@ -247,7 +247,7 @@ func setField(field reflect.Value, attr *Attr) error {
 			return fmt.Errorf("schemaspec: value of attr %q cannot be read as string: %w", attr.K, err)
 		}
 		field.SetString(s)
-	case reflect.Int:
+	case reflect.Int, reflect.Int64:
 		i, err := attr.Int()
 		if err != nil {
 			return fmt.Errorf("schemaspec: value of attr %q cannot be read as integer: %w", attr.K, err)
diff --git a/schema/schemaspec/spec.go b/schema/schemaspec/spec.go
index f7fe00de..4df6e21a 100644
--- a/schema/schemaspec/spec.go
+++ b/schema/schemaspec/spec.go
@@ -194,6 +194,16 @@ func (a *Attr) Bools() ([]bool, error) {
 	return out, nil
 }
 
+// Resource returns the first child Resource by its type and reports whether it was found.
+func (r *Resource) Resource(t string) (*Resource, bool) {
+	for i := range r.Children {
+		if r.Children[i].Type == t {
+			return r.Children[i], true
+		}
+	}
+	return nil, false
+}
+
 // Attr returns the Attr by the provided name and reports whether it was found.
 func (r *Resource) Attr(name string) (*Attr, bool) {
 	return attrVal(r.Attrs, name)
diff --git a/sql/internal/specutil/spec.go b/sql/internal/specutil/spec.go
index 1fda8080..a6c21c02 100644
--- a/sql/internal/specutil/spec.go
+++ b/sql/internal/specutil/spec.go
@@ -25,6 +25,14 @@ func BoolAttr(k string, v bool) *schemaspec.Attr {
 	}
 }
 
+// Int64Attr is a helper method for constructing *schemaspec.Attr of type int64.
+func Int64Attr(k string, v int64) *schemaspec.Attr {
+	return &schemaspec.Attr{
+		K: k,
+		V: &schemaspec.LiteralValue{V: strconv.FormatInt(v, 10)},
+	}
+}
+
 // LitAttr is a helper method for constructing *schemaspec.Attr instances that contain literal values.
 func LitAttr(k, v string) *schemaspec.Attr {
 	return &schemaspec.Attr{
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index a4baa483..2a03aa32 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -142,3 +142,9 @@ const (
 	IndexTypeGIN   = "GIN"
 	IndexTypeGiST  = "GIST"
 )
+
+// List of "GENERATED" TYPES.
+const (
+	GeneratedTypeAlways    = "ALWAYS"
+	GeneratedTypeByDefault = "BY_DEFAULT" // BY DEFAULT.
+)
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 8e37d9d2..ee9119ca 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -584,7 +584,7 @@ func (s *state) alterColumn(b *sqlx.Builder, k schema.ChangeKind, c *schema.Colu
 			}
 			// The syntax for altering identity columns is identical to sequence_options.
 			// https://www.postgresql.org/docs/current/sql-altersequence.html
-			b.P("SET START WITH", strconv.FormatInt(id.Sequence.Start, 10), "SET INCREMENT BY", strconv.FormatInt(id.Sequence.Increment, 10), "RESTART")
+			b.P("SET GENERATED", id.Generation, "SET START WITH", strconv.FormatInt(id.Sequence.Start, 10), "SET INCREMENT BY", strconv.FormatInt(id.Sequence.Increment, 10), "RESTART")
 			k &= ^schema.ChangeAttr
 		case k.Is(schema.ChangeComment):
 			// Handled separately on modifyTable.
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 43c589e6..457eddbf 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -94,6 +94,7 @@ var (
 	hclState = schemahcl.New(
 		schemahcl.WithTypes(TypeRegistry.Specs()),
 		schemahcl.WithScopedEnums("table.index.type", IndexTypeBTree, IndexTypeHash, IndexTypeGIN, IndexTypeGiST),
+		schemahcl.WithScopedEnums("table.column.identity.generated", GeneratedTypeAlways, GeneratedTypeByDefault),
 	)
 	// UnmarshalHCL unmarshals an Atlas HCL DDL document into v.
 	UnmarshalHCL = schemaspec.UnmarshalerFunc(func(bytes []byte, i interface{}) error {
@@ -179,7 +180,37 @@ func convertColumn(spec *sqlspec.Column, _ *schema.Table) (*schema.Column, error
 	if err := fixDefaultQuotes(spec.Default); err != nil {
 		return nil, err
 	}
-	return specutil.Column(spec, convertColumnType)
+	c, err := specutil.Column(spec, convertColumnType)
+	if err != nil {
+		return nil, err
+	}
+	if r, ok := spec.Extra.Resource("identity"); ok {
+		id, err := convertIdentity(r)
+		if err != nil {
+			return nil, err
+		}
+		c.Attrs = append(c.Attrs, id)
+	}
+	return c, nil
+}
+
+func convertIdentity(r *schemaspec.Resource) (*Identity, error) {
+	var spec struct {
+		Generation string `spec:"generated"`
+		Start      int64  `spec:"start"`
+		Increment  int64  `spec:"increment"`
+	}
+	if err := r.As(&spec); err != nil {
+		return nil, err
+	}
+	id := &Identity{Generation: fromVar(spec.Generation), Sequence: &Sequence{}}
+	if spec.Start != 0 {
+		id.Sequence.Start = spec.Start
+	}
+	if spec.Increment != 0 {
+		id.Sequence.Increment = spec.Increment
+	}
+	return id, nil
 }
 
 // fixDefaultQuotes fixes the quotes on the Default field to be single quotes
@@ -337,8 +368,34 @@ func indexSpec(idx *schema.Index) (*sqlspec.Index, error) {
 }
 
 // columnSpec converts from a concrete Postgres schema.Column into a sqlspec.Column.
-func columnSpec(col *schema.Column, _ *schema.Table) (*sqlspec.Column, error) {
-	return specutil.FromColumn(col, columnTypeSpec)
+func columnSpec(c *schema.Column, _ *schema.Table) (*sqlspec.Column, error) {
+	s, err := specutil.FromColumn(c, columnTypeSpec)
+	if err != nil {
+		return nil, err
+	}
+	if i := (&Identity{}); sqlx.Has(c.Attrs, i) {
+		s.Extra.Children = append(s.Extra.Children, fromIdentity(i))
+	}
+	return s, nil
+}
+
+// fromIdentity returns the resource spec for representing the identity attributes.
+func fromIdentity(i *Identity) *schemaspec.Resource {
+	id := &schemaspec.Resource{
+		Type: "identity",
+		Attrs: []*schemaspec.Attr{
+			specutil.VarAttr("generated", strings.ToUpper(asVar(i.Generation))),
+		},
+	}
+	if s := i.Sequence; s != nil {
+		if s.Start != 1 {
+			id.Attrs = append(id.Attrs, specutil.Int64Attr("start", s.Start))
+		}
+		if s.Increment != 1 {
+			id.Attrs = append(id.Attrs, specutil.Int64Attr("increment", s.Increment))
+		}
+	}
+	return id
 }
 
 // columnTypeSpec converts from a concrete Postgres schema.Type into sqlspec.Column Type.
@@ -468,3 +525,8 @@ func attr(typ *schemaspec.Type, key string) (*schemaspec.Attr, bool) {
 	}
 	return nil, false
 }
+
+// asVar and fromVar formats a string as variable to make it HCL compatible.
+// The result is simple, replace each space with underscore.
+func asVar(s string) string   { return strings.ReplaceAll(s, " ", "_") }
+func fromVar(s string) string { return strings.ReplaceAll(s, "_", " ") }


Test patch:
diff --git a/internal/integration/testdata/postgres/column-identity.txt b/internal/integration/testdata/postgres/column-identity.txt
new file mode 100644
index 00000000..e3009c66
--- /dev/null
+++ b/internal/integration/testdata/postgres/column-identity.txt
@@ -0,0 +1,51 @@
+apply 1.hcl
+cmpshow users 1.sql
+
+# Change identity generation.
+apply 2.hcl
+cmpshow users 2.sql
+
+-- 1.hcl --
+schema "$db" {}
+
+table "users" {
+  schema = schema.$db
+  column "name" {
+    null = false
+    type = int
+    identity {
+        generated = ALWAYS
+        start = 10
+        increment = 10
+    }
+  }
+}
+
+-- 1.sql --
+                  Table "script_column_identity.users"
+ Column |  Type   | Collation | Nullable |           Default
+--------+---------+-----------+----------+------------------------------
+ name   | integer |           | not null | generated always as identity
+
+
+-- 2.hcl --
+schema "$db" {}
+
+table "users" {
+  schema = schema.$db
+  column "name" {
+    null = false
+    type = int
+    identity {
+        generated = BY_DEFAULT
+        start = 10
+        increment = 10
+    }
+  }
+}
+
+-- 2.sql --
+                    Table "script_column_identity.users"
+ Column |  Type   | Collation | Nullable |             Default
+--------+---------+-----------+----------+----------------------------------
+ name   | integer |           | not null | generated by default as identity
\ No newline at end of file
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index 3a45ac0f..9a7a4d8f 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -306,7 +306,7 @@ func TestPlanChanges(t *testing.T) {
 				Reversible:    false,
 				Transactional: true,
 				Changes: []*migrate.Change{
-					{Cmd: `ALTER TABLE "users" DROP COLUMN "name", ALTER COLUMN "id" SET START WITH 1024 SET INCREMENT BY 1 RESTART`},
+					{Cmd: `ALTER TABLE "users" DROP COLUMN "name", ALTER COLUMN "id" SET GENERATED BY DEFAULT SET START WITH 1024 SET INCREMENT BY 1 RESTART`},
 				},
 			},
 		},
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index d31538e7..4a4d4769 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -287,6 +287,39 @@ table "t" {
 	})
 }
 
+func TestUnmarshalSpec_Identity(t *testing.T) {
+	f := `
+schema "s" {}
+table "t" {
+	schema = schema.s
+	column "c" {
+		type = int
+		identity {
+			generated = %s
+			start = 10
+		}
+	}
+}
+`
+	t.Run("Invalid", func(t *testing.T) {
+		f := fmt.Sprintf(f, "UNK")
+		err := UnmarshalHCL([]byte(f), &schema.Schema{})
+		require.Error(t, err)
+	})
+	t.Run("Valid", func(t *testing.T) {
+		var (
+			s schema.Schema
+			f = fmt.Sprintf(f, "ALWAYS")
+		)
+		err := UnmarshalHCL([]byte(f), &s)
+		require.NoError(t, err)
+		id := s.Tables[0].Columns[0].Attrs[0].(*Identity)
+		require.Equal(t, GeneratedTypeAlways, id.Generation)
+		require.EqualValues(t, 10, id.Sequence.Start)
+		require.Zero(t, id.Sequence.Increment)
+	})
+}
+
 func TestMarshalSpec_Enum(t *testing.T) {
 	s := schema.New("test").
 		AddTables(