Repository: ariga/atlas
Pull request #360 (2022-01-02 17:07:34)

sql/postgres: support for inspecting and diffing identity columns
Next PR adds support for migration.


Patch:
diff --git a/sql/postgres/diff.go b/sql/postgres/diff.go
index f5c3fa0b..bd6de676 100644
--- a/sql/postgres/diff.go
+++ b/sql/postgres/diff.go
@@ -50,6 +50,9 @@ func (d *diff) ColumnChange(from, to *schema.Column) (schema.ChangeKind, error)
 	if changed {
 		change |= schema.ChangeDefault
 	}
+	if identityChanged(from.Attrs, to.Attrs) {
+		change |= schema.ChangeAttr
+	}
 	return change, nil
 }
 
@@ -258,6 +261,44 @@ func (d *diff) typesEqual(x, y string) (bool, error) {
 	return b, nil
 }
 
+// Default IDENTITY attributes.
+const (
+	defaultIdentityGen  = "BY DEFAULT"
+	defaultSeqStart     = 1
+	defaultSeqIncrement = 1
+)
+
+// identityChanged reports if one of the identity attributes was changed.
+func identityChanged(from, to []schema.Attr) bool {
+	i1, ok1 := identity(from)
+	i2, ok2 := identity(to)
+	if !ok1 && !ok2 || ok1 != ok2 {
+		return ok1 != ok2
+	}
+	return i1.Generation != i2.Generation || i1.Sequence.Start != i2.Sequence.Start || i1.Sequence.Increment != i2.Sequence.Increment
+}
+
+func identity(attrs []schema.Attr) (*Identity, bool) {
+	i := &Identity{}
+	if !sqlx.Has(attrs, i) {
+		return nil, false
+	}
+	if i.Generation == "" {
+		i.Generation = defaultIdentityGen
+	}
+	if i.Sequence == nil {
+		i.Sequence = &Sequence{Start: defaultSeqStart, Increment: defaultSeqIncrement}
+		return i, true
+	}
+	if i.Sequence.Start == 0 {
+		i.Sequence.Start = defaultSeqStart
+	}
+	if i.Sequence.Increment == 0 {
+		i.Sequence.Increment = defaultSeqIncrement
+	}
+	return i, true
+}
+
 func trimCast(s string) string {
 	i := strings.LastIndex(s, "::")
 	if i == -1 {
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 72d10eaa..d2557a3e 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -156,10 +156,10 @@ func (i *inspect) columns(ctx context.Context, t *schema.Table) error {
 // addColumn scans the current row and adds a new column from it to the table.
 func (i *inspect) addColumn(t *schema.Table, rows *sql.Rows) error {
 	var (
-		typid, maxlen, precision, scale                                                                sql.NullInt64
+		typid, maxlen, precision, scale, seqstart, seqinc                                              sql.NullInt64
 		name, typ, nullable, defaults, udt, identity, generation, charset, collation, comment, typtype sql.NullString
 	)
-	if err := rows.Scan(&name, &typ, &nullable, &defaults, &maxlen, &precision, &scale, &charset, &collation, &udt, &identity, &generation, &comment, &typtype, &typid); err != nil {
+	if err := rows.Scan(&name, &typ, &nullable, &defaults, &maxlen, &precision, &scale, &charset, &collation, &udt, &identity, &seqstart, &seqinc, &generation, &comment, &typtype, &typid); err != nil {
 		return err
 	}
 	c := &schema.Column{
@@ -184,6 +184,10 @@ func (i *inspect) addColumn(t *schema.Table, rows *sql.Rows) error {
 	if identity.String == "YES" {
 		c.Attrs = append(c.Attrs, &Identity{
 			Generation: generation.String,
+			Sequence: &Sequence{
+				Start:     seqstart.Int64,
+				Increment: seqinc.Int64,
+			},
 		})
 	}
 	if sqlx.ValidString(comment) {
@@ -623,10 +627,17 @@ type (
 		T string // c, f, p, u, t, x.
 	}
 
+	// Sequence defines (the supported) sequence options.
+	// https://www.postgresql.org/docs/current/sql-createsequence.html
+	Sequence struct {
+		Start, Increment int64
+	}
+
 	// Identity defines an identity column.
 	Identity struct {
 		schema.Attr
 		Generation string // ALWAYS, BY DEFAULT.
+		Sequence   *Sequence
 	}
 
 	// IndexType represents an index type.
@@ -719,6 +730,8 @@ SELECT
 	t1.collation_name,
 	t1.udt_name,
 	t1.is_identity,
+	t1.identity_start,
+	t1.identity_increment,
 	t1.identity_generation,
 	col_description(to_regclass("table_schema" || '.' || "table_name")::oid, "ordinal_position") AS comment,
 	t2.typtype,
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index ce79c067..e5222987 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -440,7 +440,7 @@ func (s *state) column(b *sqlx.Builder, c *schema.Column) {
 			b.P("COLLATE").Ident(attr.V)
 		case *Identity:
 			if attr.Generation == "" {
-				attr.Generation = "BY DEFAULT"
+				attr.Generation = defaultIdentityGen
 			}
 			b.P("GENERATED", attr.Generation, "AS IDENTITY")
 		default:


Test patch:
diff --git a/sql/postgres/diff_test.go b/sql/postgres/diff_test.go
index 4e397511..5959a77a 100644
--- a/sql/postgres/diff_test.go
+++ b/sql/postgres/diff_test.go
@@ -28,7 +28,7 @@ func TestDiff_TableDiff(t *testing.T) {
 			to:   &schema.Table{Name: "users"},
 		},
 		{
-			name: "change primary key",
+			name: "change primary key columns",
 			from: func() *schema.Table {
 				t := &schema.Table{Name: "users", Schema: &schema.Schema{Name: "public"}, Columns: []*schema.Column{{Name: "id", Type: &schema.ColumnType{Raw: "int", Type: &schema.IntegerType{T: "int"}}}}}
 				t.PrimaryKey = &schema.Index{
@@ -39,6 +39,36 @@ func TestDiff_TableDiff(t *testing.T) {
 			to:      &schema.Table{Name: "users"},
 			wantErr: true,
 		},
+		{
+			name: "change identity attributes",
+			from: func() *schema.Table {
+				t := &schema.Table{
+					Name: "users",
+					Columns: []*schema.Column{
+						{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int"}}},
+					},
+				}
+				t.PrimaryKey = &schema.Index{Parts: []*schema.IndexPart{{C: t.Columns[0]}}}
+				return t
+			}(),
+			to: func() *schema.Table {
+				t := &schema.Table{
+					Name: "users",
+					Columns: []*schema.Column{
+						{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int"}}, Attrs: []schema.Attr{&Identity{Sequence: &Sequence{Start: 1024, Increment: 1}}}},
+					},
+				}
+				t.PrimaryKey = &schema.Index{Parts: []*schema.IndexPart{{C: t.Columns[0]}}}
+				return t
+			}(),
+			wantChanges: []schema.Change{
+				&schema.ModifyColumn{
+					From:   &schema.Column{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int"}}},
+					To:     &schema.Column{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int"}}, Attrs: []schema.Attr{&Identity{Sequence: &Sequence{Start: 1024, Increment: 1}}}},
+					Change: schema.ChangeAttr,
+				},
+			},
+		},
 		{
 			name: "add check",
 			from: &schema.Table{Name: "t1", Schema: &schema.Schema{Name: "public"}},
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index 597b6111..07f4bd89 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -57,34 +57,34 @@ func TestDriver_InspectTable(t *testing.T) {
 				m.ExpectQuery(sqltest.Escape(columnsQuery)).
 					WithArgs("public", "users").
 					WillReturnRows(sqltest.Rows(`
- column_name |      data_type      | is_nullable |         column_default          | character_maximum_length | numeric_precision | numeric_scale | character_set_name | collation_name | udt_name | is_identity | identity_generation | comment | typtype |  oid  
--------------+---------------------+-------------+---------------------------------+--------------------------+-------------------+---------------+--------------------+----------------+----------+-------------+---------------------+---------+---------+-------
- id          | bigint              | NO          |                                 |                          |                64 |             0 |                    |                | int8     | NO          |                     |         | b       |    20
- rank        | integer             | YES         |                                 |                          |                32 |             0 |                    |                | int4     | NO          |                     | rank    | b       |    23
- c1          | smallint            | NO          |           1000                  |                          |                16 |             0 |                    |                | int2     | NO          |                     |         | b       |    21
- c2          | bit                 | NO          |                                 |                        1 |                   |               |                    |                | bit      | NO          |                     |         | b       |  1560
- c3          | bit varying         | NO          |                                 |                       10 |                   |               |                    |                | varbit   | NO          |                     |         | b       |  1562
- c4          | boolean             | NO          |                                 |                          |                   |               |                    |                | bool     | NO          |                     |         | b       |    16
- c5          | bytea               | NO          |                                 |                          |                   |               |                    |                | bytea    | NO          |                     |         | b       |    17
- c6          | character           | NO          |                                 |                      100 |                   |               |                    |                | bpchar   | NO          |                     |         | b       |  1042
- c7          | character varying   | NO          |                                 |                          |                   |               |                    |                | varchar  | NO          |                     |         | b       |  1043
- c8          | cidr                | NO          |                                 |                          |                   |               |                    |                | cidr     | NO          |                     |         | b       |   650
- c9          | circle              | NO          |                                 |                          |                   |               |                    |                | circle   | NO          |                     |         | b       |   718
- c10         | date                | NO          |                                 |                          |                   |               |                    |                | date     | NO          |                     |         | b       |  1082
- c11         | time with time zone | NO          |                                 |                          |                   |               |                    |                | timetz   | NO          |                     |         | b       |  1266
- c12         | double precision    | NO          |                                 |                          |                53 |               |                    |                | float8   | NO          |                     |         | b       |   701
- c13         | real                | NO          |           random()              |                          |                24 |               |                    |                | float4   | NO          |                     |         | b       |   700
- c14         | json                | NO          |           '{}'::json            |                          |                   |               |                    |                | json     | NO          |                     |         | b       |   114
- c15         | jsonb               | NO          |           '{}'::jsonb           |                          |                   |               |                    |                | jsonb    | NO          |                     |         | b       |  3802
- c16         | money               | NO          |                                 |                          |                   |               |                    |                | money    | NO          |                     |         | b       |   790
- c17         | numeric             | NO          |                                 |                          |                   |               |                    |                | numeric  | NO          |                     |         | b       |  1700
- c18         | numeric             | NO          |                                 |                          |                 4 |             4 |                    |                | numeric  | NO          |                     |         | b       |  1700
- c19         | integer             | NO          | nextval('t1_c19_seq'::regclass) |                          |                32 |             0 |                    |                | int4     | NO          |                     |         | b       |    23
- c20         | uuid                | NO          |                                 |                          |                   |               |                    |                | uuid     | NO          |                     |         | b       |  2950
- c21         | xml                 | NO          |                                 |                          |                   |               |                    |                | xml      | NO          |                     |         | b       |   142
- c22         | ARRAY               | YES         |                                 |                          |                   |               |                    |                | _int4    | NO          |                     |         | b       |  1007
- c23         | USER-DEFINED        | YES         |                                 |                          |                   |               |                    |                | ltree    | NO          |                     |         | b       | 16535
- c24         | USER-DEFINED        | NO          |                                 |                          |                   |               |                    |                | state    | NO          |                     |         | e       | 16774
+ column_name |      data_type      | is_nullable |         column_default          | character_maximum_length | numeric_precision | numeric_scale | character_set_name | collation_name | udt_name | is_identity | identity_start | identity_increment | identity_generation | comment | typtype |  oid  
+-------------+---------------------+-------------+---------------------------------+--------------------------+-------------------+---------------+--------------------+----------------+----------+-------------+----------------+--------------------+---------------------+---------+---------+-------
+ id          | bigint              | NO          |                                 |                          |                64 |             0 |                    |                | int8     | YES         |      100       |          1         |    BY DEFAULT       |         | b       |    20
+ rank        | integer             | YES         |                                 |                          |                32 |             0 |                    |                | int4     | NO          |                |                    |                     | rank    | b       |    23
+ c1          | smallint            | NO          |           1000                  |                          |                16 |             0 |                    |                | int2     | NO          |                |                    |                     |         | b       |    21
+ c2          | bit                 | NO          |                                 |                        1 |                   |               |                    |                | bit      | NO          |                |                    |                     |         | b       |  1560
+ c3          | bit varying         | NO          |                                 |                       10 |                   |               |                    |                | varbit   | NO          |                |                    |                     |         | b       |  1562
+ c4          | boolean             | NO          |                                 |                          |                   |               |                    |                | bool     | NO          |                |                    |                     |         | b       |    16
+ c5          | bytea               | NO          |                                 |                          |                   |               |                    |                | bytea    | NO          |                |                    |                     |         | b       |    17
+ c6          | character           | NO          |                                 |                      100 |                   |               |                    |                | bpchar   | NO          |                |                    |                     |         | b       |  1042
+ c7          | character varying   | NO          |                                 |                          |                   |               |                    |                | varchar  | NO          |                |                    |                     |         | b       |  1043
+ c8          | cidr                | NO          |                                 |                          |                   |               |                    |                | cidr     | NO          |                |                    |                     |         | b       |   650
+ c9          | circle              | NO          |                                 |                          |                   |               |                    |                | circle   | NO          |                |                    |                     |         | b       |   718
+ c10         | date                | NO          |                                 |                          |                   |               |                    |                | date     | NO          |                |                    |                     |         | b       |  1082
+ c11         | time with time zone | NO          |                                 |                          |                   |               |                    |                | timetz   | NO          |                |                    |                     |         | b       |  1266
+ c12         | double precision    | NO          |                                 |                          |                53 |               |                    |                | float8   | NO          |                |                    |                     |         | b       |   701
+ c13         | real                | NO          |           random()              |                          |                24 |               |                    |                | float4   | NO          |                |                    |                     |         | b       |   700
+ c14         | json                | NO          |           '{}'::json            |                          |                   |               |                    |                | json     | NO          |                |                    |                     |         | b       |   114
+ c15         | jsonb               | NO          |           '{}'::jsonb           |                          |                   |               |                    |                | jsonb    | NO          |                |                    |                     |         | b       |  3802
+ c16         | money               | NO          |                                 |                          |                   |               |                    |                | money    | NO          |                |                    |                     |         | b       |   790
+ c17         | numeric             | NO          |                                 |                          |                   |               |                    |                | numeric  | NO          |                |                    |                     |         | b       |  1700
+ c18         | numeric             | NO          |                                 |                          |                 4 |             4 |                    |                | numeric  | NO          |                |                    |                     |         | b       |  1700
+ c19         | integer             | NO          | nextval('t1_c19_seq'::regclass) |                          |                32 |             0 |                    |                | int4     | NO          |                |                    |                     |         | b       |    23
+ c20         | uuid                | NO          |                                 |                          |                   |               |                    |                | uuid     | NO          |                |                    |                     |         | b       |  2950
+ c21         | xml                 | NO          |                                 |                          |                   |               |                    |                | xml      | NO          |                |                    |                     |         | b       |   142
+ c22         | ARRAY               | YES         |                                 |                          |                   |               |                    |                | _int4    | NO          |                |                    |                     |         | b       |  1007
+ c23         | USER-DEFINED        | YES         |                                 |                          |                   |               |                    |                | ltree    | NO          |                |                    |                     |         | b       | 16535
+ c24         | USER-DEFINED        | NO          |                                 |                          |                   |               |                    |                | state    | NO          |                |                    |                     |         | e       | 16774
 `))
 				m.ExpectQuery(sqltest.Escape(`SELECT enumtypid, enumlabel FROM pg_enum WHERE enumtypid IN ($1)`)).
 					WithArgs(16774).
@@ -102,7 +102,7 @@ func TestDriver_InspectTable(t *testing.T) {
 				require.NoError(err)
 				require.Equal("users", t.Name)
 				require.EqualValues([]*schema.Column{
-					{Name: "id", Type: &schema.ColumnType{Raw: "bigint", Type: &schema.IntegerType{T: "bigint"}}},
+					{Name: "id", Type: &schema.ColumnType{Raw: "bigint", Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&Identity{Generation: "BY DEFAULT", Sequence: &Sequence{Start: 100, Increment: 1}}}},
 					{Name: "rank", Type: &schema.ColumnType{Raw: "integer", Null: true, Type: &schema.IntegerType{T: "integer"}}, Attrs: []schema.Attr{&schema.Comment{Text: "rank"}}},
 					{Name: "c1", Type: &schema.ColumnType{Raw: "smallint", Type: &schema.IntegerType{T: "smallint"}}, Default: &schema.Literal{V: "1000"}},
 					{Name: "c2", Type: &schema.ColumnType{Raw: "bit", Type: &BitType{T: "bit", Len: 1}}},
@@ -139,10 +139,10 @@ func TestDriver_InspectTable(t *testing.T) {
 				m.ExpectQuery(sqltest.Escape(columnsQuery)).
 					WithArgs("public", "users").
 					WillReturnRows(sqltest.Rows(`
- column_name |      data_type      | is_nullable |         column_default          | character_maximum_length | numeric_precision | numeric_scale | character_set_name | collation_name | udt_name | is_identity | identity_generation | comment | typtype |  oid
--------------+---------------------+-------------+---------------------------------+--------------------------+-------------------+---------------+--------------------+----------------+----------+-------------+---------------------+---------+---------+-------
- id          | bigint              | NO          |                                 |                          |                64 |             0 |                    |                | int8     | NO          |                     |         | b       |    20
- c1          | smallint            | NO          |                                 |                          |                16 |             0 |                    |                | int2     | NO          |                     |         | b       |    21
+ column_name |      data_type      | is_nullable |         column_default          | character_maximum_length | numeric_precision | numeric_scale | character_set_name | collation_name | udt_name | is_identity | identity_start | identity_increment | identity_generation | comment | typtype |  oid
+-------------+---------------------+-------------+---------------------------------+--------------------------+-------------------+---------------+--------------------+----------------+----------+-------------+----------------+--------------------+---------------------+---------+---------+-------
+ id          | bigint              | NO          |                                 |                          |                64 |             0 |                    |                | int8     | NO          |                |                    |                     |         | b       |    20
+ c1          | smallint            | NO          |                                 |                          |                16 |             0 |                    |                | int2     | NO          |                |                    |                     |         | b       |    21
 `))
 				m.ExpectQuery(sqltest.Escape(indexesQuery)).
 					WithArgs("public", "users").
@@ -196,11 +196,11 @@ func TestDriver_InspectTable(t *testing.T) {
 					WithArgs("public", "users").
 					WithArgs("public", "users").
 					WillReturnRows(sqltest.Rows(`
- column_name |      data_type      | is_nullable |         column_default          | character_maximum_length | numeric_precision | numeric_scale | character_set_name | collation_name | udt_name | is_identity | identity_generation | comment | typtype |  oid
--------------+---------------------+-------------+---------------------------------+--------------------------+-------------------+---------------+--------------------+----------------+----------+-------------+---------------------+---------+---------+-------
- id          | integer             | NO          |                                 |                          |                32 |             0 |                    |                | int      | NO          |                     |         | b       |    20
- oid         | integer             | NO          |                                 |                          |                32 |             0 |                    |                | int      | NO          |                     |         | b       |    21
- uid         | integer             | NO          |                                 |                          |                32 |             0 |                    |                | int      | NO          |                     |         | b       |    21
+ column_name |      data_type      | is_nullable |         column_default          | character_maximum_length | numeric_precision | numeric_scale | character_set_name | collation_name | udt_name | is_identity | identity_start | identity_increment | identity_generation | comment | typtype |  oid
+-------------+---------------------+-------------+---------------------------------+--------------------------+-------------------+---------------+--------------------+----------------+----------+-------------+----------------+--------------------+---------------------+---------+---------+-------
+ id          | integer             | NO          |                                 |                          |                32 |             0 |                    |                | int      | NO          |                |                    |                     |         | b       |    20
+ oid         | integer             | NO          |                                 |                          |                32 |             0 |                    |                | int      | NO          |                |                    |                     |         | b       |    21
+ uid         | integer             | NO          |                                 |                          |                32 |             0 |                    |                | int      | NO          |                |                    |                     |         | b       |    21
 `))
 				m.noIndexes()
 				m.ExpectQuery(sqltest.Escape(fksQuery)).
@@ -246,11 +246,11 @@ func TestDriver_InspectTable(t *testing.T) {
 					WithArgs("public", "users").
 					WithArgs("public", "users").
 					WillReturnRows(sqltest.Rows(`
- column_name | data_type | is_nullable | column_default | character_maximum_length | numeric_precision | numeric_scale | character_set_name | collation_name | udt_name | is_identity | identity_generation | comment | typtype | oid
--------------+-----------+-------------+----------------+--------------------------+-------------------+---------------+--------------------+----------------+----------+-------------+---------------------+---------+---------+-----
- c1          | integer   | NO          |                |                          |                32 |             0 |                    |                | int4     | NO          |                     |         | b       |  23
- c2          | integer   | NO          |                |                          |                32 |             0 |                    |                | int4     | NO          |                     |         | b       |  23
- c3          | integer   | NO          |                |                          |                32 |             0 |                    |                | int4     | NO          |                     |         | b       |  23
+ column_name | data_type | is_nullable | column_default | character_maximum_length | numeric_precision | numeric_scale | character_set_name | collation_name | udt_name | is_identity | identity_start | identity_increment | identity_generation | comment | typtype | oid
+-------------+-----------+-------------+----------------+--------------------------+-------------------+---------------+--------------------+----------------+----------+-------------+----------------+--------------------+---------------------+---------+---------+-----
+ c1          | integer   | NO          |                |                          |                32 |             0 |                    |                | int4     | NO          |                |                    |                     |         | b       |  23
+ c2          | integer   | NO          |                |                          |                32 |             0 |                    |                | int4     | NO          |                |                    |                     |         | b       |  23
+ c3          | integer   | NO          |                |                          |                32 |             0 |                    |                | int4     | NO          |                |                    |                     |         | b       |  23
 `))
 				m.noIndexes()
 				m.noFKs()