Merge pull request #489 from k1LoW/fix-fk-parse

Fix parsing foreign key constraints
k1LoW · Jun 16, 2023 · 07c9d27 · 07c9d27
2 parents 0dddf16 + 63e8c49
commit 07c9d27
Show file tree

Hide file tree

Showing 15 changed files with 324 additions and 275 deletions.
diff --git a/drivers/mysql/mysql.go b/drivers/mysql/mysql.go
@@ -15,7 +15,7 @@ import (
 	"github.com/pkg/errors"
 )
 
-var reFK = regexp.MustCompile(`FOREIGN KEY \((.+)\) REFERENCES ([^\s]+)\s?\((.+)\)`)
+var reFK = regexp.MustCompile(`FOREIGN KEY \((.+)\) REFERENCES ([^\s\)]+)\s?\(([^\)]+)\)`)
 var reAI = regexp.MustCompile(` AUTO_INCREMENT=[\d]+`)
 var supportGeneratedColumn = true
 var supportCheckConstraint = true
@@ -498,13 +498,7 @@ WHERE t.table_schema = ?
 		return tableOrderMap[relations[i].Table.Name] < tableOrderMap[relations[j].Table.Name]
 	})
 	for _, r := range relations {
-		result := reFK.FindAllStringSubmatch(r.Def, -1)
-		if len(result) == 0 || len(result[0]) < 4 {
-			return errors.Errorf("can not parse foreign key: %s", r.Def)
-		}
-		strColumns := strings.Split(result[0][1], ", ")
-		strParentTable := result[0][2]
-		strParentColumns := strings.Split(result[0][3], ", ")
+		strColumns, strParentTable, strParentColumns, err := parseFK(r.Def)
 		for _, c := range strColumns {
 			column, err := r.Table.FindColumnByName(c)
 			if err != nil {
@@ -639,3 +633,14 @@ SELECT table_name, table_type, table_comment FROM information_schema.tables WHER
 func convertColumnNullable(str string) bool {
 	return str != "NO"
 }
+
+func parseFK(def string) ([]string, string, []string, error) {
+	result := reFK.FindAllStringSubmatch(def, -1)
+	if len(result) < 1 || len(result[0]) < 4 {
+		return nil, "", nil, errors.Errorf("can not parse foreign key: %s", def)
+	}
+	strColumns := strings.Split(result[0][1], ", ")
+	strParentTable := strings.Trim(result[0][2], `"`)
+	strParentColumns := strings.Split(result[0][3], ", ")
+	return strColumns, strParentTable, strParentColumns, nil
+}
diff --git a/drivers/postgres/postgres.go b/drivers/postgres/postgres.go
@@ -14,7 +14,7 @@ import (
 	"github.com/pkg/errors"
 )
 
-var reFK = regexp.MustCompile(`FOREIGN KEY \((.+)\) REFERENCES ([^\s]+)\s?\((.+)\)`)
+var reFK = regexp.MustCompile(`FOREIGN KEY \((.+)\) REFERENCES ([^\s\)]+)\s?\(([^\)]+)\)`)
 var reVersion = regexp.MustCompile(`([0-9]+(\.[0-9]+)*)`)
 
 // Postgres struct
@@ -57,7 +57,6 @@ func (p *Postgres) Analyze(s *schema.Schema) error {
 		s.Driver.Meta.CurrentSchema = currentSchema.String
 	}
 
-
 	// search_path
 	var searchPaths string
 	pathRows, err := p.db.Query(`SHOW search_path`)
@@ -311,18 +310,9 @@ ORDER BY tgrelid
 
 	// Relations
 	for _, r := range relations {
-		result := reFK.FindAllStringSubmatch(r.Def, -1)
-		if len(result) < 1 || len(result[0]) < 4 {
-			return errors.Errorf("can not parse foreign key: %s", r.Def)
-		}
-		strColumns := []string{}
-		for _, c := range strings.Split(result[0][1], ", ") {
-			strColumns = append(strColumns, strings.ReplaceAll(c, `"`, ""))
-		}
-		strParentTable := strings.ReplaceAll(result[0][2], `"`, "")
-		strParentColumns := []string{}
-		for _, c := range strings.Split(result[0][3], ", ") {
-			strParentColumns = append(strParentColumns, strings.ReplaceAll(c, `"`, ""))
+		strColumns, strParentTable, strParentColumns, err := parseFK(r.Def)
+		if err != nil {
+			return err
 		}
 		for _, c := range strColumns {
 			column, err := r.Table.FindColumnByName(c)
@@ -688,3 +678,20 @@ func convertConstraintType(t string) string {
 		return t
 	}
 }
+
+func parseFK(def string) ([]string, string, []string, error) {
+	result := reFK.FindAllStringSubmatch(def, -1)
+	if len(result) < 1 || len(result[0]) < 4 {
+		return nil, "", nil, errors.Errorf("can not parse foreign key: %s", def)
+	}
+	strColumns := []string{}
+	for _, c := range strings.Split(result[0][1], ", ") {
+		strColumns = append(strColumns, strings.ReplaceAll(c, `"`, ""))
+	}
+	strParentTable := strings.ReplaceAll(result[0][2], `"`, "")
+	strParentColumns := []string{}
+	for _, c := range strings.Split(result[0][3], ", ") {
+		strParentColumns = append(strParentColumns, strings.ReplaceAll(c, `"`, ""))
+	}
+	return strColumns, strParentTable, strParentColumns, nil
+}
diff --git a/drivers/postgres/postgres_test.go b/drivers/postgres/postgres_test.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	"github.com/k1LoW/tbls/schema"
 	_ "github.com/lib/pq"
 	"github.com/xo/dburl"
@@ -68,3 +69,33 @@ func TestInfo(t *testing.T) {
 		t.Errorf("got not empty string.")
 	}
 }
+
+func TestParseFK(t *testing.T) {
+	tests := []struct {
+		in              string
+		wantCols        []string
+		wantParentTable string
+		wantParentCols  []string
+	}{
+		{"FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE", []string{"user_id"}, "users", []string{"id"}},
+		{"FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE SET NULL (user_id)", []string{"user_id"}, "users", []string{"id"}},
+	}
+	for _, tt := range tests {
+		t.Run(tt.in, func(t *testing.T) {
+			gotCols, gotParentTable, gotParentCols, err := parseFK(tt.in)
+			if err != nil {
+				t.Error(err)
+				return
+			}
+			if diff := cmp.Diff(gotCols, tt.wantCols, nil); diff != "" {
+				t.Error(diff)
+			}
+			if gotParentTable != tt.wantParentTable {
+				t.Errorf("got %v want %v", gotParentTable, tt.wantParentTable)
+			}
+			if diff := cmp.Diff(gotParentCols, tt.wantParentCols, nil); diff != "" {
+				t.Error(diff)
+			}
+		})
+	}
+}
diff --git a/drivers/sqlite/sqlite.go b/drivers/sqlite/sqlite.go
@@ -13,7 +13,7 @@ import (
 	"github.com/pkg/errors"
 )
 
-var reFK = regexp.MustCompile(`FOREIGN KEY \((.+)\) REFERENCES ([^\s]+)\s?\((.+)\)`)
+var reFK = regexp.MustCompile(`FOREIGN KEY \((.+)\) REFERENCES ([^\s\)]+)\s?\(([^\)]+)\)`)
 var reFTS = regexp.MustCompile(`(?i)USING\s+fts([34])`)
 
 var shadowTables []string
@@ -362,13 +362,7 @@ SELECT name, sql FROM sqlite_master WHERE type = 'trigger' AND tbl_name = ?;
 
 	// Relations
 	for _, r := range relations {
-		result := reFK.FindAllStringSubmatch(r.Def, -1)
-		if len(result) < 1 || len(result[0]) < 4 {
-			return errors.Errorf("can not parse foreign key: %s", r.Def)
-		}
-		strColumns := strings.Split(result[0][1], ", ")
-		strParentTable := result[0][2]
-		strParentColumns := strings.Split(result[0][3], ", ")
+		strColumns, strParentTable, strParentColumns, err := parseFK(r.Def)
 		for _, c := range strColumns {
 			column, err := r.Table.FindColumnByName(c)
 			if err != nil {
@@ -499,3 +493,14 @@ func contains(s []string, e string) bool {
 	}
 	return false
 }
+
+func parseFK(def string) ([]string, string, []string, error) {
+	result := reFK.FindAllStringSubmatch(def, -1)
+	if len(result) < 1 || len(result[0]) < 4 {
+		return nil, "", nil, errors.Errorf("can not parse foreign key: %s", def)
+	}
+	strColumns := strings.Split(result[0][1], ", ")
+	strParentTable := strings.Trim(result[0][2], `"`)
+	strParentColumns := strings.Split(result[0][3], ", ")
+	return strColumns, strParentTable, strParentColumns, nil
+}
diff --git a/sample/adjust/public.posts.md b/sample/adjust/public.posts.md
@@ -22,7 +22,7 @@ Posts table
 | Name                    | Type        | Definition                                                                                                                                                              | Comment        |
 | ----------------------- | ----------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------- |
 | update_posts_updated    | TRIGGER     | CREATE CONSTRAINT TRIGGER update_posts_updated AFTER INSERT OR UPDATE ON public.posts NOT DEFERRABLE INITIALLY IMMEDIATE FOR EACH ROW EXECUTE FUNCTION update_updated() |                |
-| posts_user_id_fk        | FOREIGN KEY | FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE                                                                                                            | posts -> users |
+| posts_user_id_fk        | FOREIGN KEY | FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE SET NULL (user_id)                                                                                                 | posts -> users |
 | posts_id_pk             | PRIMARY KEY | PRIMARY KEY (id)                                                                                                                                                        |                |
 | posts_user_id_title_key | UNIQUE      | UNIQUE (user_id, title)                                                                                                                                                 |                |