From b118897c99cbd90259ab3f31f5aba59482a4de05 Mon Sep 17 00:00:00 2001
From: Shigeki Morimoto <github@morimo.to>
Date: Thu, 11 Oct 2018 12:23:42 +0900
Subject: [PATCH] Add GuessByRegexp.

---
 cmd/hakagi/hakagi.go |  6 +++++-
 src/guess/guess.go   | 34 ++++++++++++++++++++++++++++++----
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/cmd/hakagi/hakagi.go b/cmd/hakagi/hakagi.go
index 0759a73..ac4156f 100644
--- a/cmd/hakagi/hakagi.go
+++ b/cmd/hakagi/hakagi.go
@@ -15,6 +15,10 @@ import (
 var ruleToGuesser = map[string]guess.GuessOption{
 	"primarykey":     guess.GuessByPrimaryKey(),
 	"tableandcolumn": guess.GuessByTableAndColumn(),
+	"regexp": guess.GuessByRegexp([]guess.ForeignKeyRegexp{
+		guess.NewForeignKeyRegexp("^(seller|buyer|invited_friend|from_user|to_user|extracted_user|inviter_user|invitee_user|friend|reporter|reportee|related_user|invited_friend|start|user|last_user)_id$", "users", "id"),
+		guess.NewForeignKeyRegexp("^(require|coupon_attached|child)_item_id$", "items", "id"),
+	}),
 }
 
 func main() {
@@ -24,7 +28,7 @@ func main() {
 	dbPort := flag.Int("dbport", 3306, "database port")
 
 	targets := flag.String("targets", "", "analysing target databases(comma-separated)")
-	rules := flag.String("rules", "primarykey,tableandcolumn", "analysing rules(comma-separated)")
+	rules := flag.String("rules", "primarykey,tableandcolumn,regexp", "analysing rules(comma-separated)")
 
 	flag.Parse()
 
diff --git a/src/guess/guess.go b/src/guess/guess.go
index 04d6ca6..f0215e6 100644
--- a/src/guess/guess.go
+++ b/src/guess/guess.go
@@ -1,7 +1,7 @@
 package guess
 
 import (
-	"strings"
+	"regexp"
 
 	"github.com/jinzhu/inflection"
 	"github.com/syucream/hakagi/src/constraint"
@@ -15,10 +15,23 @@ const (
 
 type GuessOption func(database.Column, string, database.Column) bool
 
+type ForeignKeyRegexp struct {
+	regexp *regexp.Regexp
+	table  string
+	column string
+}
+
+func NewForeignKeyRegexp(pattern string, table string, column string) ForeignKeyRegexp {
+	return ForeignKeyRegexp{regexp.MustCompile(pattern), table, column}
+}
+
 func isAcceptableAsPrimaryKey(columnType, primaryKeyType string) bool {
-	colIsOk := strings.Index(columnType, "int") != -1
-	pkIsOk := strings.Index(primaryKeyType, "int") != -1
-	return colIsOk && pkIsOk && columnType == primaryKeyType
+	/*
+		colIsOk := strings.Index(columnType, "int") != -1
+		pkIsOk := strings.Index(primaryKeyType, "int") != -1
+		return colIsOk && pkIsOk && columnType == primaryKeyType
+	*/
+	return columnType == primaryKeyType
 }
 
 // Recongnize a column thats same name of other table's primary key is a foreign key
@@ -46,6 +59,19 @@ func GuessByTableAndColumn() GuessOption {
 	}
 }
 
+func GuessByRegexp(foreignKeyRegexpList []ForeignKeyRegexp) GuessOption {
+	return func(i database.Column, table string, pk database.Column) bool {
+		for _, foreignKeyRegexp := range foreignKeyRegexpList {
+			if table == foreignKeyRegexp.table &&
+				pk.Name == foreignKeyRegexp.column &&
+				foreignKeyRegexp.regexp.MatchString(i.Name) {
+				return true
+			}
+		}
+		return false
+	}
+}
+
 // GuessConstraints guesses foreign key constraints from primary keys and indexes.
 // NOTE composite primary keys are not supported.
 func GuessConstraints(indexes database.Indexes, primaryKeys database.PrimaryKeys, guessOptions ...GuessOption) []constraint.Constraint {