7. 数据库迁移 Migrate

day7-migrate

1. SQLite

1.1 新增字段

ALTER TABLE table_name ADD COLUMN col_name, col_type;

1.2 删除字段

CREATE TABLE new_table AS SELECT col1, col2, ... from old_table
DROP TABLE old_table
ALTER TABLE new_table RENAME TO old_table;

SQLite 删除字段需要选择需要的字段构成新表。

2. 实现

geeorm/geeorm.go

// return a - b
func difference(a []string, b []string) []string {
	setB := make(map[string]struct{})
	for _, v := range b {
		setB[v] = struct{}{}
	}

	diff := make([]string, 0)
	for _, v := range a {
		if _, ok := setB[v]; !ok {
			diff = append(diff, v)
		}
	}
	return diff
}

// Migrate table
func (engine *Engine) Migrate(value any) error {
	_, err := engine.Transaction(func(s *session.Session) (any, error) {
		if !s.Model(value).HasTable() {
			log.Infof("table %s doesn't exist, creat table", s.RefTable().Name)
			return nil, s.CreateTable()
		}
		table := s.RefTable()
		rows, _ := s.Raw(fmt.Sprintf("SELECT * FROM %s LIMIT 1", table.Name)).QueryRows()
		columns, _ := rows.Columns()
		addCols := difference(table.FieldNames, columns)
		delCols := difference(columns, table.FieldNames)
		log.Info("add cols %v, deleted cols %v", addCols, delCols)

		for _, col := range addCols {
			f := table.GetField(col)
			sqlStr := fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s;", table.Name, f.Name, f.Type)
			if _, err := s.Raw(sqlStr).Exec(); err != nil {
				return nil, err
			}
		}

		if len(delCols) == 0 {
			return nil, nil
		}
		tmp := "tmp_" + table.Name
		fieldStr := strings.Join(table.FieldNames, ", ")
		s.Raw(fmt.Sprintf("CREATE TABLE %s AS SELECT %s FROM %s;", tmp, fieldStr, table.Name))
		s.Raw(fmt.Sprintf("DROP TABLE %s;", table.Name))
		s.Raw(fmt.Sprintf("ALTER TABLE %s RENAME TO %s;", tmp, table.Name))
		_, err := s.Exec()
		return nil, err
	})

	return err
}

• difference：获取两个数据表的字段差集

3. 单元测试

func TestEngine_Migrate(t *testing.T) {
	engine := openDB(t)
	defer engine.Close()
	s := engine.NewSession()

	// Migrate User{Name string, XXX int} to User{Name string, Age int}
	_, _ = s.Raw("DROP TABLE IF EXISTS User;").Exec()
	_, _ = s.Raw("CREATE TABLE User(Name text PRIMARY KEY, XXX integer);").Exec()
	_, _ = s.Raw("INSERT INTO User(`Name`) values (?), (?)", "Tom", "Sam").Exec()
	engine.Migrate(&User{})

	rows, _ := s.Raw("SELECT * FROM User").QueryRows()
	cols, _ := rows.Columns()
	if !reflect.DeepEqual(cols, []string{"Name", "Age"}) {
		t.Fatal("failed to migrate table User, got cols:", cols)
	}

}

Reference

1. https://geektutu.com/post/geeorm-day7.html