Repository: ariga/atlas
Pull request #271 (2021-11-30 14:40:18)

sql/sqlite: extract foreign-key constraint names from create statement

a8m: Thanks for the review @rotemtam and @yonidavidson. I added unit tests for the regex as @yonidavidson recommended.

Patch:
diff --git a/sql/sqlite/driver.go b/sql/sqlite/driver.go
index ee108ace..8a930630 100644
--- a/sql/sqlite/driver.go
+++ b/sql/sqlite/driver.go
@@ -350,7 +350,7 @@ func (d *Driver) fks(ctx context.Context, t *schema.Table) error {
 	if err := d.addFKs(t, rows); err != nil {
 		return fmt.Errorf("sqlite: scan %q foreign-keys: %w", t.Name, err)
 	}
-	return nil
+	return fillConstName(t)
 }
 
 func (d *Driver) addFKs(t *schema.Table, rows *sql.Rows) error {
@@ -566,6 +566,77 @@ func autoinc(t *schema.Table) {
 	t.PrimaryKey.Parts[0].C.Attrs = append(t.PrimaryKey.Parts[0].C.Attrs, AutoIncrement{})
 }
 
+// The following regexes extract named foreign-key constraints defined in the table-constraints or inlined
+// as column-constraints. Note, we assume the SQL statements are valid as they are returned by SQLite.
+var (
+	reConstC = regexp.MustCompile("(?i)(?:[(,]\\s*)[\"`]*(\\w+)[\"`]*[^,]*\\s+CONSTRAINT\\s+[\"`]*(\\w+)[\"`]*\\s+REFERENCES\\s+[\"`]*(\\w+)[\"`]*\\s*\\(([,\"` \\w]+)\\)")
+	reConstT = regexp.MustCompile("(?i)CONSTRAINT\\s+[\"`]*(\\w+)[\"`]*\\s+FOREIGN\\s+KEY\\s*\\(([,\"` \\w]+)\\)\\s+REFERENCES\\s+[\"`]*(\\w+)[\"`]*\\s*\\(([,\"` \\w]+)\\)")
+)
+
+// fillConstName fills foreign-key constrain names from CREATE TABLE statement.
+func fillConstName(t *schema.Table) error {
+	var c CreateStmt
+	if !sqlx.Has(t.Attrs, &c) {
+		return fmt.Errorf("missing CREATE statment for table: %q", t.Name)
+	}
+	// Loop over table constraints.
+	for _, m := range reConstT.FindAllStringSubmatch(c.S, -1) {
+		if len(m) != 5 {
+			return fmt.Errorf("unexpected number of matches for a table constraint: %q", m)
+		}
+		// Pattern matches "constraint_name", "columns", "ref_table" and "ref_columns".
+		for _, fk := range t.ForeignKeys {
+			// Found a foreign-key match for the constraint.
+			if matchFK(fk, columns(m[2]), m[3], columns(m[4])) {
+				fk.Symbol = m[1]
+				break
+			}
+		}
+	}
+	// Loop over inlined column constraints.
+	for _, m := range reConstC.FindAllStringSubmatch(c.S, -1) {
+		if len(m) != 5 {
+			return fmt.Errorf("unexpected number of matches for a column constraint: %q", m)
+		}
+		// Pattern matches "column", "constraint_name", "ref_table" and "ref_columns".
+		for _, fk := range t.ForeignKeys {
+			// Found a foreign-key match for the constraint.
+			if matchFK(fk, columns(m[1]), m[3], columns(m[4])) {
+				fk.Symbol = m[2]
+				break
+			}
+		}
+	}
+	return nil
+}
+
+// columns from the matched regex above.
+func columns(s string) []string {
+	names := strings.Split(s, ",")
+	for i := range names {
+		names[i] = strings.Trim(strings.TrimSpace(names[i]), "`\"")
+	}
+	return names
+}
+
+// matchFK reports if the foreign-key matches the given attributes.
+func matchFK(fk *schema.ForeignKey, columns []string, refTable string, refColumns []string) bool {
+	if len(fk.Columns) != len(columns) || fk.RefTable.Name != refTable || len(fk.RefColumns) != len(refColumns) {
+		return false
+	}
+	for i := range columns {
+		if fk.Columns[i].Name != columns[i] {
+			return false
+		}
+	}
+	for i := range refColumns {
+		if fk.RefColumns[i].Name != refColumns[i] {
+			return false
+		}
+	}
+	return true
+}
+
 const (
 	// Query to list attached database files.
 	databasesQuery = "SELECT `name`, `file` FROM pragma_database_list()"


Test patch:
diff --git a/sql/sqlite/driver_test.go b/sql/sqlite/driver_test.go
index 8cd24f79..a65843af 100644
--- a/sql/sqlite/driver_test.go
+++ b/sql/sqlite/driver_test.go
@@ -161,13 +161,22 @@ func TestDriver_InspectTable(t *testing.T) {
 			name: "table fks",
 			before: func(m mock) {
 				m.systemVars("3.36.0")
-				m.tableExists("users", true, "CREATE TABLE users(id INTEGER PRIMARY KEY)")
+				m.tableExists("users", true, `
+CREATE TABLE users(
+	id INTEGER PRIMARY KEY,
+	c1 int,
+	c2 integer NOT NULL CONSTRAINT c2_fk REFERENCES users (c1) ON DELETE SET NULL,
+	c3 integer NOT NULL REFERENCES users (c1) ON DELETE SET NULL,
+	CONSTRAINT "c1_c2_fk" FOREIGN KEY (c1, c2) REFERENCES t2 (id, c1)
+)
+`)
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(columnsQuery, "users"))).
 					WillReturnRows(sqltest.Rows(`
  name |   type       | nullable | dflt_value  | primary 
 ------+--------------+----------+ ------------+----------
  c1   | int           |  1      |             |  0
  c2   | integer       |  0      |             |  0
+ c3   | integer       |  0      |             |  0
 `))
 				m.noIndexes("users")
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(fksQuery, "users"))).
@@ -182,15 +191,16 @@ func TestDriver_InspectTable(t *testing.T) {
 			expect: func(require *require.Assertions, t *schema.Table, err error) {
 				require.NoError(err)
 				fks := []*schema.ForeignKey{
-					{Symbol: "0", Table: t, OnUpdate: schema.NoAction, OnDelete: schema.Cascade, RefTable: &schema.Table{Name: "t2", Schema: &schema.Schema{Name: "main"}}, RefColumns: []*schema.Column{{Name: "id"}, {Name: "c1"}}},
-					{Symbol: "1", Table: t, OnUpdate: schema.NoAction, OnDelete: schema.Cascade, RefTable: t},
+					{Symbol: "c1_c2_fk", Table: t, OnUpdate: schema.NoAction, OnDelete: schema.Cascade, RefTable: &schema.Table{Name: "t2", Schema: &schema.Schema{Name: "main"}}, RefColumns: []*schema.Column{{Name: "id"}, {Name: "c1"}}},
+					{Symbol: "c2_fk", Table: t, OnUpdate: schema.NoAction, OnDelete: schema.Cascade, RefTable: t},
 				}
 				columns := []*schema.Column{
 					{Name: "c1", Type: &schema.ColumnType{Null: true, Type: &schema.IntegerType{T: "int"}, Raw: "int"}, ForeignKeys: fks[:1]},
 					{Name: "c2", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}, Raw: "integer"}, ForeignKeys: fks},
+					{Name: "c3", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}, Raw: "integer"}},
 				}
