Fix column removal in MSSQL (#3638)

* Fix column removal in MSSQL

* Use xorm session in MSSQL drop column operations

* Add transaction as MSSQL alter table is transactional

2 files changed, 61 insertions, 21 deletions

diff --git a/models/migrations/migrations.go b/models/migrations/migrations.go
index dfaef2c78d..e7f3dae646 100644
--- a/models/migrations/migrations.go
+++ b/models/migrations/migrations.go
@@ -217,6 +217,66 @@ Please try to upgrade to a lower version (>= v0.6.0) first, then upgrade to curr
 	return nil
 }
 
+func dropTableColumns(x *xorm.Engine, tableName string, columnNames ...string) (err error) {
+	if tableName == "" || len(columnNames) == 0 {
+		return nil
+	}
+
+	switch {
+	case setting.UseSQLite3:
+		log.Warn("Unable to drop columns in SQLite")
+	case setting.UseMySQL, setting.UseTiDB, setting.UsePostgreSQL:
+		cols := ""
+		for _, col := range columnNames {
+			if cols != "" {
+				cols += ", "
+			}
+			cols += "DROP COLUMN `" + col + "`"
+		}
+		if _, err := x.Exec(fmt.Sprintf("ALTER TABLE `%s` %s", tableName, columnNames)); err != nil {
+			return fmt.Errorf("Drop table `%s` columns %v: %v", tableName, columnNames, err)
+		}
+	case setting.UseMSSQL:
+		sess := x.NewSession()
+		defer sess.Close()
+
+		if err = sess.Begin(); err != nil {
+			return err
+		}
+
+		cols := ""
+		for _, col := range columnNames {
+			if cols != "" {
+				cols += ", "
+			}
+			cols += "`" + strings.ToLower(col) + "`"
+		}
+		sql := fmt.Sprintf("SELECT Name FROM SYS.DEFAULT_CONSTRAINTS WHERE PARENT_OBJECT_ID = OBJECT_ID('%[1]s') AND PARENT_COLUMN_ID IN (SELECT column_id FROM sys.columns WHERE lower(NAME) IN (%[2]s) AND object_id = OBJECT_ID('%[1]s'))",
+			tableName, strings.Replace(cols, "`", "'", -1))
+		constraints := make([]string, 0)
+		if err := sess.SQL(sql).Find(&constraints); err != nil {
+			sess.Rollback()
+			return fmt.Errorf("Find constraints: %v", err)
+		}
+		for _, constraint := range constraints {
+			if _, err := sess.Exec(fmt.Sprintf("ALTER TABLE `%s` DROP CONSTRAINT `%s`", tableName, constraint)); err != nil {
+				sess.Rollback()
+				return fmt.Errorf("Drop table `%s` constraint `%s`: %v", tableName, constraint, err)
+			}
+		}
+		if _, err := sess.Exec(fmt.Sprintf("ALTER TABLE `%s` DROP COLUMN %s", tableName, cols)); err != nil {
+			sess.Rollback()
+			return fmt.Errorf("Drop table `%s` columns %v: %v", tableName, columnNames, err)
+		}
+
+		return sess.Commit()
+	default:
+		log.Fatal(4, "Unrecognized DB")
+	}
+
+	return nil
+}
+
 func fixLocaleFileLoadPanic(_ *xorm.Engine) error {
 	cfg, err := ini.Load(setting.CustomConf)
 	if err != nil {
diff --git a/models/migrations/v56.go b/models/migrations/v56.go
index bbaee28a92..1f96cc543e 100644
--- a/models/migrations/v56.go
+++ b/models/migrations/v56.go
@@ -5,29 +5,9 @@
 package migrations
 
 import (
-	"fmt"
-
-	"code.gitea.io/gitea/modules/log"
-	"code.gitea.io/gitea/modules/setting"
-
 	"github.com/go-xorm/xorm"
 )
 
 func removeIsOwnerColumnFromOrgUser(x *xorm.Engine) (err error) {
-	switch {
-	case setting.UseSQLite3:
-		log.Warn("Unable to drop columns in SQLite")
-	case setting.UseMySQL, setting.UseTiDB, setting.UsePostgreSQL:
-		if _, err := x.Exec("ALTER TABLE org_user DROP COLUMN is_owner, DROP COLUMN num_teams"); err != nil {
-			return fmt.Errorf("DROP COLUMN org_user.is_owner, org_user.num_teams: %v", err)
-		}
-	case setting.UseMSSQL:
-		if _, err := x.Exec("ALTER TABLE org_user DROP COLUMN is_owner, num_teams"); err != nil {
-			return fmt.Errorf("DROP COLUMN org_user.is_owner, org_user.num_teams: %v", err)
-		}
-	default:
-		log.Fatal(4, "Unrecognized DB")
-	}
-
-	return nil
+	return dropTableColumns(x, "org_user", "is_owner", "num_teams")
 }