Repository: oom-ai/oomstore
Pull request #1093 (2022-01-26 07:34:34)

Jinghan/reorganize helpers of online store
This PR does:
- rename file `stream.go` to `push.go`
- reuse `online.CreateTable` in cassandra's Import
- move `OnlineBatch(Stream)TableName` to the package dbutil, so that it is consistent with offline table name helpers


Patch:
diff --git a/internal/database/dbutil/util.go b/internal/database/dbutil/util.go
index d6b0a236..ea97c2ee 100644
--- a/internal/database/dbutil/util.go
+++ b/internal/database/dbutil/util.go
@@ -50,3 +50,11 @@ func OfflineStreamSnapshotTableName(groupID int, revision int64) string {
 func OfflineStreamCdcTableName(groupID int, revision int64) string {
 	return fmt.Sprintf("offline_stream_cdc_%d_%d", groupID, revision)
 }
+
+func OnlineBatchTableName(revisionID int) string {
+	return fmt.Sprintf("online_batch_%d", revisionID)
+}
+
+func OnlineStreamTableName(groupID int) string {
+	return fmt.Sprintf("online_stream_%d", groupID)
+}
diff --git a/internal/database/online/cassandra/get.go b/internal/database/online/cassandra/get.go
index 4ef995e2..3a46ee08 100644
--- a/internal/database/online/cassandra/get.go
+++ b/internal/database/online/cassandra/get.go
@@ -10,16 +10,15 @@ import (
 
 	"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) Get(ctx context.Context, opt online.GetOpt) (dbutil.RowMap, error) {
 	var tableName string
 	if opt.Group.Category == types.CategoryBatch {
-		tableName = sqlutil.OnlineBatchTableName(*opt.RevisionID)
+		tableName = dbutil.OnlineBatchTableName(*opt.RevisionID)
 	} else {
-		tableName = sqlutil.OnlineStreamTableName(opt.Group.ID)
+		tableName = dbutil.OnlineStreamTableName(opt.Group.ID)
 	}
 
 	query := fmt.Sprintf(`SELECT %s FROM %s WHERE %s = ?`,
@@ -52,9 +51,9 @@ func (db *DB) MultiGet(ctx context.Context, opt online.MultiGetOpt) (map[string]
 	)
 
 	if opt.Group.Category == types.CategoryBatch {
-		tableName = sqlutil.OnlineBatchTableName(*opt.RevisionID)
+		tableName = dbutil.OnlineBatchTableName(*opt.RevisionID)
 	} else {
-		tableName = sqlutil.OnlineStreamTableName(opt.Group.ID)
+		tableName = dbutil.OnlineStreamTableName(opt.Group.ID)
 	}
 	entityName := opt.Group.Entity.Name
 	query := fmt.Sprintf(`SELECT %s, %s FROM %s WHERE %s in (%s)`,
diff --git a/internal/database/online/cassandra/import.go b/internal/database/online/cassandra/import.go
index 4004f089..7da7241f 100644
--- a/internal/database/online/cassandra/import.go
+++ b/internal/database/online/cassandra/import.go
@@ -9,7 +9,6 @@ import (
 
 	"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"
 )
@@ -18,23 +17,25 @@ func (db *DB) Import(ctx context.Context, opt online.ImportOpt) error {
 	// Step 0: drop existing table for streaming feature
 	var tableName string
 	if opt.Group.Category == types.CategoryBatch {
-		tableName = sqlutil.OnlineBatchTableName(*opt.RevisionID)
+		tableName = dbutil.OnlineBatchTableName(*opt.RevisionID)
 	} else {
-		tableName = sqlutil.OnlineStreamTableName(opt.Group.ID)
+		tableName = dbutil.OnlineStreamTableName(opt.Group.ID)
 		if err := db.Query(fmt.Sprintf(`DROP TABLE IF EXISTS %s;`, tableName)).Exec(); err != nil {
 			return errdefs.WithStack(err)
 		}
 	}
 
 	// Step 1: create online table
-	entity := opt.Group.Entity
-	columns := append([]string{entity.Name}, opt.Features.Names()...)
-	schema := dbutil.BuildTableSchema(tableName, entity.Name, false, opt.Features, []string{entity.Name}, Backend)
-	if err := db.Query(schema).Exec(); err != nil {
-		return errdefs.WithStack(err)
+	if err := db.CreateTable(ctx, online.CreateTableOpt{
+		EntityName: opt.Group.Entity.Name,
+		TableName:  tableName,
+		Features:   opt.Features,
+	}); err != nil {
+		return err
 	}
 
 	// Step 2: insert records to the online table
+	columns := append([]string{opt.Group.Entity.Name}, opt.Features.Names()...)
 	insertStmt := buildInsertStatement(tableName, columns)
 	batch := db.NewBatch(gocql.LoggedBatch)
 	for record := range opt.ExportStream {
diff --git a/internal/database/online/cassandra/purge.go b/internal/database/online/cassandra/purge.go
index 428d2106..445d2b17 100644
--- a/internal/database/online/cassandra/purge.go
+++ b/internal/database/online/cassandra/purge.go
@@ -4,13 +4,13 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/oom-ai/oomstore/pkg/errdefs"
+	"github.com/oom-ai/oomstore/internal/database/dbutil"
 
-	"github.com/oom-ai/oomstore/internal/database/online/sqlutil"
+	"github.com/oom-ai/oomstore/pkg/errdefs"
 )
 
 func (db *DB) Purge(ctx context.Context, revisionID int) error {
-	err := db.Query(fmt.Sprintf("DROP TABLE IF EXISTS %s", sqlutil.OnlineBatchTableName(revisionID))).
+	err := db.Query(fmt.Sprintf("DROP TABLE IF EXISTS %s", dbutil.OnlineBatchTableName(revisionID))).
 		WithContext(ctx).Exec()
 	return errdefs.WithStack(err)
 }
diff --git a/internal/database/online/cassandra/stream.go b/internal/database/online/cassandra/push.go
similarity index 86%
rename from internal/database/online/cassandra/stream.go
rename to internal/database/online/cassandra/push.go
index 4f99ee18..94824d9e 100644
--- a/internal/database/online/cassandra/stream.go
+++ b/internal/database/online/cassandra/push.go
@@ -4,13 +4,14 @@ 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/errdefs"
 )
 
 func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
-	tableName := sqlutil.OnlineStreamTableName(opt.GroupID)
+	tableName := dbutil.OnlineStreamTableName(opt.GroupID)
 
 	cond := sqlutil.BuildPushCondition(opt, Backend)
 	// cassandra's `insert` is equivalent to `insert_or_update`.
diff --git a/internal/database/online/dynamodb/get.go b/internal/database/online/dynamodb/get.go
index e737c6c8..b827edb6 100644
--- a/internal/database/online/dynamodb/get.go
+++ b/internal/database/online/dynamodb/get.go
@@ -11,7 +11,6 @@ import (
 
 	"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"
 	oomTypes "github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
@@ -23,9 +22,9 @@ const (
 func (db *DB) Get(ctx context.Context, opt online.GetOpt) (dbutil.RowMap, error) {
 	var tableName string
 	if opt.Group.Category == oomTypes.CategoryBatch {
-		tableName = sqlutil.OnlineBatchTableName(*opt.RevisionID)
+		tableName = dbutil.OnlineBatchTableName(*opt.RevisionID)
 	} else {
-		tableName = sqlutil.OnlineStreamTableName(opt.Group.ID)
+		tableName = dbutil.OnlineStreamTableName(opt.Group.ID)
 	}
 
 	entityKeyValue, err := attributevalue.Marshal(opt.EntityKey)
@@ -51,9 +50,9 @@ func (db *DB) Get(ctx context.Context, opt online.GetOpt) (dbutil.RowMap, error)
 func (db *DB) MultiGet(ctx context.Context, opt online.MultiGetOpt) (map[string]dbutil.RowMap, error) {
 	var tableName string
 	if opt.Group.Category == oomTypes.CategoryBatch {
-		tableName = sqlutil.OnlineBatchTableName(*opt.RevisionID)
+		tableName = dbutil.OnlineBatchTableName(*opt.RevisionID)
 	} else {
-		tableName = sqlutil.OnlineStreamTableName(opt.Group.ID)
+		tableName = dbutil.OnlineStreamTableName(opt.Group.ID)
 	}
 
 	entityName := opt.Group.Entity.Name
diff --git a/internal/database/online/dynamodb/import.go b/internal/database/online/dynamodb/import.go
index acb1b564..8de93833 100644
--- a/internal/database/online/dynamodb/import.go
+++ b/internal/database/online/dynamodb/import.go
@@ -11,7 +11,6 @@ import (
 	"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"
 	"github.com/oom-ai/oomstore/pkg/errdefs"
 	oomTypes "github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
@@ -24,9 +23,9 @@ 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.RevisionID)
+		tableName = dbutil.OnlineBatchTableName(*opt.RevisionID)
 	} else {
-		tableName = sqlutil.OnlineStreamTableName(opt.Group.ID)
+		tableName = dbutil.OnlineStreamTableName(opt.Group.ID)
 		_, err := db.DeleteTable(ctx, &dynamodb.DeleteTableInput{
 			TableName: aws.String(tableName),
 		})
diff --git a/internal/database/online/dynamodb/purge.go b/internal/database/online/dynamodb/purge.go
index 77b381bd..7f0b5450 100644
--- a/internal/database/online/dynamodb/purge.go
+++ b/internal/database/online/dynamodb/purge.go
@@ -3,15 +3,16 @@ package dynamodb
 import (
 	"context"
 
+	"github.com/oom-ai/oomstore/internal/database/dbutil"
+
 	"github.com/oom-ai/oomstore/pkg/errdefs"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
-	"github.com/oom-ai/oomstore/internal/database/online/sqlutil"
 )
 
 func (db *DB) Purge(ctx context.Context, revisionID int) error {
-	tableName := sqlutil.OnlineBatchTableName(revisionID)
+	tableName := dbutil.OnlineBatchTableName(revisionID)
 	_, err := db.DeleteTable(ctx, &dynamodb.DeleteTableInput{
 		TableName: aws.String(tableName),
 	})
diff --git a/internal/database/online/dynamodb/stream.go b/internal/database/online/dynamodb/push.go
similarity index 82%
rename from internal/database/online/dynamodb/stream.go
rename to internal/database/online/dynamodb/push.go
index bc1c3a8d..55df61c5 100644
--- a/internal/database/online/dynamodb/stream.go
+++ b/internal/database/online/dynamodb/push.go
@@ -11,12 +11,11 @@ import (
 	"github.com/oom-ai/oomstore/pkg/errdefs"
 
 	"github.com/oom-ai/oomstore/internal/database/online"
-	"github.com/oom-ai/oomstore/internal/database/online/sqlutil"
 )
 
 func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
 	var (
-		tableName = sqlutil.OnlineStreamTableName(opt.GroupID)
+		tableName = dbutil.OnlineStreamTableName(opt.GroupID)
 		item      = make(map[string]types.AttributeValue)
 	)
 
@@ -31,11 +30,11 @@ func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
 		if err != nil {
 			return err
 		}
-		attributevalue, err := attributevalue.Marshal(value)
+		attributeValue, err := attributevalue.Marshal(value)
 		if err != nil {
 			return errdefs.WithStack(err)
 		}
-		item[feature.Name] = attributevalue
+		item[feature.Name] = attributeValue
 	}
 
 	_, err = db.Client.PutItem(ctx, &dynamodb.PutItemInput{
diff --git a/internal/database/online/mysql/store.go b/internal/database/online/mysql/store.go
index 7a4f7c9e..c28a822f 100644
--- a/internal/database/online/mysql/store.go
+++ b/internal/database/online/mysql/store.go
@@ -51,7 +51,7 @@ func (db *DB) Purge(ctx context.Context, revisionID int) error {
 
 // TODO: refactor with text/template
 func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
-	tableName := sqlutil.OnlineStreamTableName(opt.GroupID)
+	tableName := dbutil.OnlineStreamTableName(opt.GroupID)
 
 	cond := sqlutil.BuildPushCondition(opt, Backend)
 	query := fmt.Sprintf(`INSERT INTO %s (%s) VALUES(%s) ON DUPLICATE KEY UPDATE %s`,
diff --git a/internal/database/online/postgres/store.go b/internal/database/online/postgres/store.go
index a2419b46..9b709388 100644
--- a/internal/database/online/postgres/store.go
+++ b/internal/database/online/postgres/store.go
@@ -50,7 +50,7 @@ func (db *DB) Purge(ctx context.Context, revisionID int) error {
 }
 
 func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
-	tableName := sqlutil.OnlineStreamTableName(opt.GroupID)
+	tableName := dbutil.OnlineStreamTableName(opt.GroupID)
 
 	cond := sqlutil.BuildPushCondition(opt, Backend)
 
diff --git a/internal/database/online/sqlite/store.go b/internal/database/online/sqlite/store.go
index afc49ed3..7ca8784d 100644
--- a/internal/database/online/sqlite/store.go
+++ b/internal/database/online/sqlite/store.go
@@ -49,7 +49,7 @@ func (db *DB) Purge(ctx context.Context, revisionID int) error {
 }
 
 func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
-	tableName := sqlutil.OnlineStreamTableName(opt.GroupID)
+	tableName := dbutil.OnlineStreamTableName(opt.GroupID)
 
 	cond := sqlutil.BuildPushCondition(opt, Backend)
 
diff --git a/internal/database/online/sqlutil/get.go b/internal/database/online/sqlutil/get.go
index 4f3e4430..b199937a 100644
--- a/internal/database/online/sqlutil/get.go
+++ b/internal/database/online/sqlutil/get.go
@@ -16,9 +16,9 @@ import (
 func Get(ctx context.Context, db *sqlx.DB, opt online.GetOpt, backend types.BackendType) (dbutil.RowMap, error) {
 	var tableName string
 	if opt.Group.Category == types.CategoryBatch {
-		tableName = OnlineBatchTableName(*opt.RevisionID)
+		tableName = dbutil.OnlineBatchTableName(*opt.RevisionID)
 	} else {
-		tableName = OnlineStreamTableName(opt.Group.ID)
+		tableName = dbutil.OnlineStreamTableName(opt.Group.ID)
 	}
 
 	entityName := opt.Group.Entity.Name
@@ -45,9 +45,9 @@ func Get(ctx context.Context, db *sqlx.DB, opt online.GetOpt, backend types.Back
 func MultiGet(ctx context.Context, db *sqlx.DB, opt online.MultiGetOpt, backend types.BackendType) (map[string]dbutil.RowMap, error) {
 	var tableName string
 	if opt.Group.Category == types.CategoryBatch {
-		tableName = OnlineBatchTableName(*opt.RevisionID)
+		tableName = dbutil.OnlineBatchTableName(*opt.RevisionID)
 	} else {
-		tableName = OnlineStreamTableName(opt.Group.ID)
+		tableName = dbutil.OnlineStreamTableName(opt.Group.ID)
 	}
 
 	entityName := opt.Group.Entity.Name
diff --git a/internal/database/online/sqlutil/import.go b/internal/database/online/sqlutil/import.go
index 70153ed3..ba2c9ce7 100644
--- a/internal/database/online/sqlutil/import.go
+++ b/internal/database/online/sqlutil/import.go
@@ -18,7 +18,7 @@ func Import(ctx context.Context, db *sqlx.DB, opt online.ImportOpt, backend type
 	if opt.Group.Category == types.CategoryStream {
 		tableName = dbutil.TempTable("online_stream")
 	} else {
-		tableName = OnlineBatchTableName(*opt.RevisionID)
+		tableName = dbutil.OnlineBatchTableName(*opt.RevisionID)
 	}
 	entity := opt.Group.Entity
 	columns := append([]string{entity.Name}, opt.Features.Names()...)
@@ -52,7 +52,7 @@ func Import(ctx context.Context, db *sqlx.DB, opt online.ImportOpt, backend type
 			return <-opt.ExportError
 		}
 		if opt.Group.Category == types.CategoryStream {
-			streamTableName := OnlineStreamTableName(opt.Group.ID)
+			streamTableName := dbutil.OnlineStreamTableName(opt.Group.ID)
 			if err = PurgeTx(ctx, tx, streamTableName, backend); err != nil {
 				return err
 			}
diff --git a/internal/database/online/sqlutil/purge.go b/internal/database/online/sqlutil/purge.go
index ca7fcebb..dcae09a5 100644
--- a/internal/database/online/sqlutil/purge.go
+++ b/internal/database/online/sqlutil/purge.go
@@ -13,7 +13,7 @@ import (
 
 func Purge(ctx context.Context, db *sqlx.DB, revisionID int, backend types.BackendType) error {
 	qt := dbutil.QuoteFn(backend)
-	query := fmt.Sprintf(`DROP TABLE IF EXISTS %s;`, qt(OnlineBatchTableName(revisionID)))
+	query := fmt.Sprintf(`DROP TABLE IF EXISTS %s;`, qt(dbutil.OnlineBatchTableName(revisionID)))
 	_, err := db.ExecContext(ctx, query)
 	return errdefs.WithStack(err)
 }
diff --git a/internal/database/online/sqlutil/util.go b/internal/database/online/sqlutil/util.go
deleted file mode 100644
index 65c09d53..00000000
--- a/internal/database/online/sqlutil/util.go
+++ /dev/null
@@ -1,13 +0,0 @@
-package sqlutil
-
-import (
-	"fmt"
-)
-
-func OnlineBatchTableName(revisionID int) string {
-	return fmt.Sprintf("online_batch_%d", revisionID)
-}
-
-func OnlineStreamTableName(groupID int) string {
-	return fmt.Sprintf("online_stream_%d", groupID)
-}
diff --git a/pkg/oomstore/apply.go b/pkg/oomstore/apply.go
index e56e0041..df20bf7e 100644
--- a/pkg/oomstore/apply.go
+++ b/pkg/oomstore/apply.go
@@ -4,7 +4,7 @@ import (
 	"context"
 	"io"
 
-	"github.com/oom-ai/oomstore/internal/database/online/sqlutil"
+	"github.com/oom-ai/oomstore/internal/database/dbutil"
 
 	"github.com/mitchellh/mapstructure"
 	"github.com/spf13/cast"
@@ -195,7 +195,7 @@ func (s *OomStore) applyFeature(ctx context.Context, tx metadata.DBStore, newFea
 			return func() error {
 				return s.online.CreateTable(ctx, online.CreateTableOpt{
 					EntityName: group.Entity.Name,
-					TableName:  sqlutil.OnlineStreamTableName(group.ID),
+					TableName:  dbutil.OnlineStreamTableName(group.ID),
 					Features:   features,
 				})
 			}, nil
diff --git a/pkg/oomstore/feature.go b/pkg/oomstore/feature.go
index e2ed5e08..6f1fd0d8 100644
--- a/pkg/oomstore/feature.go
+++ b/pkg/oomstore/feature.go
@@ -3,7 +3,7 @@ package oomstore
 import (
 	"context"
 
-	"github.com/oom-ai/oomstore/internal/database/online/sqlutil"
+	"github.com/oom-ai/oomstore/internal/database/dbutil"
 
 	"github.com/oom-ai/oomstore/internal/database/metadata"
 	"github.com/oom-ai/oomstore/internal/database/online"
@@ -107,7 +107,7 @@ func (s *OomStore) CreateFeature(ctx context.Context, opt types.CreateFeatureOpt
 		}
 		if err = s.online.CreateTable(ctx, online.CreateTableOpt{
 			EntityName: group.Entity.Name,
-			TableName:  sqlutil.OnlineStreamTableName(group.ID),
+			TableName:  dbutil.OnlineStreamTableName(group.ID),
 			Features:   features,
 		}); err != nil {
 			return 0, err


Test patch:
diff --git a/internal/database/online/test_impl/push.go b/internal/database/online/test_impl/push.go
index 40e85854..7b7b5105 100644
--- a/internal/database/online/test_impl/push.go
+++ b/internal/database/online/test_impl/push.go
@@ -3,7 +3,7 @@ package test_impl
 import (
 	"testing"
 
-	"github.com/oom-ai/oomstore/internal/database/online/sqlutil"
+	"github.com/oom-ai/oomstore/internal/database/dbutil"
 
 	"github.com/oom-ai/oomstore/internal/database/online"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
@@ -23,7 +23,7 @@ func TestPush(t *testing.T, prepareStore PrepareStoreFn, destoryStore DestroySto
 
 	assert.NoError(t, store.CreateTable(ctx, online.CreateTableOpt{
 		EntityName: entity.Name,
-		TableName:  sqlutil.OnlineStreamTableName(group.ID),
+		TableName:  dbutil.OnlineStreamTableName(group.ID),
 		Features:   SampleStream.Features,
 	}))