Repository: ariga/atlas
Pull request #2029 (2023-08-28 09:03:00)

sql/postgres: scan foreign-key action codes


Patch:
diff --git a/sql/internal/sqlx/sqlx.go b/sql/internal/sqlx/sqlx.go
index d810cf1b..b516cd66 100644
--- a/sql/internal/sqlx/sqlx.go
+++ b/sql/internal/sqlx/sqlx.go
@@ -91,13 +91,35 @@ func ScanStrings(rows *sql.Rows) ([]string, error) {
 	return vs, nil
 }
 
+type (
+	// ScanStringer groups the fmt.Stringer and sql.Scanner interfaces.
+	ScanStringer interface {
+		fmt.Stringer
+		sql.Scanner
+	}
+	// nullString is a sql.NullString that implements the ScanStringer interface.
+	nullString struct{ sql.NullString }
+)
+
+func (s nullString) String() string    { return s.NullString.String }
+func (s *nullString) Scan(v any) error { return s.NullString.Scan(v) }
+
 // SchemaFKs scans the rows and adds the foreign-key to the schema table.
 // Reference elements are added as stubs and should be linked manually by the
 // caller.
 func SchemaFKs(s *schema.Schema, rows *sql.Rows) error {
+	return TypedSchemaFKs[*nullString](s, rows)
+}
+
+// TypedSchemaFKs is a version of SchemaFKs that allows to specify the type of
+// used to scan update and delete actions from the database.
+func TypedSchemaFKs[T ScanStringer](s *schema.Schema, rows *sql.Rows) error {
 	for rows.Next() {
-		var name, table, column, tSchema, refTable, refColumn, refSchema, updateRule, deleteRule string
-		if err := rows.Scan(&name, &table, &column, &tSchema, &refTable, &refColumn, &refSchema, &updateRule, &deleteRule); err != nil {
+		var (
+			updateAction, deleteAction                                   = V(new(T)), V(new(T))
+			name, table, column, tSchema, refTable, refColumn, refSchema string
+		)
+		if err := rows.Scan(&name, &table, &column, &tSchema, &refTable, &refColumn, &refSchema, &updateAction, &deleteAction); err != nil {
 			return err
 		}
 		t, ok := s.Table(table)
@@ -110,8 +132,8 @@ func SchemaFKs(s *schema.Schema, rows *sql.Rows) error {
 				Symbol:   name,
 				Table:    t,
 				RefTable: t,
-				OnDelete: schema.ReferenceOption(deleteRule),
-				OnUpdate: schema.ReferenceOption(updateRule),
+				OnUpdate: schema.ReferenceOption(updateAction.String()),
+				OnDelete: schema.ReferenceOption(deleteAction.String()),
 			}
 			switch {
 			// Self reference.
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 4f2ccbd8..fbcbe411 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -529,7 +529,7 @@ func (i *inspect) fks(ctx context.Context, s *schema.Schema) error {
 		return fmt.Errorf("postgres: querying schema %q foreign keys: %w", s.Name, err)
 	}
 	defer rows.Close()
-	if err := sqlx.SchemaFKs(s, rows); err != nil {
+	if err := sqlx.TypedSchemaFKs[*ReferenceOption](s, rows); err != nil {
 		return fmt.Errorf("postgres: %w", err)
 	}
 	return rows.Err()
@@ -934,8 +934,39 @@ type (
 	Cascade struct {
 		schema.Clause
 	}
+
+	// ReferenceOption describes the ON DELETE and ON UPDATE options for foreign keys.
+	ReferenceOption schema.ReferenceOption
 )
 
+// String implements fmt.Stringer interface.
+func (o ReferenceOption) String() string {
+	return string(o)
+}
+
+// Scan implements sql.Scanner interface.
+func (o *ReferenceOption) Scan(v any) error {
+	var s sql.NullString
+	if err := s.Scan(v); err != nil {
+		return err
+	}
+	switch strings.ToLower(s.String) {
+	case "a":
+		*o = ReferenceOption(schema.NoAction)
+	case "r":
+		*o = ReferenceOption(schema.Restrict)
+	case "c":
+		*o = ReferenceOption(schema.Cascade)
+	case "n":
+		*o = ReferenceOption(schema.SetNull)
+	case "d":
+		*o = ReferenceOption(schema.SetDefault)
+	default:
+		return fmt.Errorf("unknown reference option: %q", s.String)
+	}
+	return nil
+}
+
 // IsUnique reports if the type is unique constraint.
 func (c Constraint) IsUnique() bool { return strings.ToLower(c.T) == "u" }
 
@@ -1267,8 +1298,8 @@ SELECT
     fk.referenced_table_name,
     a2.attname AS referenced_column_name,
     fk.referenced_schema_name,
-    rc.update_rule,
-    rc.delete_rule
+    fk.confupdtype,
+    fk.confdeltype
 	FROM 
 	    (
 	    	SELECT
@@ -1281,7 +1312,9 @@ SELECT
 	      		ns2.nspname AS referenced_schema_name,
 	      		generate_series(1,array_length(con.conkey,1)) as ord,
 	      		unnest(con.conkey) AS conkey,
-	      		unnest(con.confkey) AS confkey
+	      		unnest(con.confkey) AS confkey,
+	      		con.confupdtype,
+	      		con.confdeltype
 	    	FROM pg_constraint con
 	    	JOIN pg_class t1 ON t1.oid = con.conrelid
 	    	JOIN pg_class t2 ON t2.oid = con.confrelid
@@ -1293,7 +1326,6 @@ SELECT
 	) 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, fk.ord
 `


Test patch:
diff --git a/internal/integration/testdata/postgres/foreign-key-action.txt b/internal/integration/testdata/postgres/foreign-key-action.txt
new file mode 100644
index 00000000..4a59b850
--- /dev/null
+++ b/internal/integration/testdata/postgres/foreign-key-action.txt
@@ -0,0 +1,141 @@
+apply 1.hcl
+cmpshow table_a table_a.sql
+cmpshow table_b table_b.sql
+cmpshow table_c table_c.sql
+cmphcl 1.inspect.hcl
+
+-- 1.hcl --
+schema "script_foreign_key_action" {}
+
+table "table_a" {
+  schema = schema.script_foreign_key_action
+  column "id" {
+    type = text
+  }
+  primary_key {
+    columns = [column.id]
+  }
+}
+
+table "table_b" {
+  schema = schema.script_foreign_key_action
+  column "id" {
+    type = text
+  }
+  column "table_a_id" {
+    type = text
+  }
+  primary_key {
+    columns = [column.id]
+  }
+  foreign_key "table_a_fk" {
+    columns     = [column.table_a_id]
+    ref_columns = [table.table_a.column.id]
+    on_delete   = "CASCADE"
+    on_update   = "CASCADE"
+  }
+}
+
+table "table_c" {
+  schema = schema.script_foreign_key_action
+  column "id" {
+    type = text
+  }
+  column "table_a_id" {
+    type = text
+  }
+  primary_key {
+    columns = [column.id]
+  }
+  foreign_key "table_a_fk" {
+    columns     = [column.table_a_id]
+    ref_columns = [table.table_a.column.id]
+  }
+}
+
+-- table_a.sql --
+   Table "script_foreign_key_action.table_a"
+ Column | Type | Collation | Nullable | Default
+--------+------+-----------+----------+---------
+ id     | text |           | not null |
+Indexes:
+    "table_a_pkey" PRIMARY KEY, btree (id)
+Referenced by:
+    TABLE "script_foreign_key_action.table_b" CONSTRAINT "table_a_fk" FOREIGN KEY (table_a_id) REFERENCES script_foreign_key_action.table_a(id) ON UPDATE CASCADE ON DELETE CASCADE
+    TABLE "script_foreign_key_action.table_c" CONSTRAINT "table_a_fk" FOREIGN KEY (table_a_id) REFERENCES script_foreign_key_action.table_a(id)
+
+-- table_b.sql --
+     Table "script_foreign_key_action.table_b"
+   Column   | Type | Collation | Nullable | Default
+------------+------+-----------+----------+---------
+ id         | text |           | not null |
+ table_a_id | text |           | not null |
+Indexes:
+    "table_b_pkey" PRIMARY KEY, btree (id)
+Foreign-key constraints:
+    "table_a_fk" FOREIGN KEY (table_a_id) REFERENCES script_foreign_key_action.table_a(id) ON UPDATE CASCADE ON DELETE CASCADE
+
+-- table_c.sql --
+     Table "script_foreign_key_action.table_c"
+   Column   | Type | Collation | Nullable | Default
+------------+------+-----------+----------+---------
+ id         | text |           | not null |
+ table_a_id | text |           | not null |
+Indexes:
+    "table_c_pkey" PRIMARY KEY, btree (id)
+Foreign-key constraints:
+    "table_a_fk" FOREIGN KEY (table_a_id) REFERENCES script_foreign_key_action.table_a(id)
+
+-- 1.inspect.hcl --
+table "table_a" {
+  schema = schema.script_foreign_key_action
+  column "id" {
+    null = false
+    type = text
+  }
+  primary_key {
+    columns = [column.id]
+  }
+}
+table "table_b" {
+  schema = schema.script_foreign_key_action
+  column "id" {
+    null = false
+    type = text
+  }
+  column "table_a_id" {
+    null = false
+    type = text
+  }
+  primary_key {
+    columns = [column.id]
+  }
+  foreign_key "table_a_fk" {
+    columns     = [column.table_a_id]
+    ref_columns = [table.table_a.column.id]
+    on_update   = CASCADE
+    on_delete   = CASCADE
+  }
+}
+table "table_c" {
+  schema = schema.script_foreign_key_action
+  column "id" {
+    null = false
+    type = text
+  }
+  column "table_a_id" {
+    null = false
+    type = text
+  }
+  primary_key {
+    columns = [column.id]
+  }
+  foreign_key "table_a_fk" {
+    columns     = [column.table_a_id]
+    ref_columns = [table.table_a.column.id]
+    on_update   = NO_ACTION
+    on_delete   = NO_ACTION
+  }
+}
+schema "script_foreign_key_action" {
+}
\ No newline at end of file
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index 2631a186..f95331d1 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -247,13 +247,13 @@ users      | uid         | integer             | int       | NO          |
 				m.ExpectQuery(queryFKs).
 					WithArgs("public", "users").
 					WillReturnRows(sqltest.Rows(`
-constraint_name | table_name | column_name | table_schema | referenced_table_name | referenced_column_name | referenced_schema_name | update_rule | delete_rule
+constraint_name | table_name | column_name | table_schema | referenced_table_name | referenced_column_name | referenced_schema_name | confupdtype | condeltype
 -----------------+------------+-------------+--------------+-----------------------+------------------------+------------------------+-------------+-------------
-multi_column    | users      | id          | public       | t1                    | gid                    | public                 | NO ACTION   | CASCADE
-multi_column    | users      | id          | public       | t1                    | xid                    | public                 | NO ACTION   | CASCADE
-multi_column    | users      | oid         | public       | t1                    | gid                    | public                 | NO ACTION   | CASCADE
-multi_column    | users      | oid         | public       | t1                    | xid                    | public                 | NO ACTION   | CASCADE
-self_reference  | users      | uid         | public       | users                 | id                     | public                 | NO ACTION   | CASCADE
+multi_column    | users      | id          | public       | t1                    | gid                    | public                 | a            | c
+multi_column    | users      | id          | public       | t1                    | xid                    | public                 | a            | c
+multi_column    | users      | oid         | public       | t1                    | gid                    | public                 | a            | c
+multi_column    | users      | oid         | public       | t1                    | xid                    | public                 | a            | c
+self_reference  | users      | uid         | public       | users                 | id                     | public                 | a            | c
 `))
 				m.noChecks()
 				m.noEnums()