table.go 4 KB
package svmysql

import (
	"fmt"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/migrator"
	"strings"
)

func InitTable(obj IMysqlData, suffix string, notify bool) {
	info := obj.MysqlInfo(suffix)
	db := info.DbMysql
	tableName := info.TableName
	m := db.Migrator().(mysql.Migrator).Migrator
	queryTx, execTx := m.GetQueryAndExecTx()
	if HasTable(queryTx.Migrator().(mysql.Migrator).Migrator, obj, tableName) {
		if notify {
			fmt.Println(tableName + " exist, skip!")
		}
	} else {
		err := CreateTable(execTx, obj, tableName)
		if err != nil {
			fmt.Println("err", err)
		}
	}
}

func HasTable(m migrator.Migrator, obj any, tableName string) bool {
	var count int64

	m.RunWithValue(obj, func(stmt *gorm.Statement) error {
		currentDatabase := m.DB.Migrator().CurrentDatabase()
		return m.DB.Raw("SELECT count(*) FROM information_schema.tables WHERE table_schema = ? AND table_name = ? AND table_type = ?", currentDatabase, tableName, "BASE TABLE").Row().Scan(&count)
	})

	return count > 0
}

func CreateTable(execTx *gorm.DB, value any, tableName string) error {
	m := execTx.Migrator().(mysql.Migrator).Migrator
	tx := m.DB.Session(&gorm.Session{})

	if err := m.RunWithValue(value, func(stmt *gorm.Statement) (err error) {
		var (
			createTableSQL          = "CREATE TABLE ? ("
			values                  = []interface{}{clause.Table{Name: tableName}}
			hasPrimaryKeyInDataType bool
		)

		for _, dbName := range stmt.Schema.DBNames {
			field := stmt.Schema.FieldsByDBName[dbName]
			if !field.IgnoreMigration {
				createTableSQL += "? ?"
				hasPrimaryKeyInDataType = hasPrimaryKeyInDataType || strings.Contains(strings.ToUpper(m.DataTypeOf(field)), "PRIMARY KEY")
				values = append(values, clause.Column{Name: dbName}, m.DB.Migrator().FullDataTypeOf(field))
				createTableSQL += ","
			}
		}

		if !hasPrimaryKeyInDataType && len(stmt.Schema.PrimaryFields) > 0 {
			createTableSQL += "PRIMARY KEY ?,"
			primaryKeys := make([]interface{}, 0, len(stmt.Schema.PrimaryFields))
			for _, field := range stmt.Schema.PrimaryFields {
				primaryKeys = append(primaryKeys, clause.Column{Name: field.DBName})
			}

			values = append(values, primaryKeys)
		}

		for _, idx := range stmt.Schema.ParseIndexes() {
			if m.CreateIndexAfterCreateTable {
				defer func(value interface{}, name string) {
					if err == nil {
						err = tx.Migrator().CreateIndex(value, name)
					}
				}(value, idx.Name)
			} else {
				if idx.Class != "" {
					createTableSQL += idx.Class + " "
				}
				createTableSQL += "INDEX ? ?"

				if idx.Comment != "" {
					createTableSQL += fmt.Sprintf(" COMMENT '%s'", idx.Comment)
				}

				if idx.Option != "" {
					createTableSQL += " " + idx.Option
				}

				createTableSQL += ","
				values = append(values, clause.Column{Name: idx.Name}, tx.Migrator().(migrator.BuildIndexOptionsInterface).BuildIndexOptions(idx.Fields, stmt))
			}
		}

		if !m.DB.DisableForeignKeyConstraintWhenMigrating && !m.DB.IgnoreRelationshipsWhenMigrating {
			for _, rel := range stmt.Schema.Relationships.Relations {
				if rel.Field.IgnoreMigration {
					continue
				}
				if constraint := rel.ParseConstraint(); constraint != nil {
					if constraint.Schema == stmt.Schema {
						sql, vars := constraint.Build()
						createTableSQL += sql + ","
						values = append(values, vars...)
					}
				}
			}
		}

		for _, uni := range stmt.Schema.ParseUniqueConstraints() {
			createTableSQL += "CONSTRAINT ? UNIQUE (?),"
			values = append(values, clause.Column{Name: uni.Name}, clause.Expr{SQL: stmt.Quote(uni.Field.DBName)})
		}

		for _, chk := range stmt.Schema.ParseCheckConstraints() {
			createTableSQL += "CONSTRAINT ? CHECK (?),"
			values = append(values, clause.Column{Name: chk.Name}, clause.Expr{SQL: chk.Constraint})
		}

		createTableSQL = strings.TrimSuffix(createTableSQL, ",")

		createTableSQL += ")"

		if tableOption, ok := m.DB.Get("gorm:table_options"); ok {
			createTableSQL += fmt.Sprint(tableOption)
		}

		err = tx.Exec(createTableSQL, values...).Error
		return err
	}); err != nil {
		return err
	}
	return nil
}