Repository: ariga/atlas
Pull request #1748 (2023-06-16 11:55:01)

sql/mysql: improve foreign-keys inspection query
`EXPLAIN` before is described here: https://github.com/ariga/atlas/issues/1727#issuecomment-1582487319

`EXPLAIN` after:
```
+----+-------------------+------------+------------+--------+------------------------+------------+---------+-----------------------+------+----------+----------------------------------------------+
| id | select_type       | table      | partitions | type   | possible_keys          | key        | key_len | ref                   | rows | filtered | Extra                                        |
+----+-------------------+------------+------------+--------+------------------------+------------+---------+-----------------------+------+----------+----------------------------------------------+
|  1 | PRIMARY           | cat        | NULL       | index  | PRIMARY                | name       | 194     | NULL                  |    1 |   100.00 | Using index; Using temporary; Using filesort |
|  1 | PRIMARY           | cat        | NULL       | index  | PRIMARY                | name       | 194     | NULL                  |    1 |   100.00 | Using index; Using join buffer (hash join)   |
|  1 | PRIMARY           | sch        | NULL       | ref    | PRIMARY,catalog_id     | catalog_id | 8       | mysql.cat.id          |    5 |   100.00 | Using where; Using index                     |
|  1 | PRIMARY           | tbl        | NULL       | ref    | schema_id              | schema_id  | 8       | mysql.sch.id          |   76 |   100.00 | Rematerialize (<derived3>)                   |
|  1 | PRIMARY           | <derived3> | NULL       | ALL    | NULL                   | NULL       | NULL    | NULL                  |    2 |   100.00 | NULL                                         |
|  1 | PRIMARY           | fk         | NULL       | ALL    | schema_id,table_id     | NULL       | NULL    | NULL                  |  840 |   100.00 | Using where; Using join buffer (hash join)   |
|  1 | PRIMARY           | sch        | NULL       | eq_ref | PRIMARY,catalog_id     | PRIMARY    | 8       | mysql.fk.schema_id    |    1 |   100.00 | Using where                                  |
|  1 | PRIMARY           | tbl        | NULL       | eq_ref | PRIMARY                | PRIMARY    | 8       | mysql.fk.table_id     |    1 |   100.00 | Using where                                  |
|  1 | PRIMARY           | col        | NULL       | eq_ref | PRIMARY                | PRIMARY    | 8       | constraints.column_id |    1 |   100.00 | Using where                                  |
|  3 | DEPENDENT DERIVED | NULL       | NULL       | NULL   | NULL                   | NULL       | NULL    | NULL                  | NULL |     NULL | Impossible WHERE                             |
|  4 | DEPENDENT UNION   | fk         | NULL       | ref    | PRIMARY,table_id       | table_id   | 8       | mysql.tbl.id          |    1 |   100.00 | NULL                                         |
|  4 | DEPENDENT UNION   | fkcu       | NULL       | ref    | PRIMARY,foreign_key_id | PRIMARY    | 8       | mysql.fk.id           |    1 |   100.00 | NULL                                         |
+----+-------------------+------------+------------+--------+------------------------+------------+---------+-----------------------+------+----------+----------------------------------------------+
```


