Repository: oom-ai/oomstore
Pull request #1106 (2022-01-27 03:49:11)

Jinghan/create batch feature's dummy snapshot table
This PR does:
- add `snapshot` to batch feature snapshot table name
- differentiate batch and streaming feature when creating dummy snapshot table


Patch:
diff --git a/README.md b/README.md
index d1de9fcb..4d33fa09 100644
--- a/README.md
+++ b/README.md
@@ -126,7 +126,7 @@ oomcli sync --revision-id 2
 
 ```bash
 oomcli get online \
-  --entity-key 1006 \
+  --entity-keys 1006 \
   --feature account.state,account.credit_score,account.account_age_days,account.has_2fa_installed,transaction_stats.transaction_count_7d,transaction_stats.transaction_count_30d
 ```
 
diff --git a/internal/database/dbutil/util.go b/internal/database/dbutil/util.go
index ea97c2ee..a03ca5a1 100644
--- a/internal/database/dbutil/util.go
+++ b/internal/database/dbutil/util.go
@@ -39,8 +39,8 @@ func Fill(size int, elem string, sep string) string {
 	return strings.Join(r, sep)
 }
 
-func OfflineBatchTableName(groupID int, revisionID int64) string {
-	return fmt.Sprintf("offline_batch_%d_%d", groupID, revisionID)
+func OfflineBatchSnapshotTableName(groupID int, revisionID int64) string {
+	return fmt.Sprintf("offline_batch_snapshot_%d_%d", groupID, revisionID)
 }
 
 func OfflineStreamSnapshotTableName(groupID int, revision int64) string {
diff --git a/internal/database/metadata/mysql/revision.go b/internal/database/metadata/mysql/revision.go
index 9aef6f2c..898b024e 100644
--- a/internal/database/metadata/mysql/revision.go
+++ b/internal/database/metadata/mysql/revision.go
@@ -6,7 +6,6 @@ import (
 	"github.com/go-sql-driver/mysql"
 	"github.com/oom-ai/oomstore/pkg/errdefs"
 
-	"github.com/oom-ai/oomstore/internal/database/dbutil"
 	"github.com/oom-ai/oomstore/internal/database/metadata"
 )
 
@@ -34,20 +33,5 @@ func createRevision(ctx context.Context, sqlxCtx metadata.SqlxContext, opt metad
 		return 0, "", errdefs.WithStack(err)
 	}
 
-	if opt.SnapshotTable == nil {
-		updateQuery := "UPDATE feature_group_revision SET snapshot_table = ? WHERE id = ?"
-		snapshotTable = dbutil.OfflineBatchTableName(opt.GroupID, revisionID)
-		result, err := sqlxCtx.ExecContext(ctx, sqlxCtx.Rebind(updateQuery), snapshotTable, revisionID)
-		if err != nil {
-			return 0, "", errdefs.WithStack(err)
-		}
-		rowsAffected, err := result.RowsAffected()
-		if err != nil {
-			return 0, "", errdefs.WithStack(err)
-		}
-		if rowsAffected != 1 {
-			return 0, "", errdefs.Errorf("failed to update revision %d: revision not found", revisionID)
-		}
-	}
 	return int(revisionID), snapshotTable, nil
 }
diff --git a/internal/database/metadata/postgres/revision.go b/internal/database/metadata/postgres/revision.go
index bf511f17..d750aa61 100644
--- a/internal/database/metadata/postgres/revision.go
+++ b/internal/database/metadata/postgres/revision.go
@@ -7,7 +7,6 @@ import (
 	"github.com/lib/pq"
 	"github.com/oom-ai/oomstore/pkg/errdefs"
 
-	"github.com/oom-ai/oomstore/internal/database/dbutil"
 	"github.com/oom-ai/oomstore/internal/database/metadata"
 )
 
@@ -30,21 +29,6 @@ func createRevision(ctx context.Context, sqlxCtx metadata.SqlxContext, opt metad
 		}
 		return 0, "", errdefs.WithStack(err)
 	}
