package mysql import ( "context" "database/sql" "fmt" "math" "regexp" "strconv" "strings" "time" "github.com/go-sql-driver/mysql" "gorm.io/gorm" "gorm.io/gorm/callbacks" "gorm.io/gorm/clause" "gorm.io/gorm/logger" "gorm.io/gorm/migrator" "gorm.io/gorm/schema" "gorm.io/gorm/utils" ) const ( DefaultDriverName = "mysql" AutoRandomTag = "auto_random()" // Treated as an auto_random field for tidb ) type Config struct { DriverName string ServerVersion string DSN string DSNConfig *mysql.Config Conn gorm.ConnPool SkipInitializeWithVersion bool DefaultStringSize uint DefaultDatetimePrecision *int DisableWithReturning bool DisableDatetimePrecision bool DontSupportRenameIndex bool DontSupportRenameColumn bool DontSupportForShareClause bool DontSupportNullAsDefaultValue bool DontSupportRenameColumnUnique bool // As of MySQL 8.0.19, ALTER TABLE permits more general (and SQL standard) syntax // for dropping and altering existing constraints of any type. // see https://dev.mysql.com/doc/refman/8.0/en/alter-table.html DontSupportDropConstraint bool } type Dialector struct { *Config } var ( // CreateClauses create clauses CreateClauses = []string{"INSERT", "VALUES", "ON CONFLICT"} // QueryClauses query clauses QueryClauses = []string{} // UpdateClauses update clauses UpdateClauses = []string{"UPDATE", "SET", "WHERE", "ORDER BY", "LIMIT"} // DeleteClauses delete clauses DeleteClauses = []string{"DELETE", "FROM", "WHERE", "ORDER BY", "LIMIT"} defaultDatetimePrecision = 3 ) func Open(dsn string) gorm.Dialector { dsnConf, _ := mysql.ParseDSN(dsn) return &Dialector{Config: &Config{DSN: dsn, DSNConfig: dsnConf}} } func New(config Config) gorm.Dialector { switch { case config.DSN == "" && config.DSNConfig != nil: config.DSN = config.DSNConfig.FormatDSN() case config.DSN != "" && config.DSNConfig == nil: config.DSNConfig, _ = mysql.ParseDSN(config.DSN) } return &Dialector{Config: &config} } func (dialector Dialector) Name() string { return DefaultDriverName } // NowFunc return now func func (dialector Dialector) NowFunc(n int) func() time.Time { return func() time.Time { round := time.Second / time.Duration(math.Pow10(n)) return time.Now().Round(round) } } func (dialector Dialector) Apply(config *gorm.Config) error { if config.NowFunc != nil { return nil } if dialector.DefaultDatetimePrecision == nil { dialector.DefaultDatetimePrecision = &defaultDatetimePrecision } // while maintaining the readability of the code, separate the business logic from // the general part and leave it to the function to do it here. config.NowFunc = dialector.NowFunc(*dialector.DefaultDatetimePrecision) return nil } func (dialector Dialector) Initialize(db *gorm.DB) (err error) { if dialector.DriverName == "" { dialector.DriverName = DefaultDriverName } if dialector.DefaultDatetimePrecision == nil { dialector.DefaultDatetimePrecision = &defaultDatetimePrecision } if dialector.Conn != nil { db.ConnPool = dialector.Conn } else { db.ConnPool, err = sql.Open(dialector.DriverName, dialector.DSN) if err != nil { return err } } withReturning := false if !dialector.Config.SkipInitializeWithVersion { err = db.ConnPool.QueryRowContext(context.Background(), "SELECT VERSION()").Scan(&dialector.ServerVersion) if err != nil { return err } if strings.Contains(dialector.ServerVersion, "MariaDB") { dialector.Config.DontSupportRenameIndex = true dialector.Config.DontSupportRenameColumn = true dialector.Config.DontSupportForShareClause = true dialector.Config.DontSupportNullAsDefaultValue = true withReturning = checkVersion(dialector.ServerVersion, "10.5") } else if strings.HasPrefix(dialector.ServerVersion, "5.6.") { dialector.Config.DontSupportRenameIndex = true dialector.Config.DontSupportRenameColumn = true dialector.Config.DontSupportForShareClause = true dialector.Config.DontSupportDropConstraint = true } else if strings.HasPrefix(dialector.ServerVersion, "5.7.") { dialector.Config.DontSupportRenameColumn = true dialector.Config.DontSupportForShareClause = true dialector.Config.DontSupportDropConstraint = true } else if strings.HasPrefix(dialector.ServerVersion, "5.") { dialector.Config.DisableDatetimePrecision = true dialector.Config.DontSupportRenameIndex = true dialector.Config.DontSupportRenameColumn = true dialector.Config.DontSupportForShareClause = true dialector.Config.DontSupportDropConstraint = true } if strings.Contains(dialector.ServerVersion, "TiDB") { dialector.Config.DontSupportRenameColumnUnique = true } } // register callbacks callbackConfig := &callbacks.Config{ CreateClauses: CreateClauses, QueryClauses: QueryClauses, UpdateClauses: UpdateClauses, DeleteClauses: DeleteClauses, } if !dialector.Config.DisableWithReturning && withReturning { if !utils.Contains(callbackConfig.CreateClauses, "RETURNING") { callbackConfig.CreateClauses = append(callbackConfig.CreateClauses, "RETURNING") } if !utils.Contains(callbackConfig.UpdateClauses, "RETURNING") { callbackConfig.UpdateClauses = append(callbackConfig.UpdateClauses, "RETURNING") } if !utils.Contains(callbackConfig.DeleteClauses, "RETURNING") { callbackConfig.DeleteClauses = append(callbackConfig.DeleteClauses, "RETURNING") } } callbacks.RegisterDefaultCallbacks(db, callbackConfig) for k, v := range dialector.ClauseBuilders() { db.ClauseBuilders[k] = v } return } const ( // ClauseOnConflict for clause.ClauseBuilder ON CONFLICT key ClauseOnConflict = "ON CONFLICT" // ClauseValues for clause.ClauseBuilder VALUES key ClauseValues = "VALUES" // ClauseFor for clause.ClauseBuilder FOR key ClauseFor = "FOR" ) func (dialector Dialector) ClauseBuilders() map[string]clause.ClauseBuilder { clauseBuilders := map[string]clause.ClauseBuilder{ ClauseOnConflict: func(c clause.Clause, builder clause.Builder) { onConflict, ok := c.Expression.(clause.OnConflict) if !ok { c.Build(builder) return } builder.WriteString("ON DUPLICATE KEY UPDATE ") if len(onConflict.DoUpdates) == 0 { if s := builder.(*gorm.Statement).Schema; s != nil { var column clause.Column onConflict.DoNothing = false if s.PrioritizedPrimaryField != nil { column = clause.Column{Name: s.PrioritizedPrimaryField.DBName} } else if len(s.DBNames) > 0 { column = clause.Column{Name: s.DBNames[0]} } if column.Name != "" { onConflict.DoUpdates = []clause.Assignment{{Column: column, Value: column}} } builder.(*gorm.Statement).AddClause(onConflict) } } for idx, assignment := range onConflict.DoUpdates { if idx > 0 { builder.WriteByte(',') } builder.WriteQuoted(assignment.Column) builder.WriteByte('=') if column, ok := assignment.Value.(clause.Column); ok && column.Table == "excluded" { column.Table = "" builder.WriteString("VALUES(") builder.WriteQuoted(column) builder.WriteByte(')') } else { builder.AddVar(builder, assignment.Value) } } }, ClauseValues: func(c clause.Clause, builder clause.Builder) { if values, ok := c.Expression.(clause.Values); ok && len(values.Columns) == 0 { builder.WriteString("VALUES()") return } c.Build(builder) }, } if dialector.Config.DontSupportForShareClause { clauseBuilders[ClauseFor] = func(c clause.Clause, builder clause.Builder) { if values, ok := c.Expression.(clause.Locking); ok && strings.EqualFold(values.Strength, "SHARE") { builder.WriteString("LOCK IN SHARE MODE") return } c.Build(builder) } } return clauseBuilders } func (dialector Dialector) DefaultValueOf(field *schema.Field) clause.Expression { return clause.Expr{SQL: "DEFAULT"} } func (dialector Dialector) Migrator(db *gorm.DB) gorm.Migrator { return Migrator{ Migrator: migrator.Migrator{ Config: migrator.Config{ DB: db, Dialector: dialector, }, }, Dialector: dialector, } } func (dialector Dialector) BindVarTo(writer clause.Writer, stmt *gorm.Statement, v interface{}) { writer.WriteByte('?') } func (dialector Dialector) QuoteTo(writer clause.Writer, str string) { var ( underQuoted, selfQuoted bool continuousBacktick int8 shiftDelimiter int8 ) for _, v := range []byte(str) { switch v { case '`': continuousBacktick++ if continuousBacktick == 2 { writer.WriteString("``") continuousBacktick = 0 } case '.': if continuousBacktick > 0 || !selfQuoted { shiftDelimiter = 0 underQuoted = false continuousBacktick = 0 writer.WriteByte('`') } writer.WriteByte(v) continue default: if shiftDelimiter-continuousBacktick <= 0 && !underQuoted { writer.WriteByte('`') underQuoted = true if selfQuoted = continuousBacktick > 0; selfQuoted { continuousBacktick -= 1 } } for ; continuousBacktick > 0; continuousBacktick -= 1 { writer.WriteString("``") } writer.WriteByte(v) } shiftDelimiter++ } if continuousBacktick > 0 && !selfQuoted { writer.WriteString("``") } writer.WriteByte('`') } type localTimeInterface interface { In(loc *time.Location) time.Time } func (dialector Dialector) Explain(sql string, vars ...interface{}) string { if dialector.DSNConfig != nil && dialector.DSNConfig.Loc != nil { for i, v := range vars { if p, ok := v.(localTimeInterface); ok { func(i int, t localTimeInterface) { defer func() { recover() }() vars[i] = t.In(dialector.DSNConfig.Loc) }(i, p) } } } return logger.ExplainSQL(sql, nil, `'`, vars...) } func (dialector Dialector) DataTypeOf(field *schema.Field) string { switch field.DataType { case schema.Bool: return "boolean" case schema.Int, schema.Uint: return dialector.getSchemaIntAndUnitType(field) case schema.Float: return dialector.getSchemaFloatType(field) case schema.String: return dialector.getSchemaStringType(field) case schema.Time: return dialector.getSchemaTimeType(field) case schema.Bytes: return dialector.getSchemaBytesType(field) default: return dialector.getSchemaCustomType(field) } } func (dialector Dialector) getSchemaFloatType(field *schema.Field) string { if field.Precision > 0 { return fmt.Sprintf("decimal(%d, %d)", field.Precision, field.Scale) } if field.Size <= 32 { return "float" } return "double" } func (dialector Dialector) getSchemaStringType(field *schema.Field) string { size := field.Size if size == 0 { if dialector.DefaultStringSize > 0 { size = int(dialector.DefaultStringSize) } else { hasIndex := field.TagSettings["INDEX"] != "" || field.TagSettings["UNIQUE"] != "" // TEXT, GEOMETRY or JSON column can't have a default value if field.PrimaryKey || field.HasDefaultValue || hasIndex { size = 191 // utf8mb4 } } } if size >= 65536 && size <= int(math.Pow(2, 24)) { return "mediumtext" } if size > int(math.Pow(2, 24)) || size <= 0 { return "longtext" } return fmt.Sprintf("varchar(%d)", size) } func (dialector Dialector) getSchemaTimeType(field *schema.Field) string { if !dialector.DisableDatetimePrecision && field.Precision == 0 && field.TagSettings["PRECISION"] == "" { field.Precision = *dialector.DefaultDatetimePrecision } var precision string if field.Precision > 0 { precision = fmt.Sprintf("(%d)", field.Precision) } if field.NotNull || field.PrimaryKey { return "datetime" + precision } return "datetime" + precision + " NULL" } func (dialector Dialector) getSchemaBytesType(field *schema.Field) string { if field.Size > 0 && field.Size < 65536 { return fmt.Sprintf("varbinary(%d)", field.Size) } if field.Size >= 65536 && field.Size <= int(math.Pow(2, 24)) { return "mediumblob" } return "longblob" } // autoRandomType // field.DataType MUST be `schema.Int` or `schema.Uint` // Judgement logic: // 1. Is PrimaryKey; // 2. Has default value; // 3. Default value is "auto_random()"; // 4. IGNORE the field.Size, it MUST be bigint; // 5. CLEAR the default tag, and return true; // 6. Otherwise, return false. func autoRandomType(field *schema.Field) (bool, string) { if field.PrimaryKey && field.HasDefaultValue && strings.ToLower(strings.TrimSpace(field.DefaultValue)) == AutoRandomTag { field.DefaultValue = "" sqlType := "bigint" if field.DataType == schema.Uint { sqlType += " unsigned" } sqlType += " auto_random" return true, sqlType } return false, "" } func (dialector Dialector) getSchemaIntAndUnitType(field *schema.Field) string { if autoRandom, typeString := autoRandomType(field); autoRandom { return typeString } constraint := func(sqlType string) string { if field.DataType == schema.Uint { sqlType += " unsigned" } if field.AutoIncrement { sqlType += " AUTO_INCREMENT" } return sqlType } switch { case field.Size <= 8: return constraint("tinyint") case field.Size <= 16: return constraint("smallint") case field.Size <= 24: return constraint("mediumint") case field.Size <= 32: return constraint("int") default: return constraint("bigint") } } func (dialector Dialector) getSchemaCustomType(field *schema.Field) string { sqlType := string(field.DataType) if field.AutoIncrement && !strings.Contains(strings.ToLower(sqlType), " auto_increment") { sqlType += " AUTO_INCREMENT" } return sqlType } func (dialector Dialector) SavePoint(tx *gorm.DB, name string) error { return tx.Exec("SAVEPOINT " + name).Error } func (dialector Dialector) RollbackTo(tx *gorm.DB, name string) error { return tx.Exec("ROLLBACK TO SAVEPOINT " + name).Error } // checkVersion newer or equal returns true, old returns false func checkVersion(newVersion, oldVersion string) bool { if newVersion == oldVersion { return true } var ( versionTrimmerRegexp = regexp.MustCompile(`^(\d+).*$`) newVersions = strings.Split(newVersion, ".") oldVersions = strings.Split(oldVersion, ".") ) for idx, nv := range newVersions { if len(oldVersions) <= idx { return true } nvi, _ := strconv.Atoi(versionTrimmerRegexp.ReplaceAllString(nv, "$1")) ovi, _ := strconv.Atoi(versionTrimmerRegexp.ReplaceAllString(oldVersions[idx], "$1")) if nvi == ovi { continue } return nvi > ovi } return false }