123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401 |
- package mysql
- import (
- "context"
- "database/sql"
- "fmt"
- "math"
- "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"
- )
- type Config struct {
- DriverName string
- DSN string
- Conn gorm.ConnPool
- SkipInitializeWithVersion bool
- DefaultStringSize uint
- DefaultDatetimePrecision *int
- DisableDatetimePrecision bool
- DontSupportRenameIndex bool
- DontSupportRenameColumn bool
- DontSupportForShareClause bool
- }
- type Dialector struct {
- *Config
- }
- var (
- // UpdateClauses update clauses setting
- UpdateClauses = []string{"UPDATE", "SET", "WHERE", "ORDER BY", "LIMIT"}
- // DeleteClauses delete clauses setting
- DeleteClauses = []string{"DELETE", "FROM", "WHERE", "ORDER BY", "LIMIT"}
- defaultDatetimePrecision = 3
- )
- func Open(dsn string) gorm.Dialector {
- return &Dialector{Config: &Config{DSN: dsn}}
- }
- func New(config Config) gorm.Dialector {
- return &Dialector{Config: &config}
- }
- func (dialector Dialector) Name() string {
- return "mysql"
- }
- // 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().Local().Round(round)
- }
- }
- func (dialector Dialector) Apply(config *gorm.Config) error {
- if config.NowFunc == 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) {
- ctx := context.Background()
- // register callbacks
- callbacks.RegisterDefaultCallbacks(db, &callbacks.Config{
- UpdateClauses: UpdateClauses,
- DeleteClauses: DeleteClauses,
- })
- if dialector.DriverName == "" {
- dialector.DriverName = "mysql"
- }
- 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
- }
- }
- if !dialector.Config.SkipInitializeWithVersion {
- var version string
- err = db.ConnPool.QueryRowContext(ctx, "SELECT VERSION()").Scan(&version)
- if err != nil {
- return err
- }
- if strings.Contains(version, "MariaDB") {
- dialector.Config.DontSupportRenameIndex = true
- dialector.Config.DontSupportRenameColumn = true
- dialector.Config.DontSupportForShareClause = true
- } else if strings.HasPrefix(version, "5.6.") {
- dialector.Config.DontSupportRenameIndex = true
- dialector.Config.DontSupportRenameColumn = true
- dialector.Config.DontSupportForShareClause = true
- } else if strings.HasPrefix(version, "5.7.") {
- dialector.Config.DontSupportRenameColumn = true
- dialector.Config.DontSupportForShareClause = true
- } else if strings.HasPrefix(version, "5.") {
- dialector.Config.DisableDatetimePrecision = true
- dialector.Config.DontSupportRenameIndex = true
- dialector.Config.DontSupportRenameColumn = true
- dialector.Config.DontSupportForShareClause = true
- }
- }
- 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"
- // ClauseValues 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}}
- }
- }
- }
- 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.WriteString("`")
- }
- 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.WriteString("`")
- }
- func (dialector Dialector) Explain(sql string, vars ...interface{}) string {
- 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)
- }
- return string(field.DataType)
- }
- 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 {
- precision := ""
- if !dialector.DisableDatetimePrecision && field.Precision == 0 {
- field.Precision = *dialector.DefaultDatetimePrecision
- }
- 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"
- }
- func (dialector Dialector) getSchemaIntAndUnitType(field *schema.Field) string {
- sqlType := "bigint"
- switch {
- case field.Size <= 8:
- sqlType = "tinyint"
- case field.Size <= 16:
- sqlType = "smallint"
- case field.Size <= 24:
- sqlType = "mediumint"
- case field.Size <= 32:
- sqlType = "int"
- }
- if field.DataType == schema.Uint {
- sqlType += " unsigned"
- }
- if field.AutoIncrement {
- sqlType += " AUTO_INCREMENT"
- }
- return sqlType
- }
- func (dialector Dialector) SavePoint(tx *gorm.DB, name string) error {
- tx.Exec("SAVEPOINT " + name)
- return nil
- }
- func (dialector Dialector) RollbackTo(tx *gorm.DB, name string) error {
- tx.Exec("ROLLBACK TO SAVEPOINT " + name)
- return nil
- }
|