package sqlserver import ( "database/sql" "fmt" "regexp" "strings" "gorm.io/gorm" "gorm.io/gorm/clause" "gorm.io/gorm/logger" "gorm.io/gorm/migrator" "gorm.io/gorm/schema" ) const indexSQL = ` SELECT col.name AS column_name, i.name AS index_name, i.is_unique, i.is_primary_key FROM sys.indexes i LEFT JOIN sys.index_columns ic ON ic.object_id = i.object_id AND ic.index_id = i.index_id LEFT JOIN sys.all_columns col ON col.column_id = ic.column_id AND col.object_id = ic.object_id WHERE i.name IS NOT NULL AND i.is_unique_constraint = 0 AND i.object_id = OBJECT_ID(?) ` type Migrator struct { migrator.Migrator } func (m Migrator) GetTables() (tableList []string, err error) { return tableList, m.DB.Raw("SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_CATALOG = ?", m.CurrentDatabase()).Scan(&tableList).Error } func (m Migrator) CreateTable(values ...interface{}) (err error) { if err = m.Migrator.CreateTable(values...); err != nil { return } for _, value := range m.ReorderModels(values, false) { if err = m.RunWithValue(value, func(stmt *gorm.Statement) (err error) { if stmt.Schema == nil { return } for _, fieldName := range stmt.Schema.DBNames { field := stmt.Schema.FieldsByDBName[fieldName] if _, ok := field.TagSettings["COMMENT"]; !ok { continue } if err = m.setColumnComment(stmt, field, true); err != nil { return } } return }); err != nil { return } } return } func (m Migrator) setColumnComment(stmt *gorm.Statement, field *schema.Field, add bool) error { schemaName := m.getTableSchemaName(stmt.Schema) commentExpr := gorm.Expr(strings.ReplaceAll(field.Comment, "'", "''")) // add field comment if add { return m.DB.Exec( "EXEC sp_addextendedproperty 'MS_Description', N'?', 'SCHEMA', ?, 'TABLE', ?, 'COLUMN', ?", commentExpr, schemaName, stmt.Table, field.DBName, ).Error } // update field comment return m.DB.Exec( "EXEC sp_updateextendedproperty 'MS_Description', N'?', 'SCHEMA', ?, 'TABLE', ?, 'COLUMN', ?", commentExpr, schemaName, stmt.Table, field.DBName, ).Error } func (m Migrator) getTableSchemaName(schema *schema.Schema) string { // return the schema name if it is explicitly provided in the table name // otherwise return default schema name schemaName := getTableSchemaName(schema) if schemaName == "" { schemaName = m.DefaultSchema() } return schemaName } func getTableSchemaName(schema *schema.Schema) string { // return the schema name if it is explicitly provided in the table name // otherwise return a sql wildcard -> use any table_schema if schema == nil || !strings.Contains(schema.Table, ".") { return "" } _, schemaName, _ := splitFullQualifiedName(schema.Table) return schemaName } func splitFullQualifiedName(name string) (string, string, string) { nameParts := strings.Split(name, ".") if len(nameParts) == 1 { // [table_name] return "", "", nameParts[0] } else if len(nameParts) == 2 { // [table_schema].[table_name] return "", nameParts[0], nameParts[1] } else if len(nameParts) == 3 { // [table_catalog].[table_schema].[table_name] return nameParts[0], nameParts[1], nameParts[2] } return "", "", "" } func getFullQualifiedTableName(stmt *gorm.Statement) string { fullQualifiedTableName := stmt.Table if schemaName := getTableSchemaName(stmt.Schema); schemaName != "" { fullQualifiedTableName = schemaName + "." + fullQualifiedTableName } return fullQualifiedTableName } func (m Migrator) HasTable(value interface{}) bool { var count int _ = m.RunWithValue(value, func(stmt *gorm.Statement) error { schemaName := getTableSchemaName(stmt.Schema) if schemaName == "" { schemaName = "%" } return m.DB.Raw( "SELECT count(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = ? AND TABLE_CATALOG = ? and TABLE_SCHEMA like ? AND TABLE_TYPE = ?", stmt.Table, m.CurrentDatabase(), schemaName, "BASE TABLE", ).Row().Scan(&count) }) return count > 0 } func (m Migrator) DropTable(values ...interface{}) error { values = m.ReorderModels(values, false) for i := len(values) - 1; i >= 0; i-- { tx := m.DB.Session(&gorm.Session{}) if err := m.RunWithValue(values[i], func(stmt *gorm.Statement) error { type constraint struct { Name string Parent string } var constraints []constraint err := tx.Raw("SELECT name, OBJECT_NAME(parent_object_id) as parent FROM sys.foreign_keys WHERE referenced_object_id = object_id(?)", getFullQualifiedTableName(stmt)).Scan(&constraints).Error for _, c := range constraints { if err == nil { err = tx.Exec("ALTER TABLE ? DROP CONSTRAINT ?;", gorm.Expr(c.Parent), gorm.Expr(c.Name)).Error } } if err == nil { err = tx.Exec("DROP TABLE IF EXISTS ?", clause.Table{Name: stmt.Table}).Error } return err }); err != nil { return err } } return nil } func (m Migrator) RenameTable(oldName, newName interface{}) error { var oldTable, newTable string if v, ok := oldName.(string); ok { oldTable = v } else { stmt := &gorm.Statement{DB: m.DB} if err := stmt.Parse(oldName); err == nil { oldTable = stmt.Table } else { return err } } if v, ok := newName.(string); ok { newTable = v } else { stmt := &gorm.Statement{DB: m.DB} if err := stmt.Parse(newName); err == nil { newTable = stmt.Table } else { return err } } return m.DB.Exec( "sp_rename @objname = ?, @newname = ?;", clause.Table{Name: oldTable}, clause.Table{Name: newTable}, ).Error } func (m Migrator) AddColumn(value interface{}, name string) error { if err := m.Migrator.AddColumn(value, name); err != nil { return err } return m.RunWithValue(value, func(stmt *gorm.Statement) (err error) { if stmt.Schema != nil { if field := stmt.Schema.LookUpField(name); field != nil { if _, ok := field.TagSettings["COMMENT"]; !ok { return } if err = m.setColumnComment(stmt, field, true); err != nil { return } } } return }) } func (m Migrator) HasColumn(value interface{}, field string) bool { var count int64 _ = m.RunWithValue(value, func(stmt *gorm.Statement) error { currentDatabase := m.DB.Migrator().CurrentDatabase() name := field if stmt.Schema != nil { if field := stmt.Schema.LookUpField(field); field != nil { name = field.DBName } } return m.DB.Raw( "SELECT count(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_CATALOG = ? AND TABLE_NAME = ? AND COLUMN_NAME = ?", currentDatabase, stmt.Table, name, ).Row().Scan(&count) }) return count > 0 } func (m Migrator) AlterColumn(value interface{}, field string) error { return m.RunWithValue(value, func(stmt *gorm.Statement) error { if stmt.Schema != nil { if field := stmt.Schema.LookUpField(field); field != nil { fieldType := clause.Expr{SQL: m.DataTypeOf(field)} if field.NotNull { fieldType.SQL += " NOT NULL" } else { fieldType.SQL += " NULL" } return m.DB.Exec( "ALTER TABLE ? ALTER COLUMN ? ?", clause.Table{Name: getFullQualifiedTableName(stmt)}, clause.Column{Name: field.DBName}, fieldType, ).Error } } return fmt.Errorf("failed to look up field with name: %s", field) }) } func (m Migrator) RenameColumn(value interface{}, oldName, newName string) error { return m.RunWithValue(value, func(stmt *gorm.Statement) error { if stmt.Schema != nil { if field := stmt.Schema.LookUpField(oldName); field != nil { oldName = field.DBName } if field := stmt.Schema.LookUpField(newName); field != nil { newName = field.DBName } } return m.DB.Exec( "sp_rename @objname = ?, @newname = ?, @objtype = 'COLUMN';", fmt.Sprintf("%s.%s", stmt.Table, oldName), clause.Column{Name: newName}, ).Error }) } func (m Migrator) GetColumnComment(stmt *gorm.Statement, fieldDBName string) (comment sql.NullString) { queryTx := m.DB.Session(&gorm.Session{Logger: m.DB.Logger.LogMode(logger.Warn)}) if m.DB.DryRun { queryTx.DryRun = false } queryTx.Raw("SELECT value FROM [?].sys.fn_listextendedproperty('MS_Description', 'SCHEMA', ?, 'TABLE', ?, 'COLUMN', ?)", gorm.Expr(m.CurrentDatabase()), m.getTableSchemaName(stmt.Schema), stmt.Table, fieldDBName).Scan(&comment) return } func (m Migrator) MigrateColumn(value interface{}, field *schema.Field, columnType gorm.ColumnType) error { if err := m.Migrator.MigrateColumn(value, field, columnType); err != nil { return err } return m.RunWithValue(value, func(stmt *gorm.Statement) (err error) { comment := m.GetColumnComment(stmt, field.DBName) if field.Comment != comment.String { if comment.Valid { err = m.setColumnComment(stmt, field, false) } else { err = m.setColumnComment(stmt, field, true) } } return }) } var defaultValueTrimRegexp = regexp.MustCompile("^\\('?([^']*)'?\\)$") // ColumnTypes return columnTypes []gorm.ColumnType and execErr error func (m Migrator) ColumnTypes(value interface{}) ([]gorm.ColumnType, error) { columnTypes := make([]gorm.ColumnType, 0) execErr := m.RunWithValue(value, func(stmt *gorm.Statement) (err error) { rows, err := m.DB.Session(&gorm.Session{}).Table(getFullQualifiedTableName(stmt)).Limit(1).Rows() if err != nil { return err } rawColumnTypes, _ := rows.ColumnTypes() _ = rows.Close() { _, schemaName, tableName := splitFullQualifiedName(stmt.Table) query := strings.TrimSpace(` SELECT COLUMN_NAME, DATA_TYPE, COLUMN_DEFAULT, c.IS_NULLABLE, CHARACTER_MAXIMUM_LENGTH, NUMERIC_PRECISION, NUMERIC_PRECISION_RADIX, NUMERIC_SCALE, DATETIME_PRECISION, AUTO_INCREMENT = c2.is_identity FROM INFORMATION_SCHEMA.COLUMNS c LEFT JOIN sys.tables t ON c.TABLE_NAME = t.[name] LEFT JOIN sys.columns c2 ON t.object_id = c2.object_id AND c2.[name] = c.COLUMN_NAME WHERE TABLE_CATALOG = ? AND TABLE_NAME = ?`) queryParameters := []interface{}{m.CurrentDatabase(), tableName} if schemaName != "" { query += " AND TABLE_SCHEMA = ?" queryParameters = append(queryParameters, schemaName) } var ( columnTypeSQL = query columns, rowErr = m.DB.Raw(columnTypeSQL, queryParameters...).Rows() ) if rowErr != nil { return rowErr } for columns.Next() { var ( column = migrator.ColumnType{ PrimaryKeyValue: sql.NullBool{Valid: true}, UniqueValue: sql.NullBool{Valid: true}, } datetimePrecision sql.NullInt64 radixValue sql.NullInt64 nullableValue sql.NullString autoIncrementValue sql.NullBool values = []interface{}{ &column.NameValue, &column.ColumnTypeValue, &column.DefaultValueValue, &nullableValue, &column.LengthValue, &column.DecimalSizeValue, &radixValue, &column.ScaleValue, &datetimePrecision, &autoIncrementValue, } ) if scanErr := columns.Scan(values...); scanErr != nil { return scanErr } if nullableValue.Valid { column.NullableValue = sql.NullBool{Bool: strings.EqualFold(nullableValue.String, "YES"), Valid: true} } if datetimePrecision.Valid { column.DecimalSizeValue = datetimePrecision } if autoIncrementValue.Valid && autoIncrementValue.Bool { column.AutoIncrementValue = autoIncrementValue } if column.DefaultValueValue.Valid { matches := defaultValueTrimRegexp.FindStringSubmatch(column.DefaultValueValue.String) for len(matches) > 1 { column.DefaultValueValue.String = matches[1] matches = defaultValueTrimRegexp.FindStringSubmatch(column.DefaultValueValue.String) } } for _, c := range rawColumnTypes { if c.Name() == column.NameValue.String { column.SQLColumnType = c break } } columnTypes = append(columnTypes, column) } _ = columns.Close() } { _, schemaName, tableName := splitFullQualifiedName(stmt.Table) query := "SELECT c.COLUMN_NAME, t.CONSTRAINT_TYPE FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS t JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE c ON c.CONSTRAINT_NAME=t.CONSTRAINT_NAME WHERE t.CONSTRAINT_TYPE IN ('PRIMARY KEY', 'UNIQUE') AND c.TABLE_CATALOG = ? AND c.TABLE_NAME = ?" queryParameters := []interface{}{m.CurrentDatabase(), tableName} if schemaName != "" { query += " AND c.TABLE_SCHEMA = ?" queryParameters = append(queryParameters, schemaName) } columnTypeRows, err := m.DB.Raw(query, queryParameters...).Rows() if err != nil { return err } for columnTypeRows.Next() { var name, columnType string _ = columnTypeRows.Scan(&name, &columnType) for idx, c := range columnTypes { mc := c.(migrator.ColumnType) if mc.NameValue.String == name { switch columnType { case "PRIMARY KEY": mc.PrimaryKeyValue = sql.NullBool{Bool: true, Valid: true} case "UNIQUE": mc.UniqueValue = sql.NullBool{Bool: true, Valid: true} } columnTypes[idx] = mc break } } } _ = columnTypeRows.Close() } return }) return columnTypes, execErr } func (m Migrator) CreateView(name string, option gorm.ViewOption) error { if option.Query == nil { return gorm.ErrSubQueryRequired } sql := new(strings.Builder) sql.WriteString("CREATE ") if option.Replace { sql.WriteString("OR ALTER ") } sql.WriteString("VIEW ") m.QuoteTo(sql, name) sql.WriteString(" AS ") m.DB.Statement.AddVar(sql, option.Query) if option.CheckOption != "" { sql.WriteString(" ") sql.WriteString(option.CheckOption) } return m.DB.Exec(m.Explain(sql.String(), m.DB.Statement.Vars...)).Error } func (m Migrator) CreateIndex(value interface{}, name string) error { return m.RunWithValue(value, func(stmt *gorm.Statement) error { var idx *schema.Index if stmt.Schema != nil { idx = stmt.Schema.LookIndex(name) } if idx == nil { return fmt.Errorf("failed to create index with name %s", name) } opts := m.BuildIndexOptions(idx.Fields, stmt) values := []interface{}{clause.Column{Name: idx.Name}, m.CurrentTable(stmt), opts} createIndexSQL := "CREATE " if idx.Class != "" { createIndexSQL += idx.Class + " " } createIndexSQL += "INDEX ? ON ??" if idx.Where != "" { createIndexSQL += " WHERE " + idx.Where } if idx.Option != "" { createIndexSQL += " " + idx.Option } return m.DB.Exec(createIndexSQL, values...).Error }) } func (m Migrator) HasIndex(value interface{}, name string) bool { var count int _ = m.RunWithValue(value, func(stmt *gorm.Statement) error { if stmt.Schema != nil { if idx := stmt.Schema.LookIndex(name); idx != nil { name = idx.Name } } return m.DB.Raw( "SELECT count(*) FROM sys.indexes WHERE name=? AND object_id=OBJECT_ID(?)", name, getFullQualifiedTableName(stmt), ).Row().Scan(&count) }) return count > 0 } func (m Migrator) RenameIndex(value interface{}, oldName, newName string) error { return m.RunWithValue(value, func(stmt *gorm.Statement) error { return m.DB.Exec( "sp_rename @objname = ?, @newname = ?, @objtype = 'INDEX';", fmt.Sprintf("%s.%s", stmt.Table, oldName), clause.Column{Name: newName}, ).Error }) } type Index struct { ColumnName string `gorm:"column:column_name"` IndexName string `gorm:"column:index_name"` IsUnique sql.NullBool `gorm:"column:is_unique"` IsPrimaryKey sql.NullBool `gorm:"column:is_primary_key"` } func (m Migrator) GetIndexes(value interface{}) ([]gorm.Index, error) { indexes := make([]gorm.Index, 0) err := m.RunWithValue(value, func(stmt *gorm.Statement) error { result := make([]*Index, 0) if err := m.DB.Raw(indexSQL, stmt.Table).Scan(&result).Error; err != nil { return err } indexMap := make(map[string]*migrator.Index) for _, r := range result { idx, ok := indexMap[r.IndexName] if !ok { idx = &migrator.Index{ TableName: stmt.Table, NameValue: r.IndexName, ColumnList: nil, PrimaryKeyValue: r.IsPrimaryKey, UniqueValue: r.IsUnique, } } idx.ColumnList = append(idx.ColumnList, r.ColumnName) indexMap[r.IndexName] = idx } for _, idx := range indexMap { indexes = append(indexes, idx) } return nil }) return indexes, err } func (m Migrator) HasConstraint(value interface{}, name string) bool { var count int64 _ = m.RunWithValue(value, func(stmt *gorm.Statement) error { constraint, table := m.GuessConstraintInterfaceAndTable(stmt, name) if constraint != nil { name = constraint.GetName() } tableCatalog, tableSchema, tableName := splitFullQualifiedName(table) if tableCatalog == "" { tableCatalog = m.CurrentDatabase() } if tableSchema == "" { tableSchema = "%" } return m.DB.Raw( `SELECT count(*) FROM ( SELECT C.name, T.name as table_name FROM sys.check_constraints as C INNER JOIN sys.tables as T on C.parent_object_id=T.object_id INNER JOIN INFORMATION_SCHEMA.TABLES as I on I.TABLE_NAME = T.name WHERE C.name = ? AND I.TABLE_NAME = ? AND I.TABLE_SCHEMA like ? AND I.TABLE_CATALOG = ? UNION SELECT FK.name, T.name as table_name FROM sys.foreign_keys as FK INNER JOIN sys.tables as T on FK.parent_object_id=T.object_id INNER JOIN INFORMATION_SCHEMA.TABLES as I on I.TABLE_NAME = T.name WHERE FK.name = ? AND I.TABLE_NAME = ? AND I.TABLE_SCHEMA like ? AND I.TABLE_CATALOG = ? ) as constraints;`, name, tableName, tableSchema, tableCatalog, name, tableName, tableSchema, tableCatalog, ).Row().Scan(&count) }) return count > 0 } func (m Migrator) CurrentDatabase() (name string) { _ = m.DB.Raw("SELECT DB_NAME() AS [Current Database]").Row().Scan(&name) return } func (m Migrator) DefaultSchema() (name string) { _ = m.DB.Raw("SELECT SCHEMA_NAME() AS [Default Schema]").Row().Scan(&name) return }