Repository: oom-ai/oomstore
Pull request #839 (2021-12-28 10:19:04)

Feat/PrepareStreamTable
this PR does:
1. online store add method `PrepareStreamStore`
2. pg/mysql/sqlite/cassandra/dynamodb implement method `PrepareStore`
3. create stream feature table after create stream group
4. add column to stream feature table after create stream feature

TODO
add unit test

relate #801 


Patch:
diff --git a/internal/database/online/cassandra/stream.go b/internal/database/online/cassandra/stream.go
new file mode 100644
index 00000000..382e2091
--- /dev/null
+++ b/internal/database/online/cassandra/stream.go
@@ -0,0 +1,32 @@
+package cassandra
+
+import (
+	"context"
+	"fmt"
+
+	"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/oomstore/types"
+)
+
+func (db *DB) PrepareStreamTable(ctx context.Context, opt online.PrepareStreamTableOpt) error {
+	tableName := sqlutil.OnlineStreamTableName(opt.GroupID)
+
+	if opt.Feature == nil {
+		schema, err := sqlutil.CreateStreamTableSchema(ctx, tableName, opt.Entity, types.BackendMySQL)
+		if err != nil {
+			return err
+		}
+
+		return db.Query(schema).WithContext(ctx).Exec()
+	}
+
+	dbValueType, err := dbutil.DBValueType(types.BackendCassandra, opt.Feature.ValueType)
+	if err != nil {
+		return err
+	}
+
+	sql := fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", tableName, opt.Feature.Name, dbValueType)
+	return db.Query(sql).WithContext(ctx).Exec()
+}
diff --git a/internal/database/online/dynamodb/stream.go b/internal/database/online/dynamodb/stream.go
new file mode 100644
index 00000000..62781ca4
--- /dev/null
+++ b/internal/database/online/dynamodb/stream.go
@@ -0,0 +1,36 @@
+package dynamodb
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
+	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
+	"github.com/oom-ai/oomstore/internal/database/online"
+	"github.com/oom-ai/oomstore/internal/database/online/sqlutil"
+)
+
+func (db *DB) PrepareStreamTable(ctx context.Context, opt online.PrepareStreamTableOpt) error {
+	tableName := sqlutil.OnlineStreamTableName(opt.GroupID)
+
+	_, err := db.CreateTable(ctx, &dynamodb.CreateTableInput{
+		TableName: aws.String(tableName),
+		KeySchema: []types.KeySchemaElement{
+			{
+				AttributeName: aws.String(opt.Entity.Name),
+				KeyType:       types.KeyTypeHash,
+			},
+		},
+		AttributeDefinitions: []types.AttributeDefinition{
+			{
+				AttributeName: aws.String(opt.Entity.Name),
+				AttributeType: types.ScalarAttributeTypeS,
+			},
+		},
+		ProvisionedThroughput: &types.ProvisionedThroughput{
+			ReadCapacityUnits:  aws.Int64(10),
+			WriteCapacityUnits: aws.Int64(10),
+		},
+	})
+	return err
+}
diff --git a/internal/database/online/mock_online/store.go b/internal/database/online/mock_online/store.go
index 022bbda4..ddc97541 100644
--- a/internal/database/online/mock_online/store.go
+++ b/internal/database/online/mock_online/store.go
@@ -108,6 +108,20 @@ func (mr *MockStoreMockRecorder) Ping(ctx interface{}) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Ping", reflect.TypeOf((*MockStore)(nil).Ping), ctx)
 }
 
