DryRun for migrator

go-gorm · Sep 18, 2022 · c035852 · c035852
1 parent 4906259
commit c035852
Showing 1 changed file with 30 additions and 12 deletions.
diff --git a/migrator/migrator.go b/migrator/migrator.go
@@ -8,9 +8,11 @@ import (
 	"reflect"
 	"regexp"
 	"strings"
+	"time"
 
 	"gorm.io/gorm"
 	"gorm.io/gorm/clause"
+	"gorm.io/gorm/logger"
 	"gorm.io/gorm/schema"
 )
 
@@ -30,6 +32,17 @@ type Config struct {
 	gorm.Dialector
 }
 
+type printSQLLogger struct{}
+
+func (l *printSQLLogger) LogMode(ll logger.LogLevel) logger.Interface                { return l }
+func (l *printSQLLogger) Info(ctx context.Context, log string, args ...interface{})  {}
+func (l *printSQLLogger) Warn(ctx context.Context, log string, args ...interface{})  {}
+func (l *printSQLLogger) Error(ctx context.Context, log string, args ...interface{}) {}
+func (l *printSQLLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
+	sql, _ := fc()
+	fmt.Println(sql + ";")
+}
+
 // GormDataTypeInterface gorm data type interface
 type GormDataTypeInterface interface {
 	GormDBDataType(*gorm.DB, *schema.Field) string
@@ -92,14 +105,19 @@ func (m Migrator) FullDataTypeOf(field *schema.Field) (expr clause.Expr) {
 // AutoMigrate auto migrate values
 func (m Migrator) AutoMigrate(values ...interface{}) error {
 	for _, value := range m.ReorderModels(values, true) {
-		tx := m.DB.Session(&gorm.Session{})
-		if !tx.Migrator().HasTable(value) {
-			if err := tx.Migrator().CreateTable(value); err != nil {
+		queryTx := m.DB.Session(&gorm.Session{})
+		execTx := queryTx
+		if m.DB.DryRun {
+			queryTx.DryRun = false
+			execTx = m.DB.Session(&gorm.Session{Logger: &printSQLLogger{}})
+		}
+		if !queryTx.Migrator().HasTable(value) {
+			if err := execTx.Migrator().CreateTable(value); err != nil {
 				return err
 			}
 		} else {
 			if err := m.RunWithValue(value, func(stmt *gorm.Statement) (errr error) {
-				columnTypes, err := m.DB.Migrator().ColumnTypes(value)
+				columnTypes, err := queryTx.Migrator().ColumnTypes(value)
 				if err != nil {
 					return err
 				}
@@ -117,10 +135,10 @@ func (m Migrator) AutoMigrate(values ...interface{}) error {
 
 					if foundColumn == nil {
 						// not found, add column
-						if err := tx.Migrator().AddColumn(value, dbName); err != nil {
+						if err := execTx.Migrator().AddColumn(value, dbName); err != nil {
 							return err
 						}
-					} else if err := m.DB.Migrator().MigrateColumn(value, field, foundColumn); err != nil {
+					} else if err := execTx.Migrator().MigrateColumn(value, field, foundColumn); err != nil {
 						// found, smart migrate
 						return err
 					}
@@ -129,25 +147,25 @@ func (m Migrator) AutoMigrate(values ...interface{}) error {
 				for _, rel := range stmt.Schema.Relationships.Relations {
 					if !m.DB.Config.DisableForeignKeyConstraintWhenMigrating {
 						if constraint := rel.ParseConstraint(); constraint != nil &&
-							constraint.Schema == stmt.Schema && !tx.Migrator().HasConstraint(value, constraint.Name) {
-							if err := tx.Migrator().CreateConstraint(value, constraint.Name); err != nil {
+							constraint.Schema == stmt.Schema && !queryTx.Migrator().HasConstraint(value, constraint.Name) {
+							if err := execTx.Migrator().CreateConstraint(value, constraint.Name); err != nil {
 								return err
 							}
 						}
 					}
 				}
 
 				for _, chk := range stmt.Schema.ParseCheckConstraints() {
-					if !tx.Migrator().HasConstraint(value, chk.Name) {
-						if err := tx.Migrator().CreateConstraint(value, chk.Name); err != nil {
+					if !queryTx.Migrator().HasConstraint(value, chk.Name) {
+						if err := execTx.Migrator().CreateConstraint(value, chk.Name); err != nil {
 							return err
 						}
 					}
 				}
 
 				for _, idx := range stmt.Schema.ParseIndexes() {
-					if !tx.Migrator().HasIndex(value, idx.Name) {
-						if err := tx.Migrator().CreateIndex(value, idx.Name); err != nil {
+					if !queryTx.Migrator().HasIndex(value, idx.Name) {
+						if err := execTx.Migrator().CreateIndex(value, idx.Name); err != nil {
 							return err
 						}
 					}