-				fks[0].Columns = columns
-				fks[1].Columns = columns[1:]
+				fks[0].Columns = columns[:2]
+				fks[1].Columns = columns[1:2]
 				fks[1].RefColumns = columns[:1]
 				require.Equal(t.Columns, columns)
 				require.Equal(t.ForeignKeys, fks)
@@ -210,6 +220,101 @@ func TestDriver_InspectTable(t *testing.T) {
 	}
 }
 
+func TestRegex_TableConstraint(t *testing.T) {
+	tests := []struct {
+		input   string
+		matches []string
+	}{
+		{
+			input:   `CREATE TABLE pets (id int NOT NULL, owner_id int, CONSTRAINT "owner_fk" FOREIGN KEY(owner_id) REFERENCES users(id))`,
+			matches: []string{"owner_fk", "owner_id", "users", "id"},
+		},
+		{
+			input:   `CREATE TABLE pets (id int NOT NULL, owner_id int, CONSTRAINT "owner_fk" FOREIGN KEY (owner_id) REFERENCES users(id))`,
+			matches: []string{"owner_fk", "owner_id", "users", "id"},
+		},
+		{
+			input: `
+CREATE TABLE pets (
+id int NOT NULL,
+owner_id int,
+CONSTRAINT owner_fk
+	FOREIGN KEY ("owner_id") REFERENCES "users" (id)
+)`,
+			matches: []string{"owner_fk", `"owner_id"`, "users", "id"},
+		},
+		{
+			input: `
+CREATE TABLE pets (
+id int NOT NULL,
+c int,
+d int,
+CONSTRAINT "c_d_fk" FOREIGN KEY (c, d) REFERENCES "users" (a, b)
+)`,
+			matches: []string{"c_d_fk", "c, d", "users", "a, b"},
+		},
+		{
+			input:   `CREATE TABLE pets (id int NOT NULL,c int,d int,CONSTRAINT "c_d_fk" FOREIGN KEY (c, "d") REFERENCES "users" (a, "b"))`,
+			matches: []string{"c_d_fk", `c, "d"`, "users", `a, "b"`},
+		},
+		{
+			input: `CREATE TABLE pets (id int NOT NULL,c int,d int,CONSTRAINT FOREIGN KEY (c, "d") REFERENCES "users" (a, "b"))`,
+		},
+		{
+			input: `CREATE TABLE pets (id int NOT NULL,c int,d int,CONSTRAINT name FOREIGN KEY c REFERENCES "users" (a, "b"))`,
+		},
+		{
+			input: `CREATE TABLE pets (id int NOT NULL,c int,d int,CONSTRAINT name FOREIGN KEY c REFERENCES (a, "b"))`,
+		},
+	}
+	for _, tt := range tests {
+		m := reConstT.FindStringSubmatch(tt.input)
+		require.Equal(t, len(m) != 0, len(tt.matches) != 0)
+		if len(m) > 0 {
+			require.Equal(t, tt.matches, m[1:])
+		}
+	}
+}
+
+func TestRegex_ColumnConstraint(t *testing.T) {
+	tests := []struct {
+		input   string
+		matches []string
+	}{
+		{
+			input:   `CREATE TABLE pets (id int NOT NULL, owner_id int CONSTRAINT "owner_fk" REFERENCES users(id))`,
+			matches: []string{"owner_id", "owner_fk", "users", "id"},
+		},
+		{
+			input:   `CREATE TABLE pets (id int NOT NULL, owner_id int CONSTRAINT "owner_fk" REFERENCES users(id))`,
+			matches: []string{"owner_id", "owner_fk", "users", "id"},
+		},
+		{
+			input: `
+CREATE TABLE pets (
+	id int NOT NULL,
+	c int REFERENCES users(id),
+	d int CONSTRAINT "dfk" REFERENCES users(id)
+)`,
+			matches: []string{"d", "dfk", "users", "id"},
+		},
+		{
+			input: `
+CREATE TABLE t1 (
+	c int REFERENCES users(id),
+	d text CONSTRAINT "dfk" CHECK (d <> '') REFERENCES t2(d)
+)`,
+		},
+	}
+	for _, tt := range tests {
+		m := reConstC.FindStringSubmatch(tt.input)
+		require.Equal(t, len(m) != 0, len(tt.matches) != 0)
+		if len(m) > 0 {
+			require.Equal(t, tt.matches, m[1:])
+		}
+	}
+}
+
 type mock struct {
 	sqlmock.Sqlmock
 }