Repository: ariga/atlas
Pull request #1307 (2022-11-24 20:20:48)

sql/postgres: scan indexes that support foreign-key constraints


Patch:
diff --git a/sql/postgres/crdb.go b/sql/postgres/crdb.go
index 5df44fae..873929bf 100644
--- a/sql/postgres/crdb.go
+++ b/sql/postgres/crdb.go
@@ -228,7 +228,7 @@ func (i *inspect) crdbAddIndexes(s *schema.Schema, rows *sql.Rows) error {
 				idx.Attrs = append(idx.Attrs, &schema.Comment{Text: comment.String})
 			}
 			if sqlx.ValidString(contype) {
-				idx.Attrs = append(idx.Attrs, &ConType{T: contype.String})
+				idx.Attrs = append(idx.Attrs, &Constraint{T: contype.String})
 			}
 			if sqlx.ValidString(pred) {
 				idx.Attrs = append(idx.Attrs, &IndexPredicate{P: pred.String})
@@ -267,6 +267,8 @@ const (
 )
 
 // CockroachDB query for getting schema indexes.
+// Scanning constraints is disabled due to internal CockroachDB error.
+// (internal error: unexpected type *tree.DOidWrapper for key value)
 const crdbIndexesQuery = `
 SELECT
 	t.relname AS table_name,
@@ -274,7 +276,7 @@ SELECT
 	a.attname AS column_name,
 	idx.indisprimary AS primary,
 	idx.indisunique AS unique,
-	c.contype AS constraint_type,
+	NULL AS constraints,
 	pgi.indexdef create_stmt,
 	pg_get_expr(idx.indpred, idx.indrelid) AS predicate,
 	pg_get_indexdef(idx.indexrelid, idx.ord, false) AS expression,
@@ -290,13 +292,11 @@ SELECT
 	JOIN pg_class i ON i.oid = idx.indexrelid
 	JOIN pg_class t ON t.oid = idx.indrelid
 	JOIN pg_namespace n ON n.oid = t.relnamespace
-	LEFT JOIN pg_constraint c ON idx.indexrelid = c.conindid
 	LEFT JOIN pg_indexes pgi ON pgi.tablename = t.relname AND indexname = i.relname AND n.nspname = pgi.schemaname
 	LEFT JOIN pg_attribute a ON (a.attrelid, a.attnum) = (idx.indrelid, idx.key)
 WHERE
 	n.nspname = $1
 	AND t.relname IN (%s)
-	AND COALESCE(c.contype, '') <> 'f'
 ORDER BY
 	table_name, index_name, idx.ord
 `
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 59536d53..eb5ebec5 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"context"
 	"database/sql"
+	"encoding/json"
 	"fmt"
 	"regexp"
 	"strconv"
@@ -321,14 +322,14 @@ func (i *inspect) addIndexes(s *schema.Schema, rows *sql.Rows) error {
 	names := make(map[string]*schema.Index)
 	for rows.Next() {
 		var (
-			uniq, primary, included                                           bool
-			table, name, typ                                                  string
-			desc, nullsfirst, nullslast, opcdefault                           sql.NullBool
-			column, contype, pred, expr, comment, options, opcname, opcparams sql.NullString
+			uniq, primary, included                                               bool
+			table, name, typ                                                      string
+			desc, nullsfirst, nullslast, opcdefault                               sql.NullBool
+			column, constraints, pred, expr, comment, options, opcname, opcparams sql.NullString
 		)
 		if err := rows.Scan(
-			&table, &name, &typ, &column, &included, &primary, &uniq, &contype, &pred, &expr, &desc,
-			&nullsfirst, &nullslast, &comment, &options, &opcname, &opcdefault, &opcparams,
+			&table, &name, &typ, &column, &included, &primary, &uniq, &constraints, &pred, &expr,
+			&desc, &nullsfirst, &nullslast, &comment, &options, &opcname, &opcdefault, &opcparams,
 		); err != nil {
 			return fmt.Errorf("postgres: scanning indexes for schema %q: %w", s.Name, err)
 		}
@@ -349,8 +350,14 @@ func (i *inspect) addIndexes(s *schema.Schema, rows *sql.Rows) error {
 			if sqlx.ValidString(comment) {
 				idx.Attrs = append(idx.Attrs, &schema.Comment{Text: comment.String})
 			}
-			if sqlx.ValidString(contype) {
-				idx.Attrs = append(idx.Attrs, &ConType{T: contype.String})
+			if sqlx.ValidString(constraints) {
+				var m map[string]string
+				if err := json.Unmarshal([]byte(constraints.String), &m); err != nil {
+					return fmt.Errorf("postgres: unmarshaling index constraints: %w", err)
+				}
+				for n, t := range m {
+					idx.Attrs = append(idx.Attrs, &Constraint{N: n, T: t})
+				}
 			}
 			if sqlx.ValidString(pred) {
 				idx.Attrs = append(idx.Attrs, &IndexPredicate{P: pred.String})
@@ -735,10 +742,11 @@ type (
 		T string
 	}
 
-	// ConType describes constraint type.
+	// Constraint describes a postgres constraint.
 	// https://postgresql.org/docs/current/catalog-pg-constraint.html
-	ConType struct {
+	Constraint struct {
 		schema.Attr
+		N string // constraint name
 		T string // c, f, p, u, t, x.
 	}
 
@@ -855,7 +863,7 @@ type (
 )
 
 // IsUnique reports if the type is unique constraint.
-func (c ConType) IsUnique() bool { return strings.ToLower(c.T) == "u" }
+func (c Constraint) IsUnique() bool { return strings.ToLower(c.T) == "u" }
 
 // IntegerType returns the underlying integer type this serial type represents.
 func (s *SerialType) IntegerType() *schema.IntegerType {
@@ -1136,35 +1144,44 @@ WHERE
 ORDER BY
 	t1.table_name, t1.ordinal_position
 `
+
 	fksQuery = `
-SELECT
-    t1.constraint_name,
-    t1.table_name,
-    t2.column_name,
-    t1.table_schema,
-    t3.table_name AS referenced_table_name,
-    t3.column_name AS referenced_column_name,
-    t3.table_schema AS referenced_schema_name,
-    t4.update_rule,
-    t4.delete_rule
-FROM
-    information_schema.table_constraints t1
-    JOIN information_schema.key_column_usage t2
-    ON t1.constraint_name = t2.constraint_name
-    AND t1.table_schema = t2.constraint_schema
-    JOIN information_schema.constraint_column_usage t3
-    ON t1.constraint_name = t3.constraint_name
-    AND t1.table_schema = t3.constraint_schema
-    JOIN information_schema.referential_constraints t4
-    ON t1.constraint_name = t4.constraint_name
-    AND t1.table_schema = t4.constraint_schema
-WHERE
-    t1.constraint_type = 'FOREIGN KEY'
-    AND t1.table_schema = $1
-    AND t1.table_name IN (%s)
-ORDER BY
-    t1.constraint_name,
-    t2.ordinal_position
+SELECT 
+    fk.constraint_name,
+    fk.table_name,
+    a1.attname AS column_name,
+    fk.schema_name,
+    fk.referenced_table_name,
+    a2.attname AS referenced_column_name,
+    fk.referenced_schema_name,
+    rc.update_rule,
+    rc.delete_rule
+	FROM 
+	    (
+	    	SELECT
+	      		con.conname AS constraint_name,
+	      		con.conrelid,
+	      		con.confrelid,
+	      		t1.relname AS table_name,
+	      		ns1.nspname AS schema_name,
+      			t2.relname AS referenced_table_name,
+	      		ns2.nspname AS referenced_schema_name,
+	      		unnest(con.conkey) AS conkey,
+	      		unnest(con.confkey) AS confkey
+	    	FROM pg_constraint con
+	    	JOIN pg_class t1 ON t1.oid = con.conrelid
+	    	JOIN pg_class t2 ON t2.oid = con.confrelid
+	    	JOIN pg_namespace ns1 on t1.relnamespace = ns1.oid
+	    	JOIN pg_namespace ns2 on t2.relnamespace = ns2.oid
+	    	WHERE ns1.nspname = $1
+	    	AND t1.relname IN (%s)
+	    	AND con.contype = 'f'
+	) AS fk
+	JOIN pg_attribute a1 ON a1.attnum = fk.conkey AND a1.attrelid = fk.conrelid
+	JOIN pg_attribute a2 ON a2.attnum = fk.confkey AND a2.attrelid = fk.confrelid
+	JOIN information_schema.referential_constraints rc ON rc.constraint_name = fk.constraint_name AND rc.constraint_schema = fk.schema_name
+	ORDER BY
+	    fk.conrelid, fk.constraint_name
 `
 
 	// Query to list table check constraints.
@@ -1206,7 +1223,7 @@ SELECT
 	%s AS included,
 	idx.indisprimary AS primary,
 	idx.indisunique AS unique,
-	c.contype AS constraint_type,
+	con.nametypes AS constraints,
 	pg_get_expr(idx.indpred, idx.indrelid) AS predicate,
 	pg_get_indexdef(idx.indexrelid, idx.ord, false) AS expression,
 	pg_index_column_has_property(idx.indexrelid, idx.ord, 'desc') AS isdesc,
@@ -1228,7 +1245,11 @@ FROM
 	JOIN pg_class i ON i.oid = idx.indexrelid
 	JOIN pg_class t ON t.oid = idx.indrelid
 	JOIN pg_namespace n ON n.oid = t.relnamespace
-	LEFT JOIN pg_constraint c ON idx.indexrelid = c.conindid
+	LEFT JOIN (
+	    select conindid, jsonb_object_agg(conname, contype) AS nametypes
+	    from pg_constraint
+	    group by conindid
+	) con ON con.conindid = idx.indexrelid
 	LEFT JOIN pg_attribute a ON (a.attrelid, a.attnum) = (idx.indrelid, idx.key)
 	JOIN pg_am am ON am.oid = i.relam
 	LEFT JOIN pg_opclass op ON op.oid = idx.indclass[idx.ord-1]
@@ -1236,7 +1257,6 @@ FROM
 WHERE
 	n.nspname = $1
 	AND t.relname IN (%s)
-	AND COALESCE(c.contype, '') <> 'f'
 ORDER BY
 	table_name, index_name, idx.ord
 `
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 15f0d881..93171c01 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -980,7 +980,7 @@ func (s *state) index(b *sqlx.Builder, idx *schema.Index) error {
 	}
 	for _, attr := range idx.Attrs {
 		switch attr.(type) {
-		case *schema.Comment, *ConType, *IndexType, *IndexInclude, *Concurrently, *IndexPredicate, *IndexStorageParams:
+		case *schema.Comment, *IndexType, *IndexInclude, *Concurrently, *Constraint, *IndexPredicate, *IndexStorageParams:
 		default:
 			return fmt.Errorf("postgres: unexpected index attribute: %T", attr)
 		}
@@ -1100,7 +1100,15 @@ func check(b *sqlx.Builder, c *schema.Check) {
 
 // isUniqueConstraint reports if the index is a valid UNIQUE constraint.
 func isUniqueConstraint(i *schema.Index) bool {
-	if c := (ConType{}); !sqlx.Has(i.Attrs, &c) || !c.IsUnique() || !i.Unique {
+	hasC := func() bool {
+		for _, a := range i.Attrs {
+			if c, ok := a.(*Constraint); ok && c.IsUnique() {
+				return true
+			}
+		}
+		return false
+	}()
+	if !hasC || !i.Unique {
 		return false
 	}
 	// UNIQUE constraint cannot use functional indexes,


Test patch:
diff --git a/internal/integration/testdata/postgres/foreign-key.txt b/internal/integration/testdata/postgres/foreign-key.txt
new file mode 100644
index 00000000..ca2da934
--- /dev/null
+++ b/internal/integration/testdata/postgres/foreign-key.txt
@@ -0,0 +1,144 @@
+apply 1.hcl
+cmpshow t1 t1.sql
+cmpshow t2 t2.sql
+cmphcl 1.inspect.hcl
+
+-- 1.hcl --
+schema "$db" {}
+table "t1" {
+  schema = schema.$db
+  column "c1" {
+    null = false
+    type = integer
+  }
+  column "c2" {
+    null = true
+    type = integer
+  }
+  column "c3" {
+    null = true
+    type = integer
+  }
+  primary_key {
+    columns = [column.c1]
+  }
+  index "t1_c2_c3_idx" {
+    unique  = true
+    columns = [column.c2, column.c3]
+    type    = BTREE
+  }
+}
+table "t2" {
+  schema = schema.$db
+  column "c1" {
+    null = false
+    type = integer
+  }
+  column "c2" {
+    null = true
+    type = integer
+  }
+  column "c3" {
+    null = true
+    type = integer
+  }
+  primary_key {
+    columns = [column.c1]
+  }
+  foreign_key "c2_c3_1" {
+    columns     = [column.c2, column.c3]
+    ref_columns = [table.t1.column.c2, table.t1.column.c3]
+    on_update   = NO_ACTION
+    on_delete   = NO_ACTION
+  }
+  foreign_key "c2_c3_2" {
+    columns     = [column.c2, column.c3]
+    ref_columns = [table.t1.column.c2, table.t1.column.c3]
+    on_update   = NO_ACTION
+    on_delete   = NO_ACTION
+  }
+}
+
+-- t1.sql --
+           Table "script_foreign_key.t1"
+ Column |  Type   | Collation | Nullable | Default
+--------+---------+-----------+----------+---------
+ c1     | integer |           | not null |
+ c2     | integer |           |          |
+ c3     | integer |           |          |
+Indexes:
+    "t1_pkey" PRIMARY KEY, btree (c1)
+    "t1_c2_c3_idx" UNIQUE, btree (c2, c3)
+Referenced by:
+    TABLE "script_foreign_key.t2" CONSTRAINT "c2_c3_1" FOREIGN KEY (c2, c3) REFERENCES script_foreign_key.t1(c2, c3)
+    TABLE "script_foreign_key.t2" CONSTRAINT "c2_c3_2" FOREIGN KEY (c2, c3) REFERENCES script_foreign_key.t1(c2, c3)
+
+-- t2.sql --
+           Table "script_foreign_key.t2"
+ Column |  Type   | Collation | Nullable | Default
+--------+---------+-----------+----------+---------
+ c1     | integer |           | not null |
+ c2     | integer |           |          |
+ c3     | integer |           |          |
+Indexes:
+    "t2_pkey" PRIMARY KEY, btree (c1)
+Foreign-key constraints:
+    "c2_c3_1" FOREIGN KEY (c2, c3) REFERENCES script_foreign_key.t1(c2, c3)
+    "c2_c3_2" FOREIGN KEY (c2, c3) REFERENCES script_foreign_key.t1(c2, c3)
+
+-- 1.inspect.hcl --
+table "t1" {
+  schema = schema.script_foreign_key
+  column "c1" {
+    null = false
+    type = integer
+  }
+  column "c2" {
+    null = true
+    type = integer
+  }
+  column "c3" {
+    null = true
+    type = integer
+  }
+  primary_key {
+    columns = [column.c1]
+  }
+  index "t1_c2_c3_idx" {
+    unique  = true
+    columns = [column.c2, column.c3]
+    type    = BTREE
+  }
+}
+table "t2" {
+  schema = schema.script_foreign_key
+  column "c1" {
+    null = false
+    type = integer
+  }
+  column "c2" {
+    null = true
+    type = integer
+  }
+  column "c3" {
+    null = true
+    type = integer
+  }
+  primary_key {
+    columns = [column.c1]
+  }
+  foreign_key "c2_c3_1" {
+    columns     = [column.c2, column.c3]
+    ref_columns = [table.t1.column.c2, table.t1.column.c3]
+    on_update   = NO_ACTION
+    on_delete   = NO_ACTION
+  }
+  foreign_key "c2_c3_2" {
+    columns     = [column.c2, column.c3]
+    ref_columns = [table.t1.column.c2, table.t1.column.c3]
+    on_update   = NO_ACTION
+    on_delete   = NO_ACTION
+  }
+}
+schema "script_foreign_key" {
+}
\ No newline at end of file
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index 94da2c39..8b27f3e2 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -160,12 +160,12 @@ users      | ts          | tsvector            | tsvector  |  NO          |
 				m.ExpectQuery(queryIndexes).
 					WithArgs("public", "users").
 					WillReturnRows(sqltest.Rows(`
-   table_name   |    index_name   | index_type  | column_name | included | primary | unique | constraint_type | predicate             |   expression              | desc | nulls_first | nulls_last | comment   |                 options               |   opclass_name    | opclass_default | opclass_params 
+   table_name   |    index_name   | index_type  | column_name | included | primary | unique |   constraints   | predicate             |   expression              | desc | nulls_first | nulls_last | comment   |                 options               |   opclass_name    | opclass_default | opclass_params 
 ----------------+-----------------+-------------+-------------+----------+---------+--------+-----------------+-----------------------+---------------------------+------+-------------+------------+-----------+---------------------------------------+-------------------+-----------------+----------------
 users           | idx             | hash        |             | f        | f       | f      |                 |                       | "left"((c11)::text, 100)  | t    | t           | f          | boring    |                                       |     int4_ops      |        t        |
 users           | idx1            | btree       |             | f        | f       | f      |                 | (id <> NULL::integer) | "left"((c11)::text, 100)  | t    | t           | f          |           |                                       |     int4_ops      |        t        |
-users           | t1_c1_key       | btree       | c1          | f        | f       | t      | u               |                       | c1                        | t    | t           | f          |           |                                       |     int4_ops      |        t        |
-users           | t1_pkey         | btree       | id          | f        | t       | t      | p               |                       | id                        | t    | f           | f          |           |                                       |     int4_ops      |        t        |
+users           | t1_c1_key       | btree       | c1          | f        | f       | t      | {"name": "u"}   |                       | c1                        | t    | t           | f          |           |                                       |     int4_ops      |        t        |
+users           | t1_pkey         | btree       | id          | f        | t       | t      | {"t_pkey": "p"} |                       | id                        | t    | f           | f          |           |                                       |     int4_ops      |        t        |
 users           | idx4            | btree       | c1          | f        | f       | t      |                 |                       | c1                        | f    | f           | f          |           |                                       |     int4_ops      |        t        |
 users           | idx4            | btree       | id          | f        | f       | t      |                 |                       | id                        | f    | f           | t          |           |                                       |     int4_ops      |        t        |
 users           | idx5            | btree       | c1          | f        | f       | t      |                 |                       | c1                        | f    | f           | f          |           |                                       |     int4_ops      |        t        |
@@ -193,7 +193,7 @@ users           | tsx             | gist        | ts          | f        | f
 				indexes := []*schema.Index{
 					{Name: "idx", Table: t, Attrs: []schema.Attr{&IndexType{T: "hash"}, &schema.Comment{Text: "boring"}}, Parts: []*schema.IndexPart{{SeqNo: 1, X: &schema.RawExpr{X: `"left"((c11)::text, 100)`}, Desc: true, Attrs: []schema.Attr{&IndexColumnProperty{NullsFirst: true}}}}},
 					{Name: "idx1", Table: t, Attrs: []schema.Attr{&IndexType{T: "btree"}, &IndexPredicate{P: `(id <> NULL::integer)`}}, Parts: []*schema.IndexPart{{SeqNo: 1, X: &schema.RawExpr{X: `"left"((c11)::text, 100)`}, Desc: true, Attrs: []schema.Attr{&IndexColumnProperty{NullsFirst: true}}}}},
-					{Name: "t1_c1_key", Unique: true, Table: t, Attrs: []schema.Attr{&IndexType{T: "btree"}, &ConType{T: "u"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[1], Desc: true, Attrs: []schema.Attr{&IndexColumnProperty{NullsFirst: true}}}}},
+					{Name: "t1_c1_key", Unique: true, Table: t, Attrs: []schema.Attr{&IndexType{T: "btree"}, &Constraint{N: "name", T: "u"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[1], Desc: true, Attrs: []schema.Attr{&IndexColumnProperty{NullsFirst: true}}}}},
 					{Name: "idx4", Unique: true, Table: t, Attrs: []schema.Attr{&IndexType{T: "btree"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[1]}, {SeqNo: 2, C: columns[0], Attrs: []schema.Attr{&IndexColumnProperty{NullsLast: true}}}}},
 					{Name: "idx5", Unique: true, Table: t, Attrs: []schema.Attr{&IndexType{T: "btree"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[1]}, {SeqNo: 2, X: &schema.RawExpr{X: `coalesce(parent_id, 0)`}}}},
 					{Name: "idx6", Unique: true, Table: t, Attrs: []schema.Attr{&IndexType{T: "brin"}, &IndexStorageParams{AutoSummarize: true, PagesPerRange: 2}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[1]}}},
@@ -204,7 +204,7 @@ users           | tsx             | gist        | ts          | f        | f
 					Name:   "t1_pkey",
 					Unique: true,
 					Table:  t,
-					Attrs:  []schema.Attr{&IndexType{T: "btree"}, &ConType{T: "p"}},
+					Attrs:  []schema.Attr{&IndexType{T: "btree"}, &Constraint{N: "t_pkey", T: "p"}},
 					Parts:  []*schema.IndexPart{{SeqNo: 1, C: columns[0], Desc: true}},
 				}
 				columns[0].Indexes = append(columns[0].Indexes, pk, indexes[3], indexes[6])
@@ -464,7 +464,7 @@ users       | idx5       | c           | false   | false  |                 | CR
 	}
 	indexes := []*schema.Index{
 		{Name: "idx1", Table: tbl, Attrs: []schema.Attr{&IndexType{T: "btree"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[0]}}},
-		{Name: "idx2", Unique: true, Table: tbl, Attrs: []schema.Attr{&IndexType{T: "btree"}, &ConType{T: "u"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[1]}}},
+		{Name: "idx2", Unique: true, Table: tbl, Attrs: []schema.Attr{&IndexType{T: "btree"}, &Constraint{T: "u"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[1]}}},
 		{Name: "idx3", Table: tbl, Attrs: []schema.Attr{&IndexType{T: "btree"}, &schema.Comment{Text: "boring"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[2], Desc: true}}},
 		{Name: "idx4", Table: tbl, Attrs: []schema.Attr{&IndexType{T: "btree"}, &IndexPredicate{P: `d < 10`}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[3]}}},
 		{Name: "idx5", Table: tbl, Attrs: []schema.Attr{&IndexType{T: "btree"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[0]}, {SeqNo: 2, C: columns[1]}, {SeqNo: 3, C: columns[2]}}},
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index d3b73d18..a6cfaac4 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -145,7 +145,7 @@ func TestPlanChanges(t *testing.T) {
 							&schema.DropIndex{
 								I: schema.NewUniqueIndex("unique_nickname").
 									AddColumns(schema.NewColumn("nickname")).
-									AddAttrs(&ConType{T: "u"}),
+									AddAttrs(&Constraint{T: "u"}),
 							},
 						},
 					}