Patch:
diff --git a/cmd/atlas/internal/docker/docker.go b/cmd/atlas/internal/docker/docker.go
index a9828afe..bc894c07 100644
--- a/cmd/atlas/internal/docker/docker.go
+++ b/cmd/atlas/internal/docker/docker.go
@@ -203,7 +203,7 @@ func Out(w io.Writer) ConfigOption {
 //	setup("DROP SCHEMA IF EXISTS public CASCADE")
 func setup(s ...string) ConfigOption {
 	return func(c *Config) error {
-		c.setup = s
+		c.setup = append(c.setup, s...)
 		return nil
 	}
 }
@@ -280,7 +280,7 @@ func (c *Container) Wait(ctx context.Context, timeout time.Duration) error {
 			}
 			for _, s := range c.cfg.setup {
 				if _, err := db.ExecContext(ctx, s); err != nil {
-					_ = db.Close()
+					err = errors.Join(err, db.Close())
 					return fmt.Errorf("%q: %w", s, err)
 				}
 			}
diff --git a/sql/mysql/inspect.go b/sql/mysql/inspect.go
index 6f6ecfd3..2134ecb3 100644
--- a/sql/mysql/inspect.go
+++ b/sql/mysql/inspect.go
@@ -672,7 +672,11 @@ func (i *inspect) marDefaultExpr(c *schema.Column, x string) schema.Expr {
 }
 
 func (i *inspect) querySchema(ctx context.Context, query string, s *schema.Schema) (*sql.Rows, error) {
-	args := []any{s.Name}
+	// Number of times the schema name is parameterized.
+	args := make([]any, strings.Count(query, "?"))
+	for i := range args {
+		args[i] = s.Name
+	}
 	for _, t := range s.Tables {
 		args = append(args, t.Name)
 	}
@@ -721,8 +725,7 @@ WHERE
 	TABLE_SCHEMA IN (%s)
 	AND TABLE_TYPE = 'BASE TABLE'
 ORDER BY
-	TABLE_SCHEMA, TABLE_NAME
-`
+	TABLE_SCHEMA, TABLE_NAME`
 
 	tablesQueryArgs = `
 SELECT
@@ -746,8 +749,7 @@ WHERE
 	AND TABLE_NAME IN (%s)
 	AND TABLE_TYPE = 'BASE TABLE'
 ORDER BY
-	TABLE_SCHEMA, TABLE_NAME
-`
+	TABLE_SCHEMA, TABLE_NAME`
 
 	// Query to list table check constraints.
 	myChecksQuery  = `SELECT t1.TABLE_NAME, t1.CONSTRAINT_NAME, t2.CHECK_CLAUSE, t1.ENFORCED` + checksQuery
@@ -763,8 +765,7 @@ WHERE
 	AND t1.TABLE_SCHEMA = ?
 	AND t1.TABLE_NAME IN (%s)
 ORDER BY
-	t1.CONSTRAINT_NAME
-`
+	t1.CONSTRAINT_NAME`
 
 	// Query to list table foreign keys.
 	fksQuery = `
@@ -776,22 +777,20 @@ SELECT
 	t1.REFERENCED_TABLE_NAME,
 	t1.REFERENCED_COLUMN_NAME,
 	t1.REFERENCED_TABLE_SCHEMA,
-	t3.UPDATE_RULE,
-	t3.DELETE_RULE
+	t2.UPDATE_RULE,
+	t2.DELETE_RULE
 FROM
 	INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS t1
-	JOIN INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS t2
-	JOIN INFORMATION_SCHEMA.REFERENTIAL_CONSTRAINTS AS t3
+	JOIN INFORMATION_SCHEMA.REFERENTIAL_CONSTRAINTS AS t2
 	ON t1.CONSTRAINT_NAME = t2.CONSTRAINT_NAME
-	AND t1.CONSTRAINT_NAME = t3.CONSTRAINT_NAME
-	AND t1.TABLE_SCHEMA = t2.TABLE_SCHEMA
-	AND t1.TABLE_SCHEMA = t3.CONSTRAINT_SCHEMA
 WHERE
-	t2.CONSTRAINT_TYPE = 'FOREIGN KEY'
-	AND t1.TABLE_SCHEMA = ?
+	t1.REFERENCED_COLUMN_NAME IS NOT NULL
+	AND BINARY t1.TABLE_SCHEMA = ?
+	AND BINARY t2.CONSTRAINT_SCHEMA = ?
 	AND t1.TABLE_NAME IN (%s)
 ORDER BY
-	t1.CONSTRAINT_NAME,
+	BINARY t1.TABLE_NAME,
+	BINARY t1.CONSTRAINT_NAME,
 	t1.ORDINAL_POSITION`
 )
 
diff --git a/sql/mysql/tidb.go b/sql/mysql/tidb.go
index 85b23d50..6c1456d2 100644
--- a/sql/mysql/tidb.go
+++ b/sql/mysql/tidb.go
@@ -146,9 +146,6 @@ func (i *tinspect) patchSchema(ctx context.Context, s *schema.Schema) (*schema.S
 		if err := i.setCollate(t); err != nil {
 			return nil, err
 		}
-		if err := i.setFKs(s, t); err != nil {
-			return nil, err
-		}
 		if err := i.setAutoIncrement(t); err != nil {
 			return nil, err
 		}
@@ -182,69 +179,6 @@ func bytesToBitLiteral(b []byte) string {
 	return fmt.Sprintf("b'%b'", val)
 }
 
-// e.g. CONSTRAINT "" FOREIGN KEY ("foo_id") REFERENCES "foo" ("id").
-var reFK = regexp.MustCompile("(?i)CONSTRAINT\\s+[\"`]*(\\w+)[\"`]*\\s+FOREIGN\\s+KEY\\s*\\(([,\"` \\w]+)\\)\\s+REFERENCES\\s+[\"`]*(\\w+)[\"`]*\\s*\\(([,\"` \\w]+)\\).*")
-var reActions = regexp.MustCompile(fmt.Sprintf("(?i)(ON)\\s+(UPDATE|DELETE)\\s+(%s|%s|%s|%s|%s)", schema.NoAction, schema.Restrict, schema.SetNull, schema.SetDefault, schema.Cascade))
-
-func (i *tinspect) setFKs(s *schema.Schema, t *schema.Table) error {
-	var c CreateStmt
-	if !sqlx.Has(t.Attrs, &c) {
-		return fmt.Errorf("missing CREATE TABLE statement in attributes for %q", t.Name)
-	}
-	for _, m := range reFK.FindAllStringSubmatch(c.S, -1) {
-		if len(m) != 5 {
-			return fmt.Errorf("unexpected number of matches for a table constraint: %q", m)
-		}
-		stmt, ctName, clmns, refTableName, refClmns := m[0], m[1], m[2], m[3], m[4]
-		fk := &schema.ForeignKey{
-			Symbol: ctName,
-			Table:  t,
-		}
-		actions := reActions.FindAllStringSubmatch(stmt, 2)
-		for _, actionMatches := range actions {
-			actionType, actionOp := actionMatches[2], actionMatches[3]
-			switch actionType {
-			case "UPDATE":
-				fk.OnUpdate = schema.ReferenceOption(actionOp)
-			case "DELETE":
-				fk.OnDelete = schema.ReferenceOption(actionOp)
-			default:
-				return fmt.Errorf("action type %s is none of 'UPDATE'/'DELETE'", actionType)
-			}
-		}
-		refTable, ok := s.Table(refTableName)
-		if !ok {
-			return fmt.Errorf("couldn't resolve ref table %s on ", m[3])
-		}
-		fk.RefTable = refTable
-		for _, c := range columns(s, clmns) {
-			column, ok := t.Column(c)
-			if !ok {
-				return fmt.Errorf("column %q was not found for fk %q", c, ctName)
-			}
-			fk.Columns = append(fk.Columns, column)
-		}
-		for _, c := range columns(s, refClmns) {
-			column, ok := refTable.Column(c)
-			if !ok {
-				return fmt.Errorf("ref column %q was not found for fk %q", c, ctName)
-			}
-			fk.RefColumns = append(fk.RefColumns, column)
-		}
-		t.ForeignKeys = append(t.ForeignKeys, fk)
-	}
-	return nil
-}
-
-// columns from the matched regex above.
-func columns(_ *schema.Schema, s string) []string {
-	names := strings.Split(s, ",")
-	for i := range names {
-		names[i] = strings.Trim(strings.TrimSpace(names[i]), "`\"")
-	}
-	return names
-}
-
 // e.g CHARSET=utf8mb4 COLLATE=utf8mb4_bin
 var reColl = regexp.MustCompile(`(?i)CHARSET\s*=\s*(\w+)\s*COLLATE\s*=\s*(\w+)`)
 


Test patch:
diff --git a/sql/mysql/inspect_test.go b/sql/mysql/inspect_test.go
index fc7cf001..76952854 100644
--- a/sql/mysql/inspect_test.go
+++ b/sql/mysql/inspect_test.go
@@ -643,7 +643,7 @@ func TestDriver_InspectTable(t *testing.T) {
 `))
 				m.noIndexes()
 				m.ExpectQuery(queryFKs).
-					WithArgs("public", "users").
+					WithArgs("public", "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 |
@@ -840,7 +840,7 @@ func TestDriver_InspectSchema(t *testing.T) {
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(indexesExprQuery, "?, ?"))).
 					WillReturnRows(sqlmock.NewRows([]string{"table_name", "index_name", "column_name", "non_unique", "key_part", "expression"}))
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(fksQuery, "?, ?"))).
-					WithArgs("public", "users", "pets").
+					WithArgs("public", "public", "users", "pets").
 					WillReturnRows(sqltest.Rows(`
 +------------------+------------+-------------+--------------+-----------------------+------------------------+------------------------+-------------+-------------+
 | CONSTRAINT_NAME  | TABLE_NAME | COLUMN_NAME | TABLE_SCHEMA | REFERENCED_TABLE_NAME | REFERENCED_COLUMN_NAME | REFERENCED_SCHEMA_NAME | UPDATE_RULE | DELETE_RULE |