Fix guesser

Author: syucream

cmd/hakagi/hakagi.go

@@ -7,7 +7,6 @@ import (
 	"strings"
 
 	_ "github.com/go-sql-driver/mysql"
-	"github.com/syucream/hakagi/src/constraint"
 	"github.com/syucream/hakagi/src/database"
 	"github.com/syucream/hakagi/src/formatter"
 	"github.com/syucream/hakagi/src/guess"
@@ -19,7 +18,7 @@ func main() {
 	dbHost := flag.String("dbhost", "localhost", "database host")
 	dbPort := flag.Int("dbport", 3306, "database port")
 
-	targets := flag.String("targets", "", "analysing targets(comma-separated)")
+	targets := flag.String("targets", "", "analysing target databases(comma-separated)")
 
 	flag.Parse()
 
@@ -38,9 +37,7 @@ func main() {
 		log.Fatalf("Failed to fetch primary keys : %v", err)
 	}
 
-	var constraints []constraint.Constraint
-	constraints = append(constraints, guess.GuessByPrimaryKeyName(schemas, primaryKeys)...)
-	constraints = append(constraints, guess.GuessBySimilarColumn(schemas)...)
+	constraints := guess.GuessConstraints(schemas, primaryKeys, guess.GuessByPrimaryKeyName(), guess.GuessBySimilarColumn())
 
 	alterQuery := formatter.FormatSql(constraints)
 	fmt.Println(alterQuery)

src/formatter/formatter.go

@@ -8,15 +8,7 @@ import (
 )
 
 const (
-	baseSql = `ALTER TABLE
-	%s
-ADD CONSTRAINT
-  %s_ibfk_%s%s
-FOREIGN KEY
-  (%s)
-REFERENCES
-  %s(%s);
-`
+	baseSql = "ALTER TABLE %s ADD CONSTRAINT %s_ibfk_%s%s FOREIGN KEY (%s) REFERENCES %s(%s);"
 )
 
 func FormatSql(constraints []constraint.Constraint) string {

src/guess/guess.go

@@ -7,45 +7,45 @@ import (
 )
 
 const (
-	idColumn               = "id"
-	targetPrimaryKeySuffix = "_id"
+	idColumn           = "id"
+	targetColumnSuffix = "_id"
 )
 
-func isTargetPrimaryKey(column string) bool {
-	cLen := len(column)
-	tLen := len(targetPrimaryKeySuffix)
-	return cLen >= tLen && column[cLen-tLen:] == targetPrimaryKeySuffix
-}
-
-func isSimilarColumn(left, right database.Schema) bool {
-	maybeRightTable := inflection.Plural(left.Column)
-	return maybeRightTable == right.Table && right.Column == idColumn
-}
+type guessOption func(database.Schema, database.PrimaryKey) bool
 
 // Recongnize a column thats same name of other table's primary key is a foreign key
 // This base idea refers to SchemaSpy DbAnalyzer:
 //   https://github.com/schemaspy/schemaspy/blob/master/src/main/java/org/schemaspy/DbAnalyzer.java
-func GuessByPrimaryKeyName(schemas []database.Schema, primaryKeys []database.PrimaryKey) []constraint.Constraint {
-	var constraints []constraint.Constraint
+func GuessByPrimaryKeyName() guessOption {
+	return func(s database.Schema, pk database.PrimaryKey) bool {
+		return s.Column == pk.Column && pk.Column != idColumn
+	}
+}
 
-	for _, s := range schemas {
-		for _, pk := range primaryKeys {
-			if s.Table != pk.Table && s.Column == pk.Column && isTargetPrimaryKey(s.Column) {
-				constraints = append(constraints, constraint.Constraint{s.Table, s.Column, pk.Table, pk.Column})
-			}
+func GuessBySimilarColumn() guessOption {
+	return func(s database.Schema, pk database.PrimaryKey) bool {
+		cLen := len(s.Column)
+		tLen := len(targetColumnSuffix)
+
+		if !(cLen >= tLen && s.Column[cLen-tLen:] == targetColumnSuffix) {
+			return false
 		}
-	}
 
-	return constraints
+		return inflection.Plural(s.Column[:cLen-tLen]) == pk.Table && pk.Column == idColumn
+	}
 }
 
-func GuessBySimilarColumn(schemas []database.Schema) []constraint.Constraint {
+func GuessConstraints(schemas []database.Schema, primaryKeys []database.PrimaryKey, guessOptions ...guessOption) []constraint.Constraint {
 	var constraints []constraint.Constraint
 
-	for _, left := range schemas {
-		for _, right := range schemas {
-			if left.Table != right.Table && isSimilarColumn(left, right) {
-				constraints = append(constraints, constraint.Constraint{left.Table, left.Column, right.Table, right.Column})
+	for _, s := range schemas {
+		for _, pk := range primaryKeys {
+			if s.Table != pk.Table {
+				for _, guesser := range guessOptions {
+					if guesser(s, pk) {
+						constraints = append(constraints, constraint.Constraint{s.Table, s.Column, pk.Table, pk.Column})
+					}
+				}
 			}
 		}
 	}