From 93123cdc788051e917a60380eed53536c7c38a02 Mon Sep 17 00:00:00 2001
From: WangXiangUSTC <wx347249478@gmail.com>
Date: Wed, 26 Jun 2019 14:46:22 +0800
Subject: [PATCH] loader: generate db-schema-create.sql file if not exists
 (#186)

---
 loader/loader.go    | 12 +++++++++---
 loader/util.go      | 18 ++++++++++++++++++
 loader/util_test.go | 31 +++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/loader/loader.go b/loader/loader.go
index e9b0547a0e..8a03f22db5 100644
--- a/loader/loader.go
+++ b/loader/loader.go
@@ -683,10 +683,10 @@ func (l *Loader) prepareDbFiles(files map[string]struct{}) error {
 	}
 
 	if schemaFileCount == 0 {
-		return errors.New("invalid mydumper files for there are no `-schema-create.sql` files found")
+		log.Warn("invalid mydumper files for there are no `-schema-create.sql` files found, and will generate later")
 	}
 	if len(l.db2Tables) == 0 {
-		return errors.New("no available `-schema-create.sql` files, check mydumper parameter matches black-white-list in task config")
+		log.Warn("no available `-schema-create.sql` files, check mydumper parameter matches black-white-list in task config, will generate later")
 	}
 
 	return nil
@@ -713,7 +713,13 @@ func (l *Loader) prepareTableFiles(files map[string]struct{}) error {
 		}
 		tables, ok := l.db2Tables[db]
 		if !ok {
-			return errors.Errorf("invalid table schema file, cannot find db - %s", file)
+			log.Warnf("can't find schema create file for db %s, will generate one", db)
+			if err := generateSchemaCreateFile(l.cfg.Dir, db); err != nil {
+				return errors.Trace(err)
+			}
+			l.db2Tables[db] = make(Tables2DataFiles)
+			tables = l.db2Tables[db]
+			l.totalFileCount.Add(1)
 		}
 
 		if _, ok := tables[table]; ok {
diff --git a/loader/util.go b/loader/util.go
index d50b0bf26a..a2c2e420ba 100644
--- a/loader/util.go
+++ b/loader/util.go
@@ -18,8 +18,11 @@ import (
 	"crypto/sha1"
 	"fmt"
 	"os"
+	"path"
 	"path/filepath"
 	"strings"
+
+	"github.com/pingcap/errors"
 )
 
 // CollectDirFiles gets files in path
@@ -73,3 +76,18 @@ func shortSha1(s string) string {
 func percent(a int64, b int64) string {
 	return fmt.Sprintf("%.2f %%", float64(a)/float64(b)*100)
 }
+
+func generateSchemaCreateFile(dir string, schema string) error {
+	file, err := os.Create(path.Join(dir, fmt.Sprintf("%s-schema-create.sql", schema)))
+	if err != nil {
+		return errors.Trace(err)
+	}
+	defer file.Close()
+
+	_, err = fmt.Fprintf(file, "CREATE DATABASE `%s`;\n", escapeName(schema))
+	return errors.Trace(err)
+}
+
+func escapeName(name string) string {
+	return strings.Replace(name, "`", "``", -1)
+}
diff --git a/loader/util_test.go b/loader/util_test.go
index a9c13c0d47..f35f103c8a 100644
--- a/loader/util_test.go
+++ b/loader/util_test.go
@@ -14,6 +14,10 @@
 package loader
 
 import (
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path"
 	"testing"
 
 	. "github.com/pingcap/check"
@@ -61,3 +65,30 @@ func (t *testUtilSuite) TestSQLReplace(c *C) {
 func (t *testUtilSuite) TestShortSha1(c *C) {
 	c.Assert(shortSha1("/tmp/test_sha1_short_6"), Equals, "97b645")
 }
+
+func (t *testUtilSuite) TestGenerateSchemaCreateFile(c *C) {
+	dir := c.MkDir()
+	testCases := []struct {
+		schema    string
+		createSQL string
+	}{
+		{
+			"loader_test",
+			"CREATE DATABASE `loader_test`;\n",
+		}, {
+			"loader`test",
+			"CREATE DATABASE `loader``test`;\n",
+		},
+	}
+	for _, testCase := range testCases {
+		err := generateSchemaCreateFile(dir, testCase.schema)
+		c.Assert(err, IsNil)
+
+		file, err := os.Open(path.Join(dir, fmt.Sprintf("%s-schema-create.sql", testCase.schema)))
+		c.Assert(err, IsNil)
+
+		data, err := ioutil.ReadAll(file)
+		c.Assert(err, IsNil)
+		c.Assert(string(data), Equals, testCase.createSQL)
+	}
+}