restore: reset table sequences for PostgreSQL (#4357)

1 files changed, 16 insertions, 1 deletions

diff --git a/models/models.go b/models/models.go
index 98a4fd23..088d8f58 100644
--- a/models/models.go
+++ b/models/models.go
@@ -302,7 +302,9 @@ func DumpDatabase(dirPath string) (err error) {
 }
 
 // ImportDatabase imports data from backup archive.
-func ImportDatabase(dirPath string) (err error) {
+func ImportDatabase(dirPath string, verbose bool) (err error) {
+	snakeMapper := core.SnakeMapper{}
+
 	// Purposely create a local variable to not modify global variable
 	tables := append(tables, new(Version))
 	for _, table := range tables {
@@ -312,6 +314,10 @@ func ImportDatabase(dirPath string) (err error) {
 			continue
 		}
 
+		if verbose {
+			log.Trace("Importing table '%s'...", tableName)
+		}
+
 		if err = x.DropTables(table); err != nil {
 			return fmt.Errorf("fail to drop table '%s': %v", tableName, err)
 		} else if err = x.Sync2(table); err != nil {
@@ -353,6 +359,15 @@ func ImportDatabase(dirPath string) (err error) {
 				return fmt.Errorf("fail to insert strcut: %v", err)
 			}
 		}
+
+		// PostgreSQL needs manually reset table sequence for auto increment keys
+		if setting.UsePostgreSQL {
+			rawTableName := snakeMapper.Obj2Table(tableName)
+			seqName := rawTableName + "_id_seq"
+			if _, err = x.Exec(fmt.Sprintf(`SELECT setval('%s', COALESCE((SELECT MAX(id)+1 FROM "%s"), 1), false);`, seqName, rawTableName)); err != nil {
+				return fmt.Errorf("fail to reset table '%s' sequence: %v", rawTableName, err)
+			}
+		}
 	}
 	return nil
 }