go程序启动时执行安装install初始化数据

目的

程序启动执行install 安装初始化程序，如初始化mysql表，表基础数据

代码

在源目录下创建commands目录，在commands目录下创建install.go

代码如下：

package commands

import (
    "bufio"
    "eastwan.com/xy-icnc/dao"
    "eastwan.com/xy-icnc/model"
    "errors"
    "fmt"
    "github.com/astaxie/beego/logs"
    "github.com/jinzhu/gorm"
    "io/ioutil"
    "os"
    "path"
    "strings"
)

func init() {
    logs.Debug("package commands init()...")
    //获取启动参数
    if len(os.Args) >= 2 {
        switch os.Args[1] {
        case "install":
            //当启动参数的第一个参数为install时
            //如: ./test install
            Install()
            //结束后,终止程序
            os.Exit(0)
            return
        }
    }
    //TODO 模拟操作
    //Install()
    //os.Exit(0)
}
func Install() {
    //初始化mysql
    if err := InitMysql(); err != nil {
        logs.Error("系统初始化mysql失败,%v", err.Error())
    }
}

//在执行参数 install时初始化数据库，表以及表数据
func InitMysql() error {
    var err error
    var gormDB *gorm.DB
    //1.创建数据库
    var optDBName = "mysql"
    var createDBConn *gorm.DB
    var connectParam = dao.DBConfig.MysqlParam + "&multiStatements=true"
    dbConnStr := fmt.Sprintf("%s:%s@(%s)/%s?%s", dao.DBConfig.MysqlUser, dao.DBConfig.MysqlPass, dao.DBConfig.MysqlHost, optDBName, connectParam)
    createDBConn, err = gorm.Open("mysql", dbConnStr)
    if err != nil {
        logs.Error("Mysql init connect error:%v", err.Error())
        return err
    }
    var sqlStr = fmt.Sprintf("CREATE DATABASE %s DEFAULT CHARACTER SET utf8mb4 DEFAULT COLLATE utf8mb4_general_ci;",
        dao.DBConfig.MysqlDBName)
    gormDB = createDBConn.Exec(sqlStr)
    if gormDB.Error != nil {
        logs.Error("create mysql 数据库(%s) 失败,%v", dao.DBConfig.MysqlDBName, gormDB.Error.Error())
        if !strings.Contains(fmt.Sprintf("%s", gormDB.Error.Error()), "database exists") {
            return gormDB.Error
        }
        fmt.Printf("## 当前数据库(%s)已存在,是否仍然执行*.sql脚本,该操作会进行表重建导致数据丢失\n", dao.DBConfig.MysqlDBName)
        fmt.Printf("## 请输入确认?[y/n]:")
        input := bufio.NewScanner(os.Stdin)
        input.Scan()
        if input.Text() != "y" {
            fmt.Printf("\n运行结束......\n")
            return nil
        }
        fmt.Printf("\n")
    }
    //该连接记得关闭，目的是为了创建连接以至于可以创建本项目所需的数据库
    //但是由于操作的数据库是mysql自带的数据库mysql(当前未找到合适的连接数据库)，存在操作风险
    _ = createDBConn.Close()

    //2.创建表以及创建基础数据
    optDBName = dao.DBConfig.MysqlDBName
    var db *gorm.DB
    connectParam = dao.DBConfig.MysqlParam + "&multiStatements=true"
    dbConnStr = fmt.Sprintf("%s:%s@(%s)/%s?%s", dao.DBConfig.MysqlUser, dao.DBConfig.MysqlPass, dao.DBConfig.MysqlHost, optDBName, connectParam)
    db, err = gorm.Open("mysql", dbConnStr)
    if err != nil {
        logs.Error("Mysql init connect error:%v", err.Error())
        return err
    }
    //读取sql目录下的.sql文件并且执行
    sqlDirPath := model.WebConfig.ResourcePath + "/sql"
    fileInfo, err := os.Stat(sqlDirPath)
    if err != nil || !fileInfo.IsDir() {
        return errors.New("sql 脚本 不存在")
    }
    //读取脚本目录下所有的sql脚本
    files, err := ioutil.ReadDir(sqlDirPath)
    if err != nil {
        return err
    }
    if len(files) <= 0 {
        return errors.New("sql 脚本 不存在")
    }

    for _, fileInfo := range files {
        if fileInfo.IsDir() {
            //结束本层循环
            continue
        }
        //获取文件名称带后缀
        fileNameWithSuffix := path.Base(fileInfo.Name())
        //获取文件的后缀(文件类型)
        fileType := path.Ext(fileNameWithSuffix)
        //获取文件名称(不带后缀)
        //fileNameOnly := strings.TrimSuffix(fileNameWithSuffix, fileType)
        if fileType != ".sql" {
            //结束本层循环
            continue
        }
        fileBytes, err := ioutil.ReadFile(path.Join(sqlDirPath, fileNameWithSuffix))
        if err != nil {
            //结束本层循环
            continue
        }
        fmt.Printf("正在执行%s脚本....\n", fileInfo.Name())
        //执行sql文件
        //andExtFromBytes := stripCtlAndExtFromBytes(string(fileBytes))
        gormDB := db.Exec(string(fileBytes))
        if gormDB.Error != nil {
            logs.Error("%s导入失败:%v", fileInfo.Name(), gormDB.Error)
        }
        fmt.Printf("%s脚本执行结束....\n", fileInfo.Name())
    }
    ////mysql初始化结束后重新连接mysql,若该步骤在mysql创建连接之前则可省略
    //dao.InitDBConnect()
    return err
}

备注：.sql脚本文件需放在attachment/sql下，且 .sql脚本的顺序需要特别注意