-	if opt.SnapshotTable == nil {
-		updateQuery := "UPDATE feature_group_revision SET snapshot_table = $1 WHERE id = $2"
-		snapshotTable = dbutil.OfflineBatchTableName(opt.GroupID, int64(revisionID))
-		result, err := sqlxCtx.ExecContext(ctx, updateQuery, snapshotTable, revisionID)
-		if err != nil {
-			return 0, "", errdefs.WithStack(err)
-		}
-		rowsAffected, err := result.RowsAffected()
-		if err != nil {
-			return 0, "", errdefs.WithStack(err)
-		}
-		if rowsAffected != 1 {
-			return 0, "", errdefs.Errorf("failed to update revision %d: revision not found", revisionID)
-		}
-	}
 
 	return revisionID, snapshotTable, nil
 }
diff --git a/internal/database/metadata/sqlite/revision.go b/internal/database/metadata/sqlite/revision.go
index 0b540648..7f506ee4 100644
--- a/internal/database/metadata/sqlite/revision.go
+++ b/internal/database/metadata/sqlite/revision.go
@@ -7,7 +7,6 @@ import (
 	"modernc.org/sqlite"
 	sqlite3 "modernc.org/sqlite/lib"
 
-	"github.com/oom-ai/oomstore/internal/database/dbutil"
 	"github.com/oom-ai/oomstore/internal/database/metadata"
 )
 
@@ -35,20 +34,5 @@ func createRevision(ctx context.Context, sqlxCtx metadata.SqlxContext, opt metad
 		return 0, "", errdefs.WithStack(err)
 	}
 
-	if opt.SnapshotTable == nil {
-		updateQuery := "UPDATE feature_group_revision SET snapshot_table = ? WHERE id = ?"
-		snapshotTable = dbutil.OfflineBatchTableName(opt.GroupID, revisionID)
-		result, err := sqlxCtx.ExecContext(ctx, sqlxCtx.Rebind(updateQuery), snapshotTable, revisionID)
-		if err != nil {
-			return 0, "", errdefs.WithStack(err)
-		}
-		rowsAffected, err := result.RowsAffected()
-		if err != nil {
-			return 0, "", errdefs.WithStack(err)
-		}
-		if rowsAffected != 1 {
-			return 0, "", errdefs.Errorf("failed to update revision %d: revision not found", revisionID)
-		}
-	}
 	return int(revisionID), snapshotTable, nil
 }
