From b237d5e227bad40cd3b0053c81e3d26f140b7e38 Mon Sep 17 00:00:00 2001
From: martianzhang <martianzhang@hotmail.com>
Date: Sat, 22 Dec 2018 15:58:50 +0800
Subject: [PATCH] -report-type add ast-json, tiast-json

---
 ast/tidb.go      | 34 ++++++++++++++++--------------
 ast/vitess.go    | 54 ++++++++++++++++++++++++++++++++++++++++++++++++
 cmd/soar/soar.go | 22 ++++++++++----------
 3 files changed, 83 insertions(+), 27 deletions(-)
 create mode 100644 ast/vitess.go

diff --git a/ast/tidb.go b/ast/tidb.go
index 8322d88..4dd1f43 100644
--- a/ast/tidb.go
+++ b/ast/tidb.go
@@ -17,11 +17,14 @@
 package ast
 
 import (
+	"encoding/json"
+
 	"github.com/XiaoMi/soar/common"
 
 	"github.com/kr/pretty"
 	"github.com/pingcap/parser"
 	"github.com/pingcap/parser/ast"
+
 	// for pincap parser
 	_ "github.com/pingcap/tidb/types/parser_driver"
 )
@@ -43,20 +46,19 @@ func PrintPrettyStmtNode(sql, charset, collation string) {
 	}
 }
 
-// TiVisitor TODO:
-type TiVisitor struct {
-	EnterFunc func(node ast.Node) bool
-	LeaveFunc func(node ast.Node) bool
-}
-
-// Enter TODO:
-func (visitor *TiVisitor) Enter(n ast.Node) (node ast.Node, skip bool) {
-	skip = visitor.EnterFunc(n)
-	return
-}
-
-// Leave TODO:
-func (visitor *TiVisitor) Leave(n ast.Node) (node ast.Node, ok bool) {
-	ok = visitor.LeaveFunc(n)
-	return
+// StmtNode2JSON TiParse AST tree into json format
+func StmtNode2JSON(sql, charset, collation string) string {
+	var str string
+	tree, err := TiParse(sql, charset, collation)
+	if err != nil {
+		common.Log.Warning(err.Error())
+	} else {
+		b, err := json.MarshalIndent(tree, "", "  ")
+		if err != nil {
+			common.Log.Error(err.Error())
+		} else {
+			str = string(b)
+		}
+	}
+	return str
 }
diff --git a/ast/vitess.go b/ast/vitess.go
new file mode 100644
index 0000000..82906b7
--- /dev/null
+++ b/ast/vitess.go
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2018 Xiaomi, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package ast
+
+import (
+	"encoding/json"
+
+	"github.com/XiaoMi/soar/common"
+
+	"github.com/kr/pretty"
+	"vitess.io/vitess/go/vt/sqlparser"
+)
+
+// PrintPrettyVitessStmtNode print vitess AST struct data
+func PrintPrettyVitessStmtNode(sql string) {
+	tree, err := sqlparser.Parse(sql)
+	if err != nil {
+		common.Log.Warning(err.Error())
+	} else {
+		_, err = pretty.Println(tree)
+		common.LogIfWarn(err, "")
+	}
+}
+
+// VitessStmtNode2JSON vitess AST tree into json format
+func VitessStmtNode2JSON(sql string) string {
+	var str string
+	tree, err := sqlparser.Parse(sql)
+	if err != nil {
+		common.Log.Warning(err.Error())
+	} else {
+		b, err := json.MarshalIndent(tree, "", "  ")
+		if err != nil {
+			common.Log.Error(err.Error())
+		} else {
+			str = string(b)
+		}
+	}
+	return str
+}
diff --git a/cmd/soar/soar.go b/cmd/soar/soar.go
index 9548ee3..e3a9c12 100644
--- a/cmd/soar/soar.go
+++ b/cmd/soar/soar.go
@@ -31,7 +31,6 @@ import (
 	"github.com/kr/pretty"
 	"github.com/percona/go-mysql/query"
 	"github.com/ziutek/mymysql/mysql"
-	"vitess.io/vitess/go/vt/sqlparser"
 )
 
 func main() {
@@ -143,20 +142,21 @@ func main() {
 			fmt.Println(ast.Compress(sql) + common.Config.Delimiter)
 			continue
 		case "ast":
-			// SQL 抽象语法树
-			var tree sqlparser.Statement
-			tree, err = sqlparser.Parse(sql)
-			if err != nil {
-				fmt.Println(err)
-			} else {
-				_, err = pretty.Println(tree)
-				common.LogIfWarn(err, "")
-			}
+			// print vitess AST data struct
+			ast.PrintPrettyVitessStmtNode(sql)
+			continue
+		case "ast-json":
+			// print vitess SQL AST into json format
+			fmt.Println(ast.VitessStmtNode2JSON(sql))
 			continue
 		case "tiast":
-			// TiDB SQL 抽象语法树
+			// print TiDB AST data struct
 			ast.PrintPrettyStmtNode(sql, "", "")
 			continue
+		case "tiast-json":
+			// print TiDB SQL AST into json format
+			fmt.Println(ast.StmtNode2JSON(sql, "", ""))
+			continue
 		case "tokenize":
 			// SQL 切词
 			_, err = pretty.Println(ast.Tokenize(sql))
-- 
GitLab