Repository: oom-ai/oomstore
Pull request #1066 (2022-01-24 08:05:20)

Jinghan/offline method `Import` supports streaming features
This PR refactors offline method `Import`: support importing streaming features.

related to #1032 


Patch:
diff --git a/internal/database/online/cassandra/import.go b/internal/database/online/cassandra/import.go
index d1452768..478e9381 100644
--- a/internal/database/online/cassandra/import.go
+++ b/internal/database/online/cassandra/import.go
@@ -6,33 +6,41 @@ import (
 	"strings"
 
 	"github.com/gocql/gocql"
-	"github.com/oom-ai/oomstore/pkg/errdefs"
 
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
 	"github.com/oom-ai/oomstore/internal/database/online"
 	"github.com/oom-ai/oomstore/internal/database/online/sqlutil"
+	"github.com/oom-ai/oomstore/pkg/errdefs"
+	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
 func (db *DB) Import(ctx context.Context, opt online.ImportOpt) error {
-	columns := append([]string{opt.Entity.Name}, opt.Features.Names()...)
-	tableName := sqlutil.OnlineBatchTableName(opt.Revision.ID)
-
-	table := dbutil.BuildTableSchema(tableName, opt.Entity, false, opt.Features, []string{opt.Entity.Name}, Backend)
+	// Step 0: drop existing table for streaming feature
+	var tableName string
+	if opt.Group.Category == types.CategoryBatch {
+		tableName = sqlutil.OnlineBatchTableName(opt.Revision.ID)
+	} else {
+		tableName = sqlutil.OnlineStreamTableName(opt.Group.ID)
+		if err := db.Query(fmt.Sprintf(`DROP TABLE IF EXISTS %s;`, tableName)).Exec(); err != nil {
+			return errdefs.WithStack(err)
+		}
+	}
 
-	// create table
-	if err := db.Query(table).Exec(); err != nil {
+	// Step 1: create online table
+	entity := opt.Group.Entity
+	columns := append([]string{entity.Name}, opt.Features.Names()...)
+	schema := dbutil.BuildTableSchema(tableName, entity, false, opt.Features, []string{entity.Name}, Backend)
+	if err := db.Query(schema).Exec(); err != nil {
 		return errdefs.WithStack(err)
 	}
 
-	var (
-		insertStmt = buildInsertStatement(tableName, columns)
-		batch      = db.NewBatch(gocql.LoggedBatch)
-	)
+	// Step 2: insert records to the online table
+	insertStmt := buildInsertStatement(tableName, columns)
+	batch := db.NewBatch(gocql.LoggedBatch)
 	for record := range opt.ExportStream {
 		if len(record) != len(opt.Features)+1 {
 			return errdefs.Errorf("field count not matched, expected %d, got %d", len(opt.Features)+1, len(record))
 		}
-
 		if batch.Size() != BatchSize {
 			batch.Query(insertStmt, record...)
 		} else {
diff --git a/internal/database/online/dynamodb/import.go b/internal/database/online/dynamodb/import.go
index 663bd936..f05392b2 100644
--- a/internal/database/online/dynamodb/import.go
+++ b/internal/database/online/dynamodb/import.go
@@ -20,19 +20,32 @@ const (
 )
 
 func (db *DB) Import(ctx context.Context, opt online.ImportOpt) error {
+	// Step 0: clean up existing table for streaming feature
+	var tableName string
+	if opt.Group.Category == oomTypes.CategoryBatch {
+		tableName = sqlutil.OnlineBatchTableName(opt.Revision.ID)
+	} else {
+		tableName = sqlutil.OnlineStreamTableName(opt.Group.ID)
+		_, err := db.DeleteTable(ctx, &dynamodb.DeleteTableInput{
+			TableName: aws.String(tableName),
+		})
+		if err != nil {
+			return errdefs.WithStack(err)
+		}
+	}
 	// Step 1: create table
-	tableName := sqlutil.OnlineBatchTableName(opt.Revision.ID)
+	entity := opt.Group.Entity
 	_, err := db.CreateTable(ctx, &dynamodb.CreateTableInput{
 		TableName: aws.String(tableName),
 		KeySchema: []types.KeySchemaElement{
 			{
-				AttributeName: aws.String(opt.Entity.Name),
+				AttributeName: aws.String(entity.Name),
 				KeyType:       types.KeyTypeHash,
 			},
 		},
 		AttributeDefinitions: []types.AttributeDefinition{
 			{
-				AttributeName: aws.String(opt.Entity.Name),
+				AttributeName: aws.String(entity.Name),
 				AttributeType: types.ScalarAttributeTypeS,
 			},
 		},
@@ -92,7 +105,7 @@ func buildItem(record oomTypes.ExportRecord, opt online.ImportOpt) (map[string]t
 	if err != nil {
 		return nil, errdefs.WithStack(err)
 	}
-	item[opt.Entity.Name] = entityKeyValue
+	item[opt.Group.Entity.Name] = entityKeyValue
 
 	for i, feature := range opt.Features {
 		value, err := dbutil.SerializeByValueType(record.ValueAt(i), feature.ValueType, Backend)
diff --git a/internal/database/online/redis/import.go b/internal/database/online/redis/import.go
index 177d2958..44b316b0 100644
--- a/internal/database/online/redis/import.go
+++ b/internal/database/online/redis/import.go
@@ -5,9 +5,9 @@ import (
 	"fmt"
 
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
-	"github.com/oom-ai/oomstore/pkg/errdefs"
-
 	"github.com/oom-ai/oomstore/internal/database/online"
+	"github.com/oom-ai/oomstore/pkg/errdefs"
+	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
 func (db *DB) Import(ctx context.Context, opt online.ImportOpt) error {
@@ -21,7 +21,13 @@ func (db *DB) Import(ctx context.Context, opt online.ImportOpt) error {
 		}
 
 		entityKey, values := record[0], record[1:]
-		key, err := serializeRedisKeyForBatchFeature(opt.Revision.ID, entityKey)
+		var key string
+		var err error
+		if opt.Group.Category == types.CategoryBatch {
+			key, err = serializeRedisKeyForBatchFeature(opt.Revision.ID, entityKey)
+		} else {
+			key, err = serializeRedisKeyForStreamFeature(opt.Group.ID, entityKey)
+		}
 		if err != nil {
 			return err
 		}
@@ -48,7 +54,7 @@ func (db *DB) Import(ctx context.Context, opt online.ImportOpt) error {
 
 		seq++
 		if seq%PipelineBatchSize == 0 {
-			if _, err := pipe.Exec(ctx); err != nil {
+			if _, err = pipe.Exec(ctx); err != nil {
 				return errdefs.WithStack(err)
 			}
 		}
diff --git a/internal/database/online/sqlutil/import.go b/internal/database/online/sqlutil/import.go
index d18632e9..95c56adf 100644
--- a/internal/database/online/sqlutil/import.go
+++ b/internal/database/online/sqlutil/import.go
@@ -2,6 +2,7 @@ package sqlutil
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/jmoiron/sqlx"
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
@@ -13,11 +14,17 @@ import (
 const importBatchSize = 10
 
 func Import(ctx context.Context, db *sqlx.DB, opt online.ImportOpt, backend types.BackendType) error {
-	columns := append([]string{opt.Entity.Name}, opt.Features.Names()...)
+	var tableName string
+	if opt.Group.Category == types.CategoryStream {
+		tableName = dbutil.TempTable("online_stream")
+	} else {
+		tableName = OnlineBatchTableName(opt.Revision.ID)
+	}
+	entity := opt.Group.Entity
+	columns := append([]string{entity.Name}, opt.Features.Names()...)
 	err := dbutil.WithTransaction(db, ctx, func(ctx context.Context, tx *sqlx.Tx) error {
 		// create the data table
-		tableName := OnlineBatchTableName(opt.Revision.ID)
-		schema := dbutil.BuildTableSchema(tableName, opt.Entity, false, opt.Features, []string{opt.Entity.Name}, backend)
+		schema := dbutil.BuildTableSchema(tableName, entity, false, opt.Features, []string{entity.Name}, backend)
 		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return errdefs.WithStack(err)
@@ -32,19 +39,28 @@ func Import(ctx context.Context, db *sqlx.DB, opt online.ImportOpt, backend type
 			records = append(records, record)
 
 			if len(records) == importBatchSize {
-				if err := dbutil.InsertRecordsToTableTx(tx, ctx, tableName, records, columns, backend); err != nil {
+				if err = dbutil.InsertRecordsToTableTx(tx, ctx, tableName, records, columns, backend); err != nil {
 					return err
 				}
 				records = make([]interface{}, 0, importBatchSize)
 			}
 		}
-
-		if err := dbutil.InsertRecordsToTableTx(tx, ctx, tableName, records, columns, backend); err != nil {
+		if err = dbutil.InsertRecordsToTableTx(tx, ctx, tableName, records, columns, backend); err != nil {
 			return err
 		}
 		if opt.ExportError != nil {
 			return <-opt.ExportError
 		}
+		if opt.Group.Category == types.CategoryStream {
+			streamTableName := OnlineStreamTableName(opt.Group.ID)
+			if err = PurgeTx(ctx, tx, streamTableName, backend); err != nil {
+				return err
+			}
+			query := fmt.Sprintf(`ALTER TABLE %s RENAME TO %s;`, tableName, streamTableName)
+			if _, err = tx.ExecContext(ctx, query); err != nil {
+				return errdefs.WithStack(err)
+			}
+		}
 		return nil
 	})
 	return err
diff --git a/internal/database/online/sqlutil/purge.go b/internal/database/online/sqlutil/purge.go
index 41324f2c..ca7fcebb 100644
--- a/internal/database/online/sqlutil/purge.go
+++ b/internal/database/online/sqlutil/purge.go
@@ -17,3 +17,10 @@ func Purge(ctx context.Context, db *sqlx.DB, revisionID int, backend types.Backe
 	_, err := db.ExecContext(ctx, query)
 	return errdefs.WithStack(err)
 }
+
+func PurgeTx(ctx context.Context, tx *sqlx.Tx, tableName string, backend types.BackendType) error {
+	qt := dbutil.QuoteFn(backend)
+	query := fmt.Sprintf(`DROP TABLE IF EXISTS %s;`, qt(tableName))
+	_, err := tx.ExecContext(ctx, query)
+	return errdefs.WithStack(err)
+}
diff --git a/internal/database/online/store.go b/internal/database/online/store.go
index c99b6d9e..4efdbea4 100644
--- a/internal/database/online/store.go
+++ b/internal/database/online/store.go
@@ -12,7 +12,7 @@ type Store interface {
 	MultiGet(ctx context.Context, opt MultiGetOpt) (map[string]dbutil.RowMap, error)
 	Purge(ctx context.Context, revisionID int) error
 
-	// Batch import batch feature to online store
+	// Import batch / streaming features to online store
 	Import(ctx context.Context, opt ImportOpt) error
 
 	// Push streaming feature to online store
diff --git a/internal/database/online/types.go b/internal/database/online/types.go
index e0a3fada..21b2213c 100644
--- a/internal/database/online/types.go
+++ b/internal/database/online/types.go
@@ -21,11 +21,11 @@ type MultiGetOpt struct {
 }
 
 type ImportOpt struct {
+	Group        types.Group
+	Features     types.FeatureList
 	Revision     *types.Revision
-	Entity       *types.Entity
 	ExportStream <-chan types.ExportRecord
 	ExportError  <-chan error
-	Features     types.FeatureList
 }
 
 type PushOpt struct {
diff --git a/pkg/oomstore/sync.go b/pkg/oomstore/sync.go
index 5509f743..175e23c9 100644
--- a/pkg/oomstore/sync.go
+++ b/pkg/oomstore/sync.go
@@ -45,9 +45,9 @@ func (s *OomStore) Sync(ctx context.Context, opt types.SyncOpt) error {
 	}
 
 	if err = s.online.Import(ctx, online.ImportOpt{
+		Group:        *group,
 		Features:     features,
 		Revision:     revision,
-		Entity:       group.Entity,
 		ExportStream: exportResult.Data,
 		ExportError:  exportResult.GetErrorChannel(),
 	}); err != nil {


Test patch:
diff --git a/internal/database/online/test_impl/get.go b/internal/database/online/test_impl/get.go
index 27dc6c06..aeb1398f 100644
--- a/internal/database/online/test_impl/get.go
+++ b/internal/database/online/test_impl/get.go
@@ -17,7 +17,7 @@ func TestGetExisted(t *testing.T, prepareStore PrepareStoreFn, destroyStore Dest
 
 	for _, target := range s.Data {
 		opt := online.GetOpt{
-			Entity:     s.Entity,
+			Entity:     &s.Entity,
 			RevisionID: &s.Revision.ID,
 			Group:      s.Revision.Group,
 			Features:   s.Features,
@@ -41,7 +41,7 @@ func TestGetNotExistedEntityKey(t *testing.T, prepareStore PrepareStoreFn, destr
 	importSample(t, ctx, store, s)
 
 	rs, err := store.Get(ctx, online.GetOpt{
-		Entity:     s.Entity,
+		Entity:     &s.Entity,
 		RevisionID: &s.Revision.ID,
 		Group:      s.Revision.Group,
 		Features:   s.Features,
@@ -64,7 +64,7 @@ func TestMultiGet(t *testing.T, prepareStore PrepareStoreFn, destroystore Destro
 		keys = append(keys, r.EntityKey())
 	}
 	rs, err := store.MultiGet(ctx, online.MultiGetOpt{
-		Entity:     s.Entity,
+		Entity:     &s.Entity,
 		RevisionID: &s.Revision.ID,
 		Group:      s.Revision.Group,
 		Features:   s.Features,
diff --git a/internal/database/online/test_impl/purge.go b/internal/database/online/test_impl/purge.go
index 8fe2ce37..88bf4ebc 100644
--- a/internal/database/online/test_impl/purge.go
+++ b/internal/database/online/test_impl/purge.go
@@ -19,7 +19,7 @@ func TestPurgeRemovesSpecifiedRevision(t *testing.T, prepareStore PrepareStoreFn
 
 	for _, record := range SampleMedium.Data {
 		rs, err := store.Get(ctx, online.GetOpt{
-			Entity:     SampleMedium.Entity,
+			Entity:     &SampleMedium.Entity,
 			RevisionID: &SampleMedium.Revision.ID,
 			Group:      SampleMedium.Revision.Group,
 			EntityKey:  record.EntityKey(),
@@ -42,7 +42,7 @@ func TestPurgeNotRemovesOtherRevisions(t *testing.T, prepareStore PrepareStoreFn
 
 	for _, record := range SampleSmall.Data {
 		rs, err := store.Get(ctx, online.GetOpt{
-			Entity:     SampleSmall.Entity,
+			Entity:     &SampleSmall.Entity,
 			RevisionID: &SampleSmall.Revision.ID,
 			Group:      SampleSmall.Revision.Group,
 			EntityKey:  record.EntityKey(),
diff --git a/internal/database/online/test_impl/util.go b/internal/database/online/test_impl/util.go
index 223e1a71..c5135254 100644
--- a/internal/database/online/test_impl/util.go
+++ b/internal/database/online/test_impl/util.go
@@ -19,9 +19,10 @@ type PrepareStoreFn func(*testing.T) (context.Context, online.Store)
 type DestroyStoreFn func()
 
 type Sample struct {
+	Entity   types.Entity
+	Group    types.Group
 	Features types.FeatureList
-	Revision *types.Revision
-	Entity   *types.Entity
+	Revision types.Revision
 	Data     []types.ExportRecord
 }
 
@@ -30,55 +31,54 @@ var SampleMedium Sample
 
 func init() {
 	rand.Seed(time.Now().UnixNano())
-
+	entity := types.Entity{ID: 5, Name: "user"}
+	group := types.Group{ID: 1, Category: types.CategoryBatch, Entity: &entity}
 	{
 		SampleSmall = Sample{
+			Entity: entity,
+			Group:  group,
 			Features: types.FeatureList{
 				&types.Feature{
 					ID:        1,
 					Name:      "age",
 					GroupID:   1,
-					Group:     &types.Group{ID: 1, Category: types.CategoryBatch},
+					Group:     &group,
 					ValueType: types.Int64,
 				},
 				&types.Feature{
 					ID:        2,
 					Name:      "gender",
 					GroupID:   1,
-					Group:     &types.Group{ID: 1, Category: types.CategoryBatch},
+					Group:     &group,
 					ValueType: types.String,
 				},
 				&types.Feature{
 					ID:        3,
 					Name:      "account",
 					GroupID:   1,
-					Group:     &types.Group{ID: 1, Category: types.CategoryBatch},
+					Group:     &group,
 					ValueType: types.Float64,
 				},
 				&types.Feature{
 					ID:        4,
 					Name:      "is_active",
 					GroupID:   1,
-					Group:     &types.Group{ID: 1, Category: types.CategoryBatch},
+					Group:     &group,
 					ValueType: types.Bool,
 				},
 				&types.Feature{
 					ID:        5,
 					Name:      "register_time",
 					GroupID:   1,
-					Group:     &types.Group{ID: 1, Category: types.CategoryBatch},
+					Group:     &group,
 					ValueType: types.Time,
 				},
 			},
-			Revision: &types.Revision{
+			Revision: types.Revision{
 				ID:      3,
 				GroupID: 1,
-				Group: &types.Group{
-					ID:       1,
-					Category: types.CategoryBatch,
-				},
+				Group:   &group,
 			},
-			Entity: &types.Entity{ID: 5, Name: "user"},
 			Data: []types.ExportRecord{
 				[]interface{}{"3215", int64(18), "F", 1.1, true, time.Now()},
 				[]interface{}{"3216", int64(29), nil, 2.0, false, time.Now()},
@@ -99,15 +99,14 @@ func init() {
 			},
 		}
 
-		revision := &types.Revision{ID: 9, GroupID: 2, Group: &types.Group{ID: 2, Category: types.CategoryBatch}}
-		entity := &types.Entity{ID: 5, Name: "user"}
+		revision := types.Revision{ID: 9, GroupID: 2, Group: &types.Group{ID: 2, Category: types.CategoryBatch}}
 		var data []types.ExportRecord
 
 		for i := 0; i < 100; i++ {
 			record := []interface{}{dbutil.RandString(10), rand.Float64()}
 			data = append(data, record)
 		}
-		SampleMedium = Sample{features, revision, entity, data}
+		SampleMedium = Sample{entity, group, features, revision, data}
 	}
 }
 
@@ -122,9 +121,9 @@ func importSample(t *testing.T, ctx context.Context, store online.Store, samples
 		}(sample)
 
 		err := store.Import(ctx, online.ImportOpt{
+			Group:        sample.Group,
 			Features:     sample.Features,
-			Revision:     sample.Revision,
-			Entity:       sample.Entity,
+			Revision:     &sample.Revision,
 			ExportStream: stream,
 		})
 		require.NoError(t, err)