增加sql函数的导入导出 ⭐

Author: Jrohy

README.md

@@ -9,11 +9,12 @@ go语言版mysqldump, goroutine并发导sql, 比navicat等工具导出sql快！
 此项目为: https://github.com/xelabs/go-mydumper 项目的修改优化版
 
 如下变更:
-- 加入mysql source组合源命令(-m), 传参 -m user:pass@host:port 即可定义mysql源, 简化传参
-- 加入排除指定table数据导出sql 命令(-exclude)
 - 修复视图view导出后无法导入的问题
+- 增加function的导入导出
+- 加入mysql source组合源命令(-m), 传参 -m user:pass@host:port, 简化传参
+- 加入排除指定table数据导出sql 命令(-exclude)
 - 优化日志格式和运行时间的显示
-- 合并导入和导出sql功能到同一个文件(-i/-o来分区导入sql还是导出sql)
+- 合并导入和导出sql功能到同一个文件(-i/-o来分区)
 - 提供所有平台的release编译文件
 
 ## 命令行

dumper.go

@@ -19,6 +19,22 @@ func writeDBName(args *common.Args) {
 	_ = common.WriteFile(file, args.Database)
 }
 
+func dumpFunctionSchema(log *xlog.Log, conn *common.Connection, args *common.Args) {
+	qr, err := conn.Fetch(fmt.Sprintf("SELECT `name` FROM mysql.proc WHERE type = 'FUNCTION' AND db = '%s'", args.Database))
+	common.AssertNil(err)
+
+	for _, t := range qr.Rows {
+		function := t[0].String()
+		qr, err := conn.Fetch(fmt.Sprintf("SHOW CREATE FUNCTION `%s`.`%s`", args.Database, function))
+		common.AssertNil(err)
+
+		schema := qr.Rows[0][2].String() + ";\n"
+		file := fmt.Sprintf("%s/%s-schema-function.sql", args.Outdir, function)
+		_ = common.WriteFile(file, schema)
+		log.Info("dumping.function[%s.%s].schema...", args.Database, function)
+	}
+}
+
 func dumpTableSchema(log *xlog.Log, conn *common.Connection, args *common.Args, table string) {
 	qr, err := conn.Fetch(fmt.Sprintf("SHOW CREATE TABLE `%s`.`%s`", args.Database, table))
 	common.AssertNil(err)
@@ -139,10 +155,14 @@ func Dumper(log *xlog.Log, args *common.Args) {
 	// database.
 	conn := pool.Get()
 
-	// tables.
 	var wg sync.WaitGroup
 	var tables []string
 	t := time.Now()
+
+	//function
+	dumpFunctionSchema(log, conn, args)
+
+	//table
 	if args.ExcludeTables != "" {
 		excludeTable = excludeTable + args.ExcludeTables + ","
 	}

loader.go

@@ -16,15 +16,17 @@ import (
 
 // Files tuple.
 type Files struct {
-	schemas []string
-	tables  []string
+	function []string
+	schemas  []string
+	tables   []string
 }
 
 var (
-	viewSuffix   = "-schema-view.sql"
-	schemaSuffix = "-schema.sql"
-	tableSuffix  = ".sql"
-	dbName       = ""
+	functionSuffix = "-schema-function.sql"
+	viewSuffix     = "-schema-view.sql"
+	schemaSuffix   = "-schema.sql"
+	tableSuffix    = ".sql"
+	dbName         = ""
 )
 
 func loadFiles(log *xlog.Log, dir string) *Files {
@@ -39,6 +41,8 @@ func loadFiles(log *xlog.Log, dir string) *Files {
 			switch {
 			case strings.HasSuffix(path, schemaSuffix):
 				files.schemas = append(files.schemas, path)
+			case strings.HasSuffix(path, functionSuffix):
+				files.function = append(files.function, path)
 			case strings.HasSuffix(path, viewSuffix):
 				views = append(views, path)
 			default:
@@ -73,6 +77,24 @@ func restoreDatabaseSchema(log *xlog.Log, db string, conn *common.Connection) {
 	log.Info("restoring.database[%s]", dbName)
 }
 
+func restoreFunctionSchema(log *xlog.Log, functions []string, conn *common.Connection) {
+	for _, function := range functions {
+		name := strings.TrimSuffix(filepath.Base(function), functionSuffix)
+
+		dropQuery := fmt.Sprintf("DROP FUNCTION IF EXISTS %s", name)
+		err := conn.Execute(dropQuery)
+		common.AssertNil(err)
+
+		data, err := common.ReadFile(function)
+		common.AssertNil(err)
+		query := common.BytesToString(data)
+
+		err = conn.Execute(query)
+		common.AssertNil(err)
+		log.Info("restoring.function[%s]", name)
+	}
+}
+
 func restoreTableSchema(log *xlog.Log, overwrite bool, tables []string, conn *common.Connection) {
 	for _, table := range tables {
 		// use
@@ -81,10 +103,7 @@ func restoreTableSchema(log *xlog.Log, overwrite bool, tables []string, conn *co
 
 		log.Info("working.table[%s]", name)
 
-		err := conn.Execute(fmt.Sprintf("USE `%s`", dbName))
-		common.AssertNil(err)
-
-		err = conn.Execute("SET FOREIGN_KEY_CHECKS=0")
+		err := conn.Execute("SET FOREIGN_KEY_CHECKS=0")
 		common.AssertNil(err)
 
 		data, err := common.ReadFile(table)
@@ -146,16 +165,22 @@ func Loader(log *xlog.Log, args *common.Args) {
 	common.AssertNil(err)
 	defer pool.Close()
 
+	t := time.Now()
 	files := loadFiles(log, args.Outdir)
 
-	// database.
 	conn := pool.Get()
+
+	// database.
 	restoreDatabaseSchema(log, args.Database, conn)
-	pool.Put(conn)
 
+	err = conn.Execute(fmt.Sprintf("USE `%s`", dbName))
+	common.AssertNil(err)
+
+	// function.
+	restoreFunctionSchema(log, files.function, conn)
 	// tables.
-	conn = pool.Get()
 	restoreTableSchema(log, args.OverwriteTables, files.schemas, conn)
+
 	pool.Put(conn)
 
 	// Shuffle the tables
@@ -166,7 +191,6 @@ func Loader(log *xlog.Log, args *common.Args) {
 
 	var wg sync.WaitGroup
 	var bytes uint64
-	t := time.Now()
 	for _, table := range files.tables {
 		conn := pool.Get()
 		wg.Add(1)