+// PrepareStreamTable mocks base method.
+func (m *MockStore) PrepareStreamTable(ctx context.Context, opt online.PrepareStreamTableOpt) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PrepareStreamTable", ctx, opt)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// PrepareStreamTable indicates an expected call of PrepareStreamTable.
+func (mr *MockStoreMockRecorder) PrepareStreamTable(ctx, opt interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PrepareStreamTable", reflect.TypeOf((*MockStore)(nil).PrepareStreamTable), ctx, opt)
+}
+
 // Purge mocks base method.
 func (m *MockStore) Purge(ctx context.Context, revisionID int) error {
 	m.ctrl.T.Helper()
diff --git a/internal/database/online/mysql/store.go b/internal/database/online/mysql/store.go
index fad85788..3f76e8f1 100644
--- a/internal/database/online/mysql/store.go
+++ b/internal/database/online/mysql/store.go
@@ -47,3 +47,7 @@ func (db *DB) Purge(ctx context.Context, revisionID int) error {
 func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
 	panic("Implement me!")
 }
+
+func (db *DB) PrepareStreamTable(ctx context.Context, opt online.PrepareStreamTableOpt) error {
+	return sqlutil.SqlxPrapareStreamTable(ctx, db.DB, opt, types.BackendMySQL)
+}
diff --git a/internal/database/online/postgres/store.go b/internal/database/online/postgres/store.go
index 5b2c55c6..434676af 100644
--- a/internal/database/online/postgres/store.go
+++ b/internal/database/online/postgres/store.go
@@ -48,3 +48,7 @@ func (db *DB) Purge(ctx context.Context, revisionID int) error {
 func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
 	panic("Implement me!")
 }
+
+func (db *DB) PrepareStreamTable(ctx context.Context, opt online.PrepareStreamTableOpt) error {
+	return sqlutil.SqlxPrapareStreamTable(ctx, db.DB, opt, types.BackendPostgres)
+}
diff --git a/internal/database/online/redis/store.go b/internal/database/online/redis/store.go
index 32b73114..5573efec 100644
--- a/internal/database/online/redis/store.go
+++ b/internal/database/online/redis/store.go
@@ -32,6 +32,10 @@ func Open(opt *types.RedisOpt) *DB {
 	return &DB{redis.NewClient(&redisOpt)}
 }
 
+func (db *DB) PrepareStreamTable(ctx context.Context, opt online.PrepareStreamTableOpt) error {
+	return nil
+}
+
 func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
 	panic("Implement me!")
 }
diff --git a/internal/database/online/sqlite/store.go b/internal/database/online/sqlite/store.go
index 0cc7443c..334a8041 100644
--- a/internal/database/online/sqlite/store.go
+++ b/internal/database/online/sqlite/store.go
@@ -47,3 +47,7 @@ func (db *DB) Purge(ctx context.Context, revisionID int) error {
 func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
 	panic("Implement me!")
 }
+
+func (db *DB) PrepareStreamTable(ctx context.Context, opt online.PrepareStreamTableOpt) error {
+	return sqlutil.SqlxPrapareStreamTable(ctx, db.DB, opt, types.BackendSQLite)
+}
diff --git a/internal/database/online/sqlutil/util.go b/internal/database/online/sqlutil/util.go
index 9b9bbf3f..2d933bce 100644
--- a/internal/database/online/sqlutil/util.go
+++ b/internal/database/online/sqlutil/util.go
@@ -1,8 +1,13 @@
 package sqlutil
 
 import (
+	"context"
 	"fmt"
 
+	"github.com/jmoiron/sqlx"
+	"github.com/oom-ai/oomstore/internal/database/dbutil"
+	"github.com/oom-ai/oomstore/internal/database/online"
+	"github.com/oom-ai/oomstore/pkg/errdefs"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
@@ -42,3 +47,47 @@ func deserializeByTag(i interface{}, valueType types.ValueType, backend types.Ba
 		return i, nil
 	}
 }
