123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363 |
- package sqlite
- import (
- "fmt"
- "regexp"
- "strings"
- "gorm.io/gorm"
- "gorm.io/gorm/clause"
- "gorm.io/gorm/migrator"
- "gorm.io/gorm/schema"
- )
- type Migrator struct {
- migrator.Migrator
- }
- func (m *Migrator) RunWithoutForeignKey(fc func() error) error {
- var enabled int
- m.DB.Raw("PRAGMA foreign_keys").Scan(&enabled)
- if enabled == 1 {
- m.DB.Exec("PRAGMA foreign_keys = OFF")
- defer m.DB.Exec("PRAGMA foreign_keys = ON")
- }
- return fc()
- }
- func (m Migrator) HasTable(value interface{}) bool {
- var count int
- m.Migrator.RunWithValue(value, func(stmt *gorm.Statement) error {
- return m.DB.Raw("SELECT count(*) FROM sqlite_master WHERE type='table' AND name=?", stmt.Table).Row().Scan(&count)
- })
- return count > 0
- }
- func (m Migrator) DropTable(values ...interface{}) error {
- return m.RunWithoutForeignKey(func() error {
- values = m.ReorderModels(values, false)
- tx := m.DB.Session(&gorm.Session{})
- for i := len(values) - 1; i >= 0; i-- {
- if err := m.RunWithValue(values[i], func(stmt *gorm.Statement) error {
- return tx.Exec("DROP TABLE IF EXISTS ?", clause.Table{Name: stmt.Table}).Error
- }); err != nil {
- return err
- }
- }
- return nil
- })
- }
- func (m Migrator) HasColumn(value interface{}, name string) bool {
- var count int
- m.Migrator.RunWithValue(value, func(stmt *gorm.Statement) error {
- if stmt.Schema != nil {
- if field := stmt.Schema.LookUpField(name); field != nil {
- name = field.DBName
- }
- }
- if name != "" {
- m.DB.Raw(
- "SELECT count(*) FROM sqlite_master WHERE type = ? AND tbl_name = ? AND (sql LIKE ? OR sql LIKE ? OR sql LIKE ?)",
- "table", stmt.Table, `%"`+name+`" %`, `%`+name+` %`, "%`"+name+"`%",
- ).Row().Scan(&count)
- }
- return nil
- })
- return count > 0
- }
- func (m Migrator) AlterColumn(value interface{}, name string) error {
- return m.RunWithoutForeignKey(func() error {
- return m.recreateTable(value, nil, func(rawDDL string, stmt *gorm.Statement) (sql string, sqlArgs []interface{}, err error) {
- if field := stmt.Schema.LookUpField(name); field != nil {
- reg, err := regexp.Compile("(`|'|\"| )" + field.DBName + "(`|'|\"| ) .*?,")
- if err != nil {
- return "", nil, err
- }
- createSQL := reg.ReplaceAllString(rawDDL, fmt.Sprintf("`%v` ?,", field.DBName))
- return createSQL, []interface{}{m.FullDataTypeOf(field)}, nil
- } else {
- return "", nil, fmt.Errorf("failed to alter field with name %v", name)
- }
- })
- })
- }
- func (m Migrator) DropColumn(value interface{}, name string) error {
- return m.recreateTable(value, nil, func(rawDDL string, stmt *gorm.Statement) (sql string, sqlArgs []interface{}, err error) {
- if field := stmt.Schema.LookUpField(name); field != nil {
- name = field.DBName
- }
- reg, err := regexp.Compile("(`|'|\"| )" + name + "(`|'|\"| ) .*?,")
- if err != nil {
- return "", nil, err
- }
- createSQL := reg.ReplaceAllString(rawDDL, "")
- return createSQL, nil, nil
- })
- }
- func (m Migrator) CreateConstraint(value interface{}, name string) error {
- return m.RunWithValue(value, func(stmt *gorm.Statement) error {
- constraint, chk, table := m.GuessConstraintAndTable(stmt, name)
- return m.recreateTable(value, &table,
- func(rawDDL string, stmt *gorm.Statement) (sql string, sqlArgs []interface{}, err error) {
- var (
- constraintName string
- constraintSql string
- constraintValues []interface{}
- )
- if constraint != nil {
- constraintName = constraint.Name
- constraintSql, constraintValues = buildConstraint(constraint)
- } else if chk != nil {
- constraintName = chk.Name
- constraintSql = "CONSTRAINT ? CHECK (?)"
- constraintValues = []interface{}{clause.Column{Name: chk.Name}, clause.Expr{SQL: chk.Constraint}}
- } else {
- return "", nil, nil
- }
- createDDL, err := parseDDL(rawDDL)
- if err != nil {
- return "", nil, err
- }
- createDDL.addConstraint(constraintName, constraintSql)
- createSQL := createDDL.compile()
- return createSQL, constraintValues, nil
- })
- })
- }
- func (m Migrator) DropConstraint(value interface{}, name string) error {
- return m.RunWithValue(value, func(stmt *gorm.Statement) error {
- constraint, chk, table := m.GuessConstraintAndTable(stmt, name)
- if constraint != nil {
- name = constraint.Name
- } else if chk != nil {
- name = chk.Name
- }
- return m.recreateTable(value, &table,
- func(rawDDL string, stmt *gorm.Statement) (sql string, sqlArgs []interface{}, err error) {
- createDDL, err := parseDDL(rawDDL)
- if err != nil {
- return "", nil, err
- }
- createDDL.removeConstraint(name)
- createSQL := createDDL.compile()
- return createSQL, nil, nil
- })
- })
- }
- func (m Migrator) HasConstraint(value interface{}, name string) bool {
- var count int64
- m.RunWithValue(value, func(stmt *gorm.Statement) error {
- constraint, chk, table := m.GuessConstraintAndTable(stmt, name)
- if constraint != nil {
- name = constraint.Name
- } else if chk != nil {
- name = chk.Name
- }
- m.DB.Raw(
- "SELECT count(*) FROM sqlite_master WHERE type = ? AND tbl_name = ? AND (sql LIKE ? OR sql LIKE ? OR sql LIKE ?)",
- "table", table, `%CONSTRAINT "`+name+`" %`, `%CONSTRAINT `+name+` %`, "%CONSTRAINT `"+name+"`%",
- ).Row().Scan(&count)
- return nil
- })
- return count > 0
- }
- func (m Migrator) CurrentDatabase() (name string) {
- var null interface{}
- m.DB.Raw("PRAGMA database_list").Row().Scan(&null, &name, &null)
- return
- }
- func (m Migrator) BuildIndexOptions(opts []schema.IndexOption, stmt *gorm.Statement) (results []interface{}) {
- for _, opt := range opts {
- str := stmt.Quote(opt.DBName)
- if opt.Expression != "" {
- str = opt.Expression
- }
- if opt.Collate != "" {
- str += " COLLATE " + opt.Collate
- }
- if opt.Sort != "" {
- str += " " + opt.Sort
- }
- results = append(results, clause.Expr{SQL: str})
- }
- return
- }
- func (m Migrator) CreateIndex(value interface{}, name string) error {
- return m.RunWithValue(value, func(stmt *gorm.Statement) error {
- if idx := stmt.Schema.LookIndex(name); idx != nil {
- opts := m.BuildIndexOptions(idx.Fields, stmt)
- values := []interface{}{clause.Column{Name: idx.Name}, clause.Table{Name: stmt.Table}, opts}
- createIndexSQL := "CREATE "
- if idx.Class != "" {
- createIndexSQL += idx.Class + " "
- }
- createIndexSQL += "INDEX ?"
- if idx.Type != "" {
- createIndexSQL += " USING " + idx.Type
- }
- createIndexSQL += " ON ??"
- if idx.Where != "" {
- createIndexSQL += " WHERE " + idx.Where
- }
- return m.DB.Exec(createIndexSQL, values...).Error
- }
- return fmt.Errorf("failed to create index with name %v", name)
- })
- }
- func (m Migrator) HasIndex(value interface{}, name string) bool {
- var count int
- m.RunWithValue(value, func(stmt *gorm.Statement) error {
- if idx := stmt.Schema.LookIndex(name); idx != nil {
- name = idx.Name
- }
- if name != "" {
- m.DB.Raw(
- "SELECT count(*) FROM sqlite_master WHERE type = ? AND tbl_name = ? AND name = ?", "index", stmt.Table, name,
- ).Row().Scan(&count)
- }
- return nil
- })
- return count > 0
- }
- func (m Migrator) RenameIndex(value interface{}, oldName, newName string) error {
- return m.RunWithValue(value, func(stmt *gorm.Statement) error {
- var sql string
- m.DB.Raw("SELECT sql FROM sqlite_master WHERE type = ? AND tbl_name = ? AND name = ?", "index", stmt.Table, oldName).Row().Scan(&sql)
- if sql != "" {
- return m.DB.Exec(strings.Replace(sql, oldName, newName, 1)).Error
- }
- return fmt.Errorf("failed to find index with name %v", oldName)
- })
- }
- func (m Migrator) DropIndex(value interface{}, name string) error {
- return m.RunWithValue(value, func(stmt *gorm.Statement) error {
- if idx := stmt.Schema.LookIndex(name); idx != nil {
- name = idx.Name
- }
- return m.DB.Exec("DROP INDEX ?", clause.Column{Name: name}).Error
- })
- }
- func buildConstraint(constraint *schema.Constraint) (sql string, results []interface{}) {
- sql = "CONSTRAINT ? FOREIGN KEY ? REFERENCES ??"
- if constraint.OnDelete != "" {
- sql += " ON DELETE " + constraint.OnDelete
- }
- if constraint.OnUpdate != "" {
- sql += " ON UPDATE " + constraint.OnUpdate
- }
- var foreignKeys, references []interface{}
- for _, field := range constraint.ForeignKeys {
- foreignKeys = append(foreignKeys, clause.Column{Name: field.DBName})
- }
- for _, field := range constraint.References {
- references = append(references, clause.Column{Name: field.DBName})
- }
- results = append(results, clause.Table{Name: constraint.Name}, foreignKeys, clause.Table{Name: constraint.ReferenceSchema.Table}, references)
- return
- }
- func (m Migrator) getRawDDL(table string) (string, error) {
- var createSQL string
- m.DB.Raw("SELECT sql FROM sqlite_master WHERE type = ? AND tbl_name = ? AND name = ?", "table", table, table).Row().Scan(&createSQL)
- if m.DB.Error != nil {
- return "", m.DB.Error
- }
- return createSQL, nil
- }
- func (m Migrator) recreateTable(value interface{}, tablePtr *string,
- getCreateSQL func(rawDDL string, stmt *gorm.Statement) (sql string, sqlArgs []interface{}, err error)) error {
- return m.RunWithValue(value, func(stmt *gorm.Statement) error {
- table := stmt.Table
- if tablePtr != nil {
- table = *tablePtr
- }
- rawDDL, err := m.getRawDDL(table)
- if err != nil {
- return err
- }
- newTableName := table + "__temp"
- createSQL, sqlArgs, err := getCreateSQL(rawDDL, stmt)
- if err != nil {
- return err
- }
- if createSQL == "" {
- return nil
- }
- tableReg, err := regexp.Compile(" ('|`|\"| )" + table + "('|`|\"| ) ")
- if err != nil {
- return err
- }
- createSQL = tableReg.ReplaceAllString(createSQL, fmt.Sprintf(" `%v` ", newTableName))
- createDDL, err := parseDDL(createSQL)
- if err != nil {
- return err
- }
- columns := createDDL.getColumns()
- return m.DB.Transaction(func(tx *gorm.DB) error {
- queries := []string{
- createSQL,
- fmt.Sprintf("INSERT INTO `%v`(%v) SELECT %v FROM `%v`", newTableName, strings.Join(columns, ","), strings.Join(columns, ","), table),
- fmt.Sprintf("DROP TABLE `%v`", table),
- fmt.Sprintf("ALTER TABLE `%v` RENAME TO `%v`", newTableName, table),
- }
- for _, query := range queries {
- if err := tx.Exec(query, sqlArgs...).Error; err != nil {
- return err
- }
- }
- return nil
- })
- })
- }
|