parser package: add TableDBListMap to Result

romberli · Sep 7, 2021 · 665eca9 · 665eca9
1 parent b07a7fd
commit 665eca9
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 23 deletions.
diff --git a/middleware/sql/parser/parser_test.go b/middleware/sql/parser/parser_test.go
@@ -1,6 +1,7 @@
 package parser
 
 import (
+	"encoding/json"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -62,14 +63,21 @@ func TestParser_Parse(t *testing.T) {
 	 PRIMARY KEY (id),
 	 KEY idx_col1_col2_col3 (col1, col2, col3)
 	 ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 ;`
+	// sql = `
+	// 	select *
+	// 	from t01
+	// 			 inner join db1.t01 dt01 on t01.id = dt01.id
+	// 			 inner join t02 on t01.id = t02.id
+	// 			 inner join db2.t02 dt02 on dt01.id = dt02.id
+	// `
 	p := NewParserWithDefault()
 
 	result, err := p.Parse(sql)
 	asst.Nil(err, "test Parse() failed")
 	asst.Equal("t01", result.TableNames[0])
 
 	// print result
-	jsonBytes, err := result.Marshal()
+	jsonBytes, err := json.Marshal(result)
 	asst.Nil(err, "test Parse() failed")
 	t.Log(string(jsonBytes))
 }

diff --git a/middleware/sql/parser/result.go b/middleware/sql/parser/result.go
@@ -1,27 +1,28 @@
 package parser
 
 import (
-	"encoding/json"
-
 	"github.com/romberli/go-util/common"
 	"github.com/romberli/go-util/constant"
 )
 
 type Result struct {
-	SQLType        string            `json:"sql_type"`
-	DBNames        []string          `json:"db_names"`
-	TableNames     []string          `json:"table_names"`
-	TableComments  map[string]string `json:"table_comments"`
-	ColumnNames    []string          `json:"column_names"`
-	ColumnTypes    map[string]string `json:"column_types"`
-	ColumnComments map[string]string `json:"column_comments"`
+	SQLType        string              `json:"sql_type"`
+	TableDBListMap map[string][]string `json:"table_db_list_map"`
+	DBNames        []string            `json:"db_names"`
+	TableNames     []string            `json:"table_names"`
+	TableComments  map[string]string   `json:"table_comments"`
+	ColumnNames    []string            `json:"column_names"`
+	ColumnTypes    map[string]string   `json:"column_types"`
+	ColumnComments map[string]string   `json:"column_comments"`
 }
 
 // NewResult returns a new *Result
-func NewResult(sqlType string, dbNames []string, tableNames []string, tableComments map[string]string,
-	columnNames []string, columnTypes map[string]string, columnComments map[string]string) *Result {
+func NewResult(sqlType string, TableDBListMap map[string][]string, dbNames []string, tableNames []string,
+	tableComments map[string]string, columnNames []string, columnTypes map[string]string,
+	columnComments map[string]string) *Result {
 	return &Result{
 		SQLType:        sqlType,
+		TableDBListMap: TableDBListMap,
 		DBNames:        dbNames,
 		TableNames:     tableNames,
 		TableComments:  tableComments,
@@ -35,6 +36,7 @@ func NewResult(sqlType string, dbNames []string, tableNames []string, tableComme
 func NewEmptyResult() *Result {
 	return &Result{
 		SQLType:        constant.EmptyString,
+		TableDBListMap: make(map[string][]string),
 		DBNames:        []string{},
 		TableNames:     []string{},
 		TableComments:  make(map[string]string),
@@ -49,6 +51,11 @@ func (r *Result) GetSQLType() string {
 	return r.SQLType
 }
 
+// GetTableDBListMap returns table db list map
+func (r *Result) GetTableDBListMap() map[string][]string {
+	return r.TableDBListMap
+}
+
 // GetDBNames returns the db names
 func (r *Result) GetDBNames() []string {
 	return r.DBNames
@@ -84,6 +91,17 @@ func (r *Result) SetSQLType(sqlType string) {
 	r.SQLType = sqlType
 }
 
+// AddTableDBListMap adds db name to the result
+func (r *Result) AddTableDBListMap(tableName string, dbName string) {
+	dbList, ok := r.TableDBListMap[tableName]
+	if !ok {
+		r.TableDBListMap[tableName] = []string{dbName}
+	}
+	if ok && !common.StringInSlice(dbList, dbName) {
+		r.TableDBListMap[tableName] = append(dbList, dbName)
+	}
+}
+
 // AddDBName adds db name to the result
 func (r *Result) AddDBName(dbName string) {
 	if !common.StringInSlice(r.DBNames, dbName) {
@@ -119,8 +137,3 @@ func (r *Result) SetColumnType(columnName string, columnType string) {
 func (r *Result) SetColumnComment(columnName string, columnComment string) {
 	r.ColumnComments[columnName] = columnComment
 }
-
-// Marshal marshals result to json bytes
-func (r *Result) Marshal() ([]byte, error) {
-	return json.Marshal(r)
-}
diff --git a/middleware/sql/parser/visitor.go b/middleware/sql/parser/visitor.go
@@ -123,13 +123,12 @@ func (v *Visitor) Leave(in ast.Node) (out ast.Node, ok bool) {
 
 // visitTableName visits the given node which type is *ast.TableName
 func (v *Visitor) visitTableName(node *ast.TableName) {
-	v.result.AddTableName(node.Name.L)
-
-	// get db name
+	tableName := node.Name.L
 	dbName := node.Schema.L
-	if dbName != constant.EmptyString {
-		v.result.AddDBName(dbName)
-	}
+
+	v.result.AddTableDBListMap(tableName, dbName)
+	v.result.AddDBName(dbName)
+	v.result.AddTableName(tableName)
 }
 
 // visitCreateTableStmt visits the given node which type is *ast.CreateTableStmt