+
+func CreateStreamTableSchema(ctx context.Context, tableName string, entity *types.Entity, backend types.BackendType) (string, error) {
+	var (
+		entityFormat string
+	)
+	switch backend {
+	case types.BackendMySQL:
+		entityFormat = fmt.Sprintf("`%s` VARCHAR(%d) PRIMARY KEY", entity.Name, entity.Length)
+	case types.BackendPostgres:
+		entityFormat = fmt.Sprintf(`"%s" VARCHAR(%d) PRIMARY KEY`, entity.Name, entity.Length)
+	case types.BackendCassandra, types.BackendSQLite:
+		entityFormat = fmt.Sprintf(`"%s" TEXY PRIMARY KEY`, entity.Name)
+	default:
+		return "", errdefs.InvalidAttribute(fmt.Errorf("backend %s not support", backend))
+	}
+
+	schema := fmt.Sprintf(`
+CREATE TABLE %s (
+   %s
+)`, tableName, entityFormat)
+	return schema, nil
+}
+
+func SqlxPrapareStreamTable(ctx context.Context, db *sqlx.DB, opt online.PrepareStreamTableOpt, backend types.BackendType) error {
+	tableName := OnlineStreamTableName(opt.GroupID)
+
+	if opt.Feature == nil {
+		schema, err := CreateStreamTableSchema(ctx, tableName, opt.Entity, backend)
+		if err != nil {
+			return err
+		}
+		_, err = db.ExecContext(ctx, schema)
+		return err
+	}
+
+	dbValueType, err := dbutil.DBValueType(backend, opt.Feature.ValueType)
+	if err != nil {
+		return err
+	}
+
+	sql := fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", tableName, opt.Feature.Name, dbValueType)
+	_, err = db.ExecContext(ctx, sql)
+	return err
+}
diff --git a/internal/database/online/store.go b/internal/database/online/store.go
index b0b4f35d..c99b6d9e 100644
--- a/internal/database/online/store.go
+++ b/internal/database/online/store.go
@@ -16,8 +16,11 @@ type Store interface {
 	Import(ctx context.Context, opt ImportOpt) error
 
 	// Push streaming feature to online store
+	// Note: Make sure that the table corresponding to the stream feature already exists before executing this method
 	Push(ctx context.Context, opt PushOpt) error
 
+	PrepareStreamTable(ctx context.Context, opt PrepareStreamTableOpt) error
+
 	Ping(ctx context.Context) error
 	io.Closer
 }
diff --git a/internal/database/online/tikv/store.go b/internal/database/online/tikv/store.go
index 09d6a0e0..1b8d891a 100644
--- a/internal/database/online/tikv/store.go
+++ b/internal/database/online/tikv/store.go
@@ -3,11 +3,12 @@ package tikv
 import (
 	"context"
 
+	"github.com/tikv/client-go/v2/config"
+	"github.com/tikv/client-go/v2/rawkv"
+
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
 	"github.com/oom-ai/oomstore/internal/database/online"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
-	"github.com/tikv/client-go/v2/config"
-	"github.com/tikv/client-go/v2/rawkv"
 )
 
 const BackendType = types.BackendTiKV
@@ -54,3 +55,7 @@ func (db *DB) Purge(ctx context.Context, revisionID int) error {
 func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
 	panic("implement me")
 }
+
+func (db *DB) PrepareStreamTable(ctx context.Context, opt online.PrepareStreamTableOpt) error {
+	return nil
+}
diff --git a/internal/database/online/types.go b/internal/database/online/types.go
index e2d8e475..d3cbb1bf 100644
--- a/internal/database/online/types.go
+++ b/internal/database/online/types.go
@@ -30,3 +30,13 @@ type PushOpt struct {
 	Entity  *types.Entity
 	Feature types.Feature
 }
