support import multi datasheet

Author: medcl

csv2sql.yml

@@ -53,6 +53,18 @@ modules:
             - 'INSERT INTO `ifish`.`aquatic_base_info` (`id`, `outer_code`, `common_name`, `scientific_name`, `english_name`, `chinese_name`, `region_name`, `aquatic_category_id`)'
             - 'VALUES (<{id: }>, <{outer_code: }>, <{common_name: }>, <{scientific_name: }>, <{english_name: }>, <{chinese_name: }>, <{region_name: }>, <{aquatic_category_id: }>);'
 
+        - joint: convert_sql
+          enabled: true
+          parameters:
+            sheet_name: aquatic_region
+            data_start_from_index: 1
+            column_name:
+            - id
+            - code
+            - name
+            - is_homemade
+            row_format:
+            - 'INSERT INTO `aquatic_region` (`id`, `code`, `name`, `is_homemade`) VALUES (<{id: }>, <{code: }>, <{name: }>, <{is_homemade: }>);'
 
         - joint: import_sql
           enabled: true

pipelines/convert_sql.go

@@ -22,8 +22,8 @@ import (
 	"github.com/360EntSecGroup-Skylar/excelize"
 	log "github.com/cihub/seelog"
 	"github.com/infinitbyte/framework/core/pipeline"
-	"strings"
 	"github.com/infinitbyte/framework/core/util"
+	"strings"
 )
 
 type ConvertSQLJoint struct {
@@ -36,10 +36,10 @@ func (joint ConvertSQLJoint) Name() string {
 
 const excelBytesKey = "excelBytes"
 
-var rowFormat pipeline.ParaKey = "row_format"
-var sheetName pipeline.ParaKey = "sheet_name"
-var columnName pipeline.ParaKey = "column_name"
-var dataFromIndex pipeline.ParaKey = "data_start_from_index"
+var rowFormatKey pipeline.ParaKey = "row_format"
+var sheetNameKey pipeline.ParaKey = "sheet_name"
+var columnNameKey pipeline.ParaKey = "column_name"
+var dataFromIndexKey pipeline.ParaKey = "data_start_from_index"
 
 var sqlKey pipeline.ParaKey = "sql"
 
@@ -48,7 +48,7 @@ func (joint ConvertSQLJoint) Process(c *pipeline.Context) error {
 	excelBytes := c.MustGetBytes(excelBytesKey)
 	r := bytes.NewReader(excelBytes)
 
-	templates := joint.MustGetStringArray(rowFormat)
+	templates := joint.MustGetStringArray(rowFormatKey)
 	log.Debug("row templates: ", templates)
 
 	xlsx, err := excelize.OpenReader(r)
@@ -59,10 +59,12 @@ func (joint ConvertSQLJoint) Process(c *pipeline.Context) error {
 	sheetMap := xlsx.GetSheetMap()
 	log.Debug("sheets: ", sheetMap)
 
-	colNames := joint.MustGetStringArray(columnName)
-	dataOffset := joint.MustGetInt(dataFromIndex)
+	colNames := joint.MustGetStringArray(columnNameKey)
+	dataOffset := joint.MustGetInt(dataFromIndexKey)
+
+	sheetName := joint.MustGetString(sheetNameKey)
 
-	rows := xlsx.GetRows(joint.MustGetString(sheetName))
+	rows := xlsx.GetRows(sheetName)
 
 	var sqlBuffer bytes.Buffer
 	for offset, row := range rows {
@@ -100,13 +102,19 @@ func (joint ConvertSQLJoint) Process(c *pipeline.Context) error {
 		}
 	}
 
-	c.Set(sqlKey, sqlBuffer.String())
+	sqlMap, ok := c.GetMap(sqlKey)
+	if !ok {
+		sqlMap = map[string]interface{}{}
+	}
+
+	sqlText := sqlBuffer.String()
+	sqlMap[sheetName] = sqlText
+	c.Set(sqlKey, sqlMap)
 
 	log.Trace(sqlBuffer.String())
 	return nil
 }
 
-
 func formatString(str string) string {
 	str = strings.Replace(str, "\"", "", -1)
 	str = strings.Replace(str, "'", "", -1)

pipelines/import_sql.go

@@ -35,52 +35,54 @@ func (joint ImportSQLJoint) Name() string {
 
 func (joint ImportSQLJoint) Process(c *pipeline.Context) error {
 
-	sqlText := c.MustGetString(sqlKey)
+	sqlTextMap := c.MustGetMap(sqlKey)
 
-	log.Debug("start execute: ", sqlText)
+	for k, sqlText := range sqlTextMap {
 
-	db, err := sql.Open("mysql", joint.MustGetString(mysqlConn))
-	if err != nil {
-		log.Error(err)
-		panic(err)
-	}
+		log.Debug("start execute: ", k, ",sql: ", sqlText)
+		db, err := sql.Open("mysql", joint.MustGetString(mysqlConn))
+		if err != nil {
+			log.Error(err)
+			panic(err)
+		}
 
-	tx, err := db.Begin()
-	if err != nil {
-		log.Error(err)
-		panic(err)
-	}
+		tx, err := db.Begin()
+		if err != nil {
+			log.Error(err)
+			panic(err)
+		}
 
-	defer func() {
-		if r := recover(); r != nil {
-			log.Info("the database is rolled back.")
-			err = tx.Rollback()
-			if err != nil {
-				log.Error(err)
-				panic(err)
+		defer func() {
+			if r := recover(); r != nil {
+				log.Info("the database is rolled back.")
+				err = tx.Rollback()
+				if err != nil {
+					log.Error(err)
+					panic(err)
+				}
 			}
+		}()
+
+		//插入数据
+		result, err := tx.Exec(sqlText.(string))
+
+		if err != nil {
+			log.Error(err, result)
+			panic(err)
 		}
-	}()
 
-	//插入数据
-	result, err := tx.Exec(sqlText)
+		rc, _ := result.RowsAffected()
+		l, _ := result.RowsAffected()
 
-	if err != nil {
-		log.Error(err, result)
-		panic(err)
-	}
+		err = tx.Commit()
+		if err != nil {
+			log.Error(err)
+			panic(err)
+		}
 
-	rc, _ := result.RowsAffected()
-	l, _ := result.RowsAffected()
-	log.Debugf("%v rows affected, lastInsertID: %v", rc, l)
+		log.Infof("sql execute success, %v rows affected, lastInsertID: %v", rc, l)
 
-	err = tx.Commit()
-	if err != nil {
-		log.Error(err)
-		panic(err)
 	}
 
-	log.Info("sql execute success")
-
 	return nil
 }

pipelines/logging.go

@@ -15,12 +15,16 @@ func (joint LoggingJoint) Name() string {
 
 func (joint LoggingJoint) Process(c *pipeline.Context) error {
 
-	sql := c.MustGetString(sqlKey)
+	sqlTextMap := c.MustGetMap(sqlKey)
 
 	util.FileAppendNewLine("log/executed_sql.log", "")
 	util.FileAppendNewLine("log/executed_sql.log", time.Now().String())
-	util.FileAppendNewLine("log/executed_sql.log", sql)
-	util.FileAppendNewLine("log/executed_sql.log", "")
+
+	for k, sqlText := range sqlTextMap {
+		util.FileAppendNewLine("log/executed_sql.log", "process data sheet: "+k)
+		util.FileAppendNewLine("log/executed_sql.log", sqlText.(string))
+		util.FileAppendNewLine("log/executed_sql.log", "")
+	}
 
 	return nil
 }

pipelines/read_csv.go

@@ -45,4 +45,3 @@ func (joint ReadCsvJoint) Process(c *pipeline.Context) error {
 
 	return nil
 }
-

pipelines/read_csv_test.go

@@ -63,9 +63,9 @@ func TestSeek(t *testing.T) {
 	e.IsDebug = true
 	joint := ReadCsvJoint{}
 	c := pipeline.Context{}
-	c.Set(rowFormat, []string{"insert into mytable(a,b,c) values(<{colA: }>,<{colB: }>,<{colC: }>);"})
-	c.Set(sheetName, "Sheet1")
-	c.Set(dataFromIndex, 2)
-	c.Set(columnName, []string{"colA", "colB", "colC", "colD", "colE"})
+	c.Set(rowFormatKey, []string{"insert into mytable(a,b,c) values(<{colA: }>,<{colB: }>,<{colC: }>);"})
+	c.Set(sheetNameKey, "Sheet1")
+	c.Set(dataFromIndexKey, 2)
+	c.Set(columnNameKey, []string{"colA", "colB", "colC", "colD", "colE"})
 	joint.Process(&c)
 }