make dropTableColumns drop columns on sqlite and constraints on all (#6849)

1 files changed, 90 insertions, 2 deletions

diff --git a/models/migrations/migrations.go b/models/migrations/migrations.go
index 7a5b0c4184..ef43b0453b 100644
--- a/models/migrations/migrations.go
+++ b/models/migrations/migrations.go
@@ -13,6 +13,7 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"time"
 
@@ -284,11 +285,98 @@ func dropTableColumns(sess *xorm.Session, tableName string, columnNames ...strin
 	if tableName == "" || len(columnNames) == 0 {
 		return nil
 	}
+	// TODO: This will not work if there are foreign keys
 
 	switch {
 	case setting.UseSQLite3:
-		log.Warn("Unable to drop columns in SQLite")
-	case setting.UseMySQL, setting.UseTiDB, setting.UsePostgreSQL:
+		// First drop the indexes on the columns
+		res, errIndex := sess.Query(fmt.Sprintf("PRAGMA index_list(`%s`)", tableName))
+		if errIndex != nil {
+			return errIndex
+		}
+		for _, row := range res {
+			indexName := row["name"]
+			indexRes, err := sess.Query(fmt.Sprintf("PRAGMA index_info(`%s`)", indexName))
+			if err != nil {
+				return err
+			}
+			if len(indexRes) != 1 {
+				continue
+			}
+			indexColumn := string(indexRes[0]["name"])
+			for _, name := range columnNames {
+				if name == indexColumn {
+					_, err := sess.Exec(fmt.Sprintf("DROP INDEX `%s`", indexName))
+					if err != nil {
+						return err
+					}
+				}
+			}
+		}
+
+		// Here we need to get the columns from the original table
+		sql := fmt.Sprintf("SELECT sql FROM sqlite_master WHERE tbl_name='%s' and type='table'", tableName)
+		res, err := sess.Query(sql)
+		if err != nil {
+			return err
+		}
+		tableSQL := string(res[0]["sql"])
+		tableSQL = tableSQL[strings.Index(tableSQL, "("):]
+		for _, name := range columnNames {
+			tableSQL = regexp.MustCompile(regexp.QuoteMeta("`"+name+"`")+"[^`,)]*[,)]").ReplaceAllString(tableSQL, "")
+		}
+
+		columns := regexp.MustCompile("`([^`]*)`").FindAllString(tableSQL, -1)
+
+		tableSQL = fmt.Sprintf("CREATE TABLE `new_%s_new` ", tableName) + tableSQL
+		if _, err := sess.Exec(tableSQL); err != nil {
+			return err
+		}
+
+		// Now restore the data
+		columnsSeparated := strings.Join(columns, ",")
+		insertSQL := fmt.Sprintf("INSERT INTO `new_%s_new` (%s) SELECT %s FROM %s", tableName, columnsSeparated, columnsSeparated, tableName)
+		if _, err := sess.Exec(insertSQL); err != nil {
+			return err
+		}
+
+		// Now drop the old table
+		if _, err := sess.Exec(fmt.Sprintf("DROP TABLE `%s`", tableName)); err != nil {
+			return err
+		}
+
+		// Rename the table
+		if _, err := sess.Exec(fmt.Sprintf("ALTER TABLE `new_%s_new` RENAME TO `%s`", tableName, tableName)); err != nil {
+			return err
+		}
+
+	case setting.UsePostgreSQL:
+		cols := ""
+		for _, col := range columnNames {
+			if cols != "" {
+				cols += ", "
+			}
+			cols += "DROP COLUMN `" + col + "` CASCADE"
+		}
+		if _, err := sess.Exec(fmt.Sprintf("ALTER TABLE `%s` %s", tableName, cols)); err != nil {
+			return fmt.Errorf("Drop table `%s` columns %v: %v", tableName, columnNames, err)
+		}
+	case setting.UseMySQL, setting.UseTiDB:
+		// Drop indexes on columns first
+		sql := fmt.Sprintf("SHOW INDEX FROM %s WHERE column_name IN ('%s')", tableName, strings.Join(columnNames, "','"))
+		res, err := sess.Query(sql)
+		if err != nil {
+			return err
+		}
+		for _, index := range res {
+			indexName := index["column_name"]
+			_, err := sess.Exec(fmt.Sprintf("DROP INDEX `%s` ON `%s`", indexName, tableName))
+			if err != nil {
+				return err
+			}
+		}
+
+		// Now drop the columns
 		cols := ""
 		for _, col := range columnNames {
 			if cols != "" {