+
+type PrepareStreamTableOpt struct {
+	Entity *types.Entity
+
+	GroupID int
+
+	// Feature is not nil to add a new row to the stream table;
+	// otherwise it means the stream table will be created.
+	Feature *types.Feature
+}
diff --git a/pkg/oomstore/apply.go b/pkg/oomstore/apply.go
index 34e26014..e819bc99 100644
--- a/pkg/oomstore/apply.go
+++ b/pkg/oomstore/apply.go
@@ -10,6 +10,7 @@ import (
 	"gopkg.in/yaml.v3"
 
 	"github.com/oom-ai/oomstore/internal/database/metadata"
+	"github.com/oom-ai/oomstore/internal/database/online"
 	"github.com/oom-ai/oomstore/pkg/errdefs"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types/apply"
@@ -90,13 +91,24 @@ func (s *OomStore) applyGroup(ctx context.Context, tx metadata.DBStore, newGroup
 		if !errdefs.IsNotFound(err) {
 			return err
 		}
-		_, err = tx.CreateGroup(ctx, metadata.CreateGroupOpt{
+
+		id, err := tx.CreateGroup(ctx, metadata.CreateGroupOpt{
 			GroupName:   newGroup.Name,
 			EntityID:    entity.ID,
 			Category:    newGroup.Category,
 			Description: newGroup.Description,
 		})
-		return err
+		if err != nil {
+			return err
+		}
+
+		if newGroup.Category == types.CategoryStream {
+			return s.online.PrepareStreamTable(ctx, online.PrepareStreamTableOpt{
+				Entity:  entity,
+				GroupID: id,
+			})
+		}
+		return nil
 	}
 
 	if newGroup.Description != group.Description {
@@ -118,7 +130,9 @@ func (s *OomStore) applyFeature(ctx context.Context, tx metadata.DBStore, newFea
 	if err != nil {
 		return err
 	}
-	feature, err := tx.GetFeatureByName(ctx, newFeature.Name)
+
+	featureFullName := fmt.Sprintf("%s.%s", newFeature.GroupName, newFeature.Name)
+	feature, err := tx.GetFeatureByName(ctx, featureFullName)
 	if err != nil {
 		if !errdefs.IsNotFound(err) {
 			return err
@@ -127,14 +141,29 @@ func (s *OomStore) applyFeature(ctx context.Context, tx metadata.DBStore, newFea
 		if err != nil {
 			return err
 		}
-		_, err = tx.CreateFeature(ctx, metadata.CreateFeatureOpt{
+		id, err := tx.CreateFeature(ctx, metadata.CreateFeatureOpt{
 			FeatureName: newFeature.Name,
-			FullName:    fmt.Sprintf("%s.%s", group.Name, newFeature.Name),
+			FullName:    featureFullName,
 			GroupID:     group.ID,
 			Description: newFeature.Description,
 			ValueType:   valueType,
 		})
-		return err
+		if err != nil {
+			return err
+		}
+
+		if group.Category == types.CategoryStream {
+			feature, err := tx.GetFeature(ctx, id)
+			if err != nil {
+				return err
+			}
+			return s.online.PrepareStreamTable(ctx, online.PrepareStreamTableOpt{
+				Entity:  group.Entity,
+				GroupID: group.ID,
+				Feature: feature,
+			})
+		}
+		return nil
 	}
 
 	if newFeature.Description != feature.Description {
diff --git a/pkg/oomstore/feature.go b/pkg/oomstore/feature.go
index 780ca650..3da245ca 100644
--- a/pkg/oomstore/feature.go
+++ b/pkg/oomstore/feature.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"github.com/oom-ai/oomstore/internal/database/metadata"
+	"github.com/oom-ai/oomstore/internal/database/online"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
@@ -59,11 +60,30 @@ func (s *OomStore) CreateFeature(ctx context.Context, opt types.CreateFeatureOpt
 		return 0, err
 	}
 
-	return s.metadata.CreateFeature(ctx, metadata.CreateFeatureOpt{
+	id, err := s.metadata.CreateFeature(ctx, metadata.CreateFeatureOpt{
 		FeatureName: opt.FeatureName,
 		FullName:    fmt.Sprintf("%s.%s", group.Name, opt.FeatureName),
 		GroupID:     group.ID,
 		ValueType:   opt.ValueType,
 		Description: opt.Description,
 	})
+	if err != nil {
+		return 0, err
+	}
+
+	if group.Category == types.CategoryStream {
+		feature, err := s.metadata.GetFeature(ctx, id)
+		if err != nil {
+			return 0, err
+		}
+		if err := s.online.PrepareStreamTable(ctx, online.PrepareStreamTableOpt{
+			Entity:  group.Entity,
+			GroupID: group.ID,
+			Feature: feature,
+		}); err != nil {
+			return 0, err
+		}
+	}
+
+	return id, nil
 }
diff --git a/pkg/oomstore/group.go b/pkg/oomstore/group.go
index b4633268..5e75d7b3 100644
--- a/pkg/oomstore/group.go
+++ b/pkg/oomstore/group.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/oom-ai/oomstore/internal/database/metadata"
+	"github.com/oom-ai/oomstore/internal/database/online"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
@@ -14,12 +15,25 @@ func (s *OomStore) CreateGroup(ctx context.Context, opt types.CreateGroupOpt) (i
 		return 0, err
 	}
 
-	return s.metadata.CreateGroup(ctx, metadata.CreateGroupOpt{
+	id, err := s.metadata.CreateGroup(ctx, metadata.CreateGroupOpt{
 		GroupName:   opt.GroupName,
 		EntityID:    entity.ID,
 		Description: opt.Description,
 		Category:    opt.Category,
 	})
+	if err != nil {
+		return 0, err
+	}
+
+	if opt.Category == types.CategoryStream {
+		if err = s.online.PrepareStreamTable(ctx, online.PrepareStreamTableOpt{
+			Entity:  entity,
+			GroupID: id,
+		}); err != nil {
+			return 0, err
+		}
+	}
+	return id, nil
 }
 
 // Get metadata of a feature group by ID.


Test patch:
diff --git a/pkg/oomstore/feature_test.go b/pkg/oomstore/feature_test.go
index fe75406a..ec67e8bf 100644
--- a/pkg/oomstore/feature_test.go
+++ b/pkg/oomstore/feature_test.go
@@ -11,6 +11,8 @@ import (
 	"github.com/oom-ai/oomstore/internal/database/metadata"
 	"github.com/oom-ai/oomstore/internal/database/metadata/mock_metadata"
 	"github.com/oom-ai/oomstore/internal/database/offline/mock_offline"
+	"github.com/oom-ai/oomstore/internal/database/online"
+	"github.com/oom-ai/oomstore/internal/database/online/mock_online"
 	"github.com/oom-ai/oomstore/pkg/oomstore"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
@@ -20,9 +22,10 @@ func TestCreateFeature(t *testing.T) {
 	defer ctrl.Finish()
 	ctx := context.Background()
 
+	onlineStore := mock_online.NewMockStore(ctrl)
 	offlineStore := mock_offline.NewMockStore(ctrl)
 	metadataStore := mock_metadata.NewMockStore(ctrl)
-	store := oomstore.TEST__New(nil, offlineStore, metadataStore)
+	store := oomstore.TEST__New(onlineStore, offlineStore, metadataStore)
 
 	testCases := []struct {
 		description string
@@ -72,8 +75,17 @@ func TestCreateFeature(t *testing.T) {
 				ValueType:   tc.valueType,
 				Description: tc.opt.Description,
 			}
-			metadataStore.EXPECT().CreateFeature(ctx, metadataOpt).Return(0, nil)
+			metadataStore.EXPECT().CreateFeature(ctx, metadataOpt).Return(1, nil)
 
+			if tc.group.Category == types.CategoryStream {
+				var feature types.Feature
+				metadataStore.EXPECT().GetFeature(ctx, 1).Return(&feature, nil)
+				onlineStore.EXPECT().PrepareStreamTable(ctx, online.PrepareStreamTableOpt{
+					Entity:  tc.group.Entity,
+					GroupID: tc.group.ID,
+					Feature: &feature,
+				})
+			}
 			_, err := store.CreateFeature(ctx, tc.opt)
 			assert.NoError(t, err)
 		})