diff --git a/oomcli/cmd/get_online.go b/oomcli/cmd/get_online.go
index 28ad7ddb..c2b991b3 100644
--- a/oomcli/cmd/get_online.go
+++ b/oomcli/cmd/get_online.go
@@ -41,8 +41,8 @@ func init() {
 
 	flags := getOnlineCmd.Flags()
 
-	flags.StringSliceVarP(&multiGetOpt.EntityKeys, "entity-key", "k", nil, "entity keys")
-	_ = getOnlineCmd.MarkFlagRequired("entity-key")
+	flags.StringSliceVarP(&multiGetOpt.EntityKeys, "entity-keys", "k", nil, "entity keys")
+	_ = getOnlineCmd.MarkFlagRequired("entity-keys")
 
 	flags.StringSliceVar(&multiGetOpt.FeatureNames, "feature", nil, "feature full names")
 	_ = getOnlineCmd.MarkFlagRequired("feature")
diff --git a/pkg/oomstore/import.go b/pkg/oomstore/import.go
index 38a82312..cc939dbd 100644
--- a/pkg/oomstore/import.go
+++ b/pkg/oomstore/import.go
@@ -57,7 +57,7 @@ func (s *OomStore) csvReaderImport(ctx context.Context, opt *importOpt, dataSour
 		return 0, err
 	}
 
-	newRevisionID, snapshotTableName, err := s.createRevision(ctx, metadata.CreateRevisionOpt{
+	newRevisionID, _, err := s.createRevision(ctx, metadata.CreateRevisionOpt{
 		Revision:      time.Now().UnixMilli(),
 		GroupID:       opt.group.ID,
 		SnapshotTable: nil,
@@ -67,7 +67,7 @@ func (s *OomStore) csvReaderImport(ctx context.Context, opt *importOpt, dataSour
 	if err != nil {
 		return 0, err
 	}
-
+	snapshotTableName := dbutil.OfflineBatchSnapshotTableName(opt.group.ID, int64(newRevisionID))
 	revision, err := s.offline.Import(ctx, offline.ImportOpt{
 		EntityName:        opt.entityName,
 		Features:          opt.features,
@@ -83,9 +83,10 @@ func (s *OomStore) csvReaderImport(ctx context.Context, opt *importOpt, dataSour
 	if opt.Revision != nil {
 		revision = *opt.Revision
 	}
-	if err := s.metadata.UpdateRevision(ctx, metadata.UpdateRevisionOpt{
-		RevisionID:  newRevisionID,
-		NewRevision: &revision,
+	if err = s.metadata.UpdateRevision(ctx, metadata.UpdateRevisionOpt{
+		RevisionID:       newRevisionID,
+		NewRevision:      &revision,
+		NewSnapshotTable: &snapshotTableName,
 	}); err != nil {
 		return 0, err
 	}
diff --git a/pkg/oomstore/revision.go b/pkg/oomstore/revision.go
index 2ee4656a..d4e3dee0 100644
--- a/pkg/oomstore/revision.go
+++ b/pkg/oomstore/revision.go
@@ -30,14 +30,14 @@ func (s *OomStore) createRevision(ctx context.Context, opt metadata.CreateRevisi
 	var snapshotTable string
 	var dummyRevision *types.Revision
 
-	if err := s.metadata.WithTransaction(ctx, func(c context.Context, db metadata.DBStore) error {
-		_, err := s.metadata.GetRevisionBy(ctx, opt.GroupID, 0)
+	if err := s.metadata.WithTransaction(ctx, func(c context.Context, tx metadata.DBStore) error {
+		_, err := tx.GetRevisionBy(ctx, opt.GroupID, 0)
 		if err != nil {
 			if !errdefs.IsNotFound(err) {
 				return err
 			}
 
-			if _, _, err = s.metadata.CreateRevision(ctx, metadata.CreateRevisionOpt{
+			if _, _, err = tx.CreateRevision(ctx, metadata.CreateRevisionOpt{
 				Revision:    0,
 				GroupID:     opt.GroupID,
 				Description: "dummy revision will be used at Join and Export",
@@ -45,12 +45,13 @@ func (s *OomStore) createRevision(ctx context.Context, opt metadata.CreateRevisi
 				return err
 			}
 
-			if dummyRevision, err = s.metadata.GetRevisionBy(ctx, opt.GroupID, 0); err != nil {
+			dummyRevision, err = tx.GetRevisionBy(ctx, opt.GroupID, 0)
+			if err != nil {
 				return err
 			}
 		}
 
-		revisionID, snapshotTable, err = s.metadata.CreateRevision(ctx, opt)
+		revisionID, snapshotTable, err = tx.CreateRevision(ctx, opt)
 		return err
 	}); err != nil {
 		return 0, "", err
@@ -66,7 +67,12 @@ func (s *OomStore) createRevision(ctx context.Context, opt metadata.CreateRevisi
 }
 
 func (s *OomStore) createSnapshotAndCdcTable(ctx context.Context, revision *types.Revision) error {
-	snapshotTable := dbutil.OfflineStreamSnapshotTableName(revision.GroupID, revision.Revision)
+	var snapshotTableName string
+	if revision.Group.Category == types.CategoryStream {
+		snapshotTableName = dbutil.OfflineStreamSnapshotTableName(revision.GroupID, revision.Revision)
+	} else {
+		snapshotTableName = dbutil.OfflineBatchSnapshotTableName(revision.GroupID, int64(revision.ID))
+	}
 
 	// Create snapshot table in offline store
 	features, err := s.metadata.ListFeature(ctx, metadata.ListFeatureOpt{
@@ -77,7 +83,7 @@ func (s *OomStore) createSnapshotAndCdcTable(ctx context.Context, revision *type
 	}
 
 	if err = s.offline.CreateTable(ctx, offline.CreateTableOpt{
-		TableName:  snapshotTable,
+		TableName:  snapshotTableName,
 		EntityName: revision.Group.Entity.Name,
 		Features:   features,
 		TableType:  types.TableStreamSnapshot,
@@ -102,7 +108,7 @@ func (s *OomStore) createSnapshotAndCdcTable(ctx context.Context, revision *type
 	// Update snapshot_table in feature_group_revision table
 	return s.metadata.UpdateRevision(ctx, metadata.UpdateRevisionOpt{
 		RevisionID:       revision.ID,
-		NewSnapshotTable: &snapshotTable,
+		NewSnapshotTable: &snapshotTableName,
 		NewCdcTable:      cdcTable,
 	})
 }


Test patch:
diff --git a/internal/database/metadata/test_impl/revision.go b/internal/database/metadata/test_impl/revision.go
index 1eeb6bf7..1bce4f80 100644
--- a/internal/database/metadata/test_impl/revision.go
+++ b/internal/database/metadata/test_impl/revision.go
@@ -62,7 +62,7 @@ func TestCreateRevision(t *testing.T, prepareStore PrepareStoreFn, destroyStore
 			expectedRevision: &types.Revision{
 				ID:            2,
 				Revision:      2000,
-				SnapshotTable: "offline_batch_1_2",
+				SnapshotTable: "",
 				Anchored:      false,
 				Description:   "description",
 				GroupID:       groupID,
diff --git a/oomcli/test/test_get_meta_revision.sh b/oomcli/test/test_get_meta_revision.sh
index 5944aa73..797d4097 100755
--- a/oomcli/test/test_get_meta_revision.sh
+++ b/oomcli/test/test_get_meta_revision.sh
@@ -9,8 +9,8 @@ import_device_sample > /dev/null
 case='oomcli get meta revision works'
 expected='
 ID,REVISION,GROUP,SNAPSHOT-TABLE,CDC-TABLE,DESCRIPTION,ANCHORED,CREATE-TIME,MODIFY-TIME
-1,0,phone,offline_stream_snapshot_1_0,,dummy revision will be used at Join and Export
-2,1638519905556,phone,offline_batch_1_2,,test data,true,2021-12-03T08:25:05Z,2021-12-03T08:25:05Z
+1,0,phone,offline_batch_snapshot_1_1,,dummy revision will be used at Join and Export
+2,1638519905556,phone,offline_batch_snapshot_1_2,,test data,true,2021-12-03T08:25:05Z,2021-12-03T08:25:05Z
 '
 actual=$(oomcli get meta revision --group phone -o csv --wide)
 filter() { cut -d ',' -f 1,3-6 <<<"$1"; }
diff --git a/oomcli/test/test_import_table_link.sh b/oomcli/test/test_import_table_link.sh
index d94ae161..101e4c06 100755
--- a/oomcli/test/test_import_table_link.sh
+++ b/oomcli/test/test_import_table_link.sh
@@ -9,13 +9,13 @@ import_device_sample
 case='oomcli import using table link'
 oomcli import \
     --group phone \
-    --table-link "offline_batch_1_2" \
+    --table-link "offline_batch_snapshot_1_2" \
     --description 'linked table' > /dev/null
 expected='
 ID,REVISION,GROUP,SNAPSHOT-TABLE,CDC-TABLE,DESCRIPTION
-1,0,phone,offline_stream_snapshot_1_0,
-2,1639047117470,phone,offline_batch_1_2,,test data
-3,1639047117552,phone,offline_batch_1_2,,linked table
+1,0,phone,offline_batch_snapshot_1_1,
+2,1639047117470,phone,offline_batch_snapshot_1_2,,test data
+3,1639047117552,phone,offline_batch_snapshot_1_2,,linked table
 '
 filter() { cut -d, -f 1,3,4,5; }
 actual=$(oomcli get meta revision -o csv)