feat(internal/database): add Revision for db method Import

oom-ai · Nov 8, 2021 · 1a44029 · 1a44029
1 parent 3a2d242
commit 1a44029
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/internal/database/offline/postgres/import.go b/internal/database/offline/postgres/import.go
@@ -63,8 +63,13 @@ func (db *DB) Import(ctx context.Context, opt offline.ImportOpt) (int64, string,
 			return err
 		}
 
-		// generate revision using current timestamp
-		revision = time.Now().Unix()
+		if opt.Revision != nil {
+			// use user-defined revision
+			revision = *opt.Revision
+		} else {
+			// generate revision using current timestamp
+			revision = time.Now().Unix()
+		}
 
 		// generate final data table name
 		finalTableName = opt.GroupName + "_" + strconv.FormatInt(revision, 10)

diff --git a/internal/database/offline/postgres/import_test.go b/internal/database/offline/postgres/import_test.go
@@ -50,7 +50,9 @@ func TestImport(t *testing.T) {
 	})
 
 	t.Run("normal import call", func(t *testing.T) {
+		revision := int64(1234)
 		opt.Features[0].DBValueType = "varchar(32)"
+		opt.Revision = &revision
 		_, tableName, err := db.Import(context.Background(), opt)
 		assert.Nil(t, err)
 

diff --git a/internal/database/offline/types.go b/internal/database/offline/types.go
@@ -25,6 +25,7 @@ type ImportOpt struct {
 	Entity    *types.Entity
 	Features  types.FeatureList
 	Header    []string
+	Revision  *int64
 
 	// CsvReader must not contain header
 	CsvReader *csv.Reader