Repository: oom-ai/oomstore
Pull request #892 (2022-01-04 07:22:32)

feature: offline table index
Template for table schema becomes very complicated when supporting optional primary key and index. So I rewrite it using normal functions.


Patch:
diff --git a/internal/database/database.go b/internal/database/database.go
index 30266104..0a9db934 100644
--- a/internal/database/database.go
+++ b/internal/database/database.go
@@ -37,8 +37,6 @@ func OpenOnlineStore(opt types.OnlineStoreConfig) (online.Store, error) {
 		return onlineRedis.Open(opt.Redis), nil
 	case types.BackendMySQL:
 		return onlineMySQL.Open(opt.MySQL)
-	case types.BackendTiDB:
-		return onlineMySQL.Open(opt.TiDB)
 	case types.BackendDynamoDB:
 		return onlineDynamoDB.Open(opt.DynamoDB)
 	case types.BackendCassandra:
@@ -58,8 +56,6 @@ func OpenMetadataStore(opt types.MetadataStoreConfig) (metadata.Store, error) {
 		return metadataPG.Open(context.Background(), opt.Postgres)
 	case types.BackendMySQL:
 		return metadataMySQL.Open(context.Background(), opt.MySQL)
-	case types.BackendTiDB:
-		return metadataMySQL.Open(context.Background(), opt.TiDB)
 	case types.BackendSQLite:
 		return metadataSQLite.Open(context.Background(), opt.SQLite)
 	default:
@@ -73,8 +69,6 @@ func CreateMetadataDatabase(ctx context.Context, opt types.MetadataStoreConfig)
 		return metadataPG.CreateDatabase(ctx, *opt.Postgres)
 	case types.BackendMySQL:
 		return metadataMySQL.CreateDatabase(ctx, *opt.MySQL)
-	case types.BackendTiDB:
-		return metadataMySQL.CreateDatabase(ctx, *opt.TiDB)
 	case types.BackendSQLite:
 		return metadataSQLite.CreateDatabase(ctx, *opt.SQLite)
 	default:
@@ -88,8 +82,6 @@ func OpenOfflineStore(ctx context.Context, opt types.OfflineStoreConfig) (offlin
 		return offlinePG.Open(opt.Postgres)
 	case types.BackendMySQL:
 		return offlineMySQL.Open(opt.MySQL)
-	case types.BackendTiDB:
-		return offlineMySQL.Open(opt.TiDB)
 	case types.BackendSnowflake:
 		return offlineSnowflake.Open(opt.Snowflake)
 	case types.BackendBigQuery:
diff --git a/internal/database/dbutil/schema.go b/internal/database/dbutil/schema.go
index cae4bbe1..82ffe28d 100644
--- a/internal/database/dbutil/schema.go
+++ b/internal/database/dbutil/schema.go
@@ -1,47 +1,12 @@
 package dbutil
 
 import (
-	"bytes"
 	"fmt"
 	"strings"
-	"text/template"
 
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
-const tableSchemaTmpl = `
-CREATE TABLE {{ .TableName }} (
-	{{ entity .Entity .Backend }},
-	{{ fields .Fields .Backend }}
-)`
-
-// TODO: Add back `PRIMARY KEY` back when we have functions for
-// creating cdc table
-var (
-	tableSchemaTmplFuncMap = template.FuncMap{
-		"entity": func(entity types.Entity, backend types.BackendType) string {
-			entityName := QuoteFn(backend)(entity.Name)
-			switch backend {
-			case types.BackendCassandra, types.BackendSQLite:
-				return fmt.Sprintf(`%s TEXT`, entityName)
-			case types.BackendPostgres, types.BackendRedshift, types.BackendSnowflake, types.BackendMySQL:
-				return fmt.Sprintf(`%s VARCHAR(%d)`, entityName, entity.Length)
-			case types.BackendBigQuery:
-				return fmt.Sprintf(`%s STRING`, entityName)
-			default:
-				panic(fmt.Sprintf("unsupported backend type %s", backend))
-			}
-		},
-		"fields": func(columns []Column, backend types.BackendType) string {
-			rs := make([]string, 0, len(columns))
-			for _, column := range columns {
-				rs = append(rs, fmt.Sprintf("%s %s", QuoteFn(backend)(column.Name), column.DbType))
-			}
-			return strings.Join(rs, ",\n\t")
-		},
-	}
-)
-
 type Column struct {
 	Name      string
 	DbType    string
@@ -58,61 +23,96 @@ func (c ColumnList) Names() []string {
 	return names
 }
 
-type TableSchemaTmplOpts struct {
-	TableName string
-	Entity    types.Entity
-	Fields    []Column
-	Backend   types.BackendType
-}
-
 func BuildTableSchema(
 	tableName string,
 	entity *types.Entity,
-	withUnixMillis bool,
+	isCDC bool,
 	features types.FeatureList,
+	pkFields []string,
 	backend types.BackendType,
 ) string {
-	buf := bytes.NewBuffer(nil)
-	opt := tableSchemaTmplOpts(tableName, *entity, withUnixMillis, features, backend)
-	tmpl := template.Must(template.New("schema").Funcs(tableSchemaTmplFuncMap).Parse(tableSchemaTmpl))
-	if err := tmpl.Execute(buf, opt); err != nil {
-		panic(err)
+	columns := parseColumns(entity, isCDC, features, backend)
+	return createTableDDL(tableName, columns, pkFields, backend)
+}
+
+func createTableDDL(tableName string, columns ColumnList, pkFields []string, backend types.BackendType) string {
+	qt := QuoteFn(backend)
+	columnDefs := make([]string, 0, len(columns))
+	for _, column := range columns {
+		columnDefs = append(columnDefs, fmt.Sprintf("\t%s %s", qt(column.Name), column.DbType))
 	}
-	return buf.String()
+	tableDef := strings.Join(columnDefs, ",\n")
+	if len(pkFields) != 0 {
+		switch backend {
+		case types.BackendBigQuery:
+			// big query does not support primary key
+		case types.BackendCassandra,
+			types.BackendPostgres,
+			types.BackendMySQL,
+			types.BackendSQLite,
+			types.BackendSnowflake,
+			types.BackendRedshift:
+			tableDef += fmt.Sprintf(",\n\tPRIMARY KEY (%s)", qt(pkFields...))
+		default:
+			panic(fmt.Sprintf("unsupported backend type %s", backend))
+		}
+	}
+	return fmt.Sprintf("CREATE TABLE %s (\n%s\n)", qt(tableName), tableDef)
 }
 
-func tableSchemaTmplOpts(tableName string,
-	entity types.Entity,
-	withUnixMillis bool,
-	features types.FeatureList,
-	backend types.BackendType,
-) *TableSchemaTmplOpts {
-	fields := make([]Column, 0, len(features))
-	if withUnixMillis {
-		dbType, err := DBValueType(backend, types.Int64)
-		if err != nil {
-			panic(err)
+func BuildIndexDDL(tableName string, indexName string, fields []string, backend types.BackendType) string {
+	qt := QuoteFn(backend)
+	switch backend {
+	case types.BackendCassandra,
+		types.BackendPostgres,
+		types.BackendMySQL,
+		types.BackendSQLite:
+		return fmt.Sprintf("CREATE INDEX %s ON %s (%s)", qt(indexName), qt(tableName), qt(fields...))
+	case types.BackendSnowflake:
+		return fmt.Sprintf("ALTER TABLE %s CLUSTER BY (%s)", qt(tableName), qt(fields...))
+	default:
+		panic(fmt.Sprintf("unsupported backend type %s", backend))
+	}
+}
+
+func parseColumns(entity *types.Entity, isCDC bool, features types.FeatureList, backend types.BackendType) (rs []Column) {
+	// entity column
+	{
+		c := Column{Name: entity.Name, ValueType: types.String}
+		switch backend {
+		case types.BackendCassandra, types.BackendSQLite:
+			c.DbType = "TEXT"
+		case types.BackendPostgres, types.BackendRedshift, types.BackendSnowflake, types.BackendMySQL:
+			c.DbType = fmt.Sprintf(`VARCHAR(%d)`, entity.Length)
+		case types.BackendBigQuery:
+			c.DbType = "STRING"
+		default:
+			panic(fmt.Sprintf("unsupported backend type %s", backend))
 		}
-		fields = append(fields, Column{
-			Name:   "unix_milli",
-			DbType: dbType,
-		})
+		rs = append(rs, c)
 	}
-	for _, feature := range features {
-		dbType, err := DBValueType(backend, feature.ValueType)
-		if err != nil {
-			panic(err)
+
+	// unix_milli column
+	{
+		if isCDC {
+			valueType := types.Int64
+			dbType, err := DBValueType(backend, valueType)
+			if err != nil {
+				panic(err)
+			}
+			rs = append(rs, Column{Name: "unix_milli", DbType: dbType, ValueType: valueType})
 		}
-		fields = append(fields, Column{
-			Name:   feature.Name,
-			DbType: dbType,
-		})
 	}
 
-	return &TableSchemaTmplOpts{
-		TableName: QuoteFn(backend)(tableName),
-		Entity:    entity,
-		Fields:    fields,
-		Backend:   backend,
+	// feature columns
+	{
+		for _, f := range features {
+			dbType, err := DBValueType(backend, f.ValueType)
+			if err != nil {
+				panic(err)
+			}
+			rs = append(rs, Column{Name: f.Name, DbType: dbType, ValueType: f.ValueType})
+		}
 	}
+	return
 }
diff --git a/internal/database/offline/bigquery/store.go b/internal/database/offline/bigquery/store.go
index a02a7314..d0d59f81 100644
--- a/internal/database/offline/bigquery/store.go
+++ b/internal/database/offline/bigquery/store.go
@@ -83,7 +83,7 @@ func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
 }
 
 func (db *DB) CreateTable(ctx context.Context, opt offline.CreateTableOpt) error {
-	schema := dbutil.BuildTableSchema(opt.TableName, opt.Entity, opt.WithUnixMillis, opt.Features, Backend)
+	schema := dbutil.BuildTableSchema(opt.TableName, opt.Entity, opt.IsCDC, opt.Features, nil, Backend)
 	if _, err := db.Client.Query(schema).Run(ctx); err != nil {
 		return err
 	}
diff --git a/internal/database/offline/sqlutil/create_table.go b/internal/database/offline/sqlutil/create_table.go
index 0edc2b79..e472a54e 100644
--- a/internal/database/offline/sqlutil/create_table.go
+++ b/internal/database/offline/sqlutil/create_table.go
@@ -10,9 +10,24 @@ import (
 )
 
 func CreateTable(ctx context.Context, db *sqlx.DB, opt offline.CreateTableOpt, backend types.BackendType) error {
-	schema := dbutil.BuildTableSchema(opt.TableName, opt.Entity, opt.WithUnixMillis, opt.Features, backend)
-	if _, err := db.ExecContext(ctx, schema); err != nil {
-		return err
+	if opt.IsCDC {
+		// Create index (entity_key, unix_milli) on cdc table
+		schema := dbutil.BuildTableSchema(opt.TableName, opt.Entity, false, opt.Features, nil, backend)
+		if _, err := db.ExecContext(ctx, schema); err != nil {
+			return err
+		}
+		indexFields := []string{opt.Entity.Name, "unix_milli"}
+		indexDDL := dbutil.BuildIndexDDL(opt.TableName, "idx", indexFields, backend)
+		if _, err := db.ExecContext(ctx, indexDDL); err != nil {
+			return err
+		}
+	} else {
+		// Create primary key (entity_key) on snapshot table
+		pkFields := []string{opt.Entity.Name}
+		schema := dbutil.BuildTableSchema(opt.TableName, opt.Entity, false, opt.Features, pkFields, backend)
+		if _, err := db.ExecContext(ctx, schema); err != nil {
+			return err
+		}
 	}
 	return nil
 }
diff --git a/internal/database/offline/sqlutil/import.go b/internal/database/offline/sqlutil/import.go
index abd7f9e3..99761f53 100644
--- a/internal/database/offline/sqlutil/import.go
+++ b/internal/database/offline/sqlutil/import.go
@@ -16,7 +16,11 @@ func Import(ctx context.Context, db *sqlx.DB, opt offline.ImportOpt, loadData Lo
 	var revision int64
 	err := dbutil.WithTransaction(db, ctx, func(ctx context.Context, tx *sqlx.Tx) error {
 		// create the data table
-		schema := dbutil.BuildTableSchema(opt.SnapshotTableName, opt.Entity, false, opt.Features, backendType)
+		pkFields := []string{opt.Entity.Name}
+		if opt.NoPK {
+			pkFields = nil
+		}
+		schema := dbutil.BuildTableSchema(opt.SnapshotTableName, opt.Entity, false, opt.Features, pkFields, backendType)
 		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return err
diff --git a/internal/database/offline/sqlutil/join.go b/internal/database/offline/sqlutil/join.go
index 9500eba0..5ebaa666 100644
--- a/internal/database/offline/sqlutil/join.go
+++ b/internal/database/offline/sqlutil/join.go
@@ -44,18 +44,13 @@ func Join(ctx context.Context, db *sqlx.DB, opt offline.JoinOpt, backendType typ
 }
 
 func DoJoin(ctx context.Context, dbOpt dbutil.DBOpt, opt DoJoinOpt) (*types.JoinResult, error) {
-	data := make(chan []interface{})
-	defer close(data)
-	emptyResult := &types.JoinResult{
-		Data: data,
-	}
 	// Step 1: prepare temporary table entity_rows
 	features := types.FeatureList{}
 	for _, featureList := range opt.FeatureMap {
 		features = append(features, featureList...)
 	}
 	if len(features) == 0 {
-		return emptyResult, nil
+		return nil, nil
 	}
 	entityRowsTableName, err := PrepareEntityRowsTable(ctx, dbOpt, opt.Entity, opt.EntityRows, opt.ValueNames)
 	if err != nil {
@@ -174,15 +169,10 @@ func JoinOneGroup(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.JoinOneGr
 }
 
 func ReadJoinedTable(ctx context.Context, dbOpt dbutil.DBOpt, opt ReadJoinedTableOpt) (*types.JoinResult, error) {
-	data := make(chan []interface{})
-	defer close(data)
-	emptyResult := &types.JoinResult{
-		Data: data,
-	}
 	tableNames := opt.TableNames
 	dropTableNames := append([]string{opt.EntityRowsTableName}, opt.AllTableNames...)
 	if len(tableNames) == 0 {
-		return emptyResult, nil
+		return nil, nil
 	}
 	qt := dbutil.QuoteFn(dbOpt.Backend)
 
diff --git a/internal/database/offline/sqlutil/snapshot.go b/internal/database/offline/sqlutil/snapshot.go
index 728f447c..8be3fd22 100644
--- a/internal/database/offline/sqlutil/snapshot.go
+++ b/internal/database/offline/sqlutil/snapshot.go
@@ -69,7 +69,7 @@ func Snapshot(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.SnapshotOpt)
 		currCdcTableName = fmt.Sprintf("%s.%s", *dbOpt.DatasetID, currCdcTableName)
 	}
 
-	schema := dbutil.BuildTableSchema(currSnapshotTableName, opt.Group.Entity, false, opt.Features, dbOpt.Backend)
+	schema := dbutil.BuildTableSchema(currSnapshotTableName, opt.Group.Entity, false, opt.Features, []string{opt.Group.Entity.Name}, dbOpt.Backend)
 	if err := dbOpt.ExecContext(ctx, schema, nil); err != nil {
 		return err
 	}
diff --git a/internal/database/offline/types.go b/internal/database/offline/types.go
index c0882aef..6f2f4966 100644
--- a/internal/database/offline/types.go
+++ b/internal/database/offline/types.go
@@ -45,6 +45,7 @@ type ImportOpt struct {
 	Revision          *int64
 	SnapshotTableName string
 	Source            *CSVSource
+	NoPK              bool // TODO: to import cdc data temporarily for testing
 }
 
 type CSVSource struct {
@@ -60,8 +61,8 @@ type SnapshotOpt struct {
 }
 
 type CreateTableOpt struct {
-	TableName      string
-	Entity         *types.Entity
-	Features       types.FeatureList
-	WithUnixMillis bool
+	TableName string
+	Entity    *types.Entity
+	Features  types.FeatureList
+	IsCDC     bool
 }
diff --git a/internal/database/online/cassandra/get.go b/internal/database/online/cassandra/get.go
index 6bcedd21..58da07ae 100644
--- a/internal/database/online/cassandra/get.go
+++ b/internal/database/online/cassandra/get.go
@@ -25,7 +25,7 @@ func (db *DB) Get(ctx context.Context, opt online.GetOpt) (dbutil.RowMap, error)
 	rs := make(map[string]interface{})
 	if err := db.Query(query, opt.EntityKey).WithContext(ctx).MapScan(rs); err != nil {
 		if err == gocql.ErrNotFound || isTableNotFoundError(err, tableName) {
-			return rs, nil
+			return nil, nil
 		}
 		return nil, err
 	}
@@ -57,7 +57,7 @@ func (db *DB) MultiGet(ctx context.Context, opt online.MultiGetOpt) (map[string]
 		Iter().SliceMap()
 	if err != nil {
 		if err == gocql.ErrNotFound || isTableNotFoundError(err, tableName) {
-			return rs, nil
+			return nil, nil
 		}
 		return nil, err
 	}
diff --git a/internal/database/online/cassandra/import.go b/internal/database/online/cassandra/import.go
index 3a1fe5a5..1f20861e 100644
--- a/internal/database/online/cassandra/import.go
+++ b/internal/database/online/cassandra/import.go
@@ -16,7 +16,7 @@ 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, Backend)
+	table := dbutil.BuildTableSchema(tableName, opt.Entity, false, opt.Features, []string{opt.Entity.Name}, Backend)
 
 	// create table
 	if err := db.Query(table).Exec(); err != nil {
diff --git a/internal/database/online/dynamodb/get.go b/internal/database/online/dynamodb/get.go
index 3ed91fc4..c2e06059 100644
--- a/internal/database/online/dynamodb/get.go
+++ b/internal/database/online/dynamodb/get.go
@@ -33,7 +33,7 @@ func (db *DB) Get(ctx context.Context, opt online.GetOpt) (dbutil.RowMap, error)
 	})
 	if err != nil {
 		if apiErr := new(types.ResourceNotFoundException); errors.As(err, &apiErr) {
-			return make(dbutil.RowMap), nil
+			return nil, nil
 		}
 		return nil, err
 	}
@@ -100,7 +100,7 @@ func batchGetItem(ctx context.Context, db *DB, keys []map[string]types.Attribute
 
 func deserializeFeatureValues(features oomTypes.FeatureList, item map[string]types.AttributeValue) (dbutil.RowMap, error) {
 	if item == nil {
-		return make(dbutil.RowMap), nil
+		return nil, nil
 	}
 	rowMap := make(dbutil.RowMap)
 	var value interface{}
diff --git a/internal/database/online/dynamodb/store.go b/internal/database/online/dynamodb/store.go
index 1d9e0148..2972e4f0 100644
--- a/internal/database/online/dynamodb/store.go
+++ b/internal/database/online/dynamodb/store.go
@@ -53,3 +53,7 @@ func (db *DB) Ping(ctx context.Context) error {
 func (db *DB) Close() error {
 	return nil
 }
+
+func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
+	panic("Implement me!")
+}
diff --git a/internal/database/online/dynamodb/stream.go b/internal/database/online/dynamodb/stream.go
index f8bb7c9a..234ce3bd 100644
--- a/internal/database/online/dynamodb/stream.go
+++ b/internal/database/online/dynamodb/stream.go
@@ -4,7 +4,6 @@ import (
 	"context"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
-	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/attributevalue"
 	"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"
@@ -40,34 +39,3 @@ func (db *DB) PrepareStreamTable(ctx context.Context, opt online.PrepareStreamTa
 	})
 	return err
 }
-
-func (db *DB) Push(ctx context.Context, opt online.PushOpt) error {
-	var (
-		tableName = sqlutil.OnlineStreamTableName(opt.GroupID)
-		item      = make(map[string]types.AttributeValue)
-	)
-
-	entityKeyValue, err := attributevalue.Marshal(opt.EntityKey)
-	if err != nil {
-		return err
-	}
-	item[opt.Entity.Name] = entityKeyValue
-
-	for i, feature := range opt.Features {
-		value, err := serializeByTag(opt.FeatureValues[i], feature.ValueType)
-		if err != nil {
-			return err
-		}
-		attributevalue, err := attributevalue.Marshal(value)
-		if err != nil {
-			return err
-		}
-		item[feature.Name] = attributevalue
-	}
-
-	_, err = db.Client.PutItem(ctx, &dynamodb.PutItemInput{
-		TableName: &tableName,
-		Item:      item,
-	})
-	return err
-}
diff --git a/internal/database/online/sqlutil/get.go b/internal/database/online/sqlutil/get.go
index b8344552..009bdff9 100644
--- a/internal/database/online/sqlutil/get.go
+++ b/internal/database/online/sqlutil/get.go
@@ -22,7 +22,7 @@ func Get(ctx context.Context, db *sqlx.DB, opt online.GetOpt, backend types.Back
 	record, err := db.QueryRowxContext(ctx, db.Rebind(query), opt.EntityKey).SliceScan()
 	if err != nil {
 		if err == sql.ErrNoRows || dbutil.IsTableNotFoundError(err, backend) {
-			return make(dbutil.RowMap), nil
+			return nil, nil
 		}
 		return nil, err
 	}
diff --git a/internal/database/online/sqlutil/import.go b/internal/database/online/sqlutil/import.go
index 80075eac..ecbd3b7a 100644
--- a/internal/database/online/sqlutil/import.go
+++ b/internal/database/online/sqlutil/import.go
@@ -17,7 +17,7 @@ func Import(ctx context.Context, db *sqlx.DB, opt online.ImportOpt, backend type
 	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, backend)
+		schema := dbutil.BuildTableSchema(tableName, opt.Entity, false, opt.Features, []string{opt.Entity.Name}, backend)
 		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return err
diff --git a/internal/database/online/tikv/get.go b/internal/database/online/tikv/get.go
index 131767ba..45343a3c 100644
--- a/internal/database/online/tikv/get.go
+++ b/internal/database/online/tikv/get.go
@@ -26,7 +26,7 @@ func (db *DB) Get(ctx context.Context, opt online.GetOpt) (dbutil.RowMap, error)
 		return nil, err
 	}
 	if rowMap, ok := res[opt.EntityKey]; !ok {
-		return make(dbutil.RowMap), nil
+		return nil, nil
 	} else {
 		return rowMap, nil
 	}
diff --git a/internal/database/online/tikv/purge.go b/internal/database/online/tikv/purge.go
index 43d9768b..ec644a15 100644
--- a/internal/database/online/tikv/purge.go
+++ b/internal/database/online/tikv/purge.go
@@ -7,12 +7,13 @@ import (
 )
 
 func (db *DB) Purge(ctx context.Context, revisionID int) error {
-	serializedRevisionID, err := kvutil.SerializeByValue(revisionID)
+	startKey, err := kvutil.SerializeByValue(revisionID)
 	if err != nil {
 		return err
 	}
-	startKey := append([]byte(kvutil.KeyPrefixForBatchFeature+serializedRevisionID), byte(keyDelimiter))
-	endKey := append([]byte(kvutil.KeyPrefixForBatchFeature+serializedRevisionID), byte(keyDelimiter+1))
-
-	return db.DeleteRange(ctx, startKey, endKey)
+	endKey, err := kvutil.SerializeByValue(revisionID + 1)
+	if err != nil {
+		return err
+	}
+	return db.DeleteRange(ctx, []byte(kvutil.KeyPrefixForBatchFeature+startKey), []byte(kvutil.KeyPrefixForBatchFeature+endKey))
 }
diff --git a/internal/database/online/tikv/util.go b/internal/database/online/tikv/util.go
index c04eb3a0..33b2326d 100644
--- a/internal/database/online/tikv/util.go
+++ b/internal/database/online/tikv/util.go
@@ -6,7 +6,7 @@ import (
 	"github.com/oom-ai/oomstore/internal/database/online/kvutil"
 )
 
-const keyDelimiter = ':'
+const keyDelimiter = ":"
 
 func getKeyOfBatchFeature(revisionID string, entityKey string, featureID string) []byte {
 	return []byte(kvutil.KeyPrefixForBatchFeature + getKey(revisionID, entityKey, featureID))
@@ -17,5 +17,5 @@ func getKeyOfStreamFeature(groupID string, entityKey string, featureID string) [
 }
 
 func getKey(revisionID string, entityKey string, featureID string) string {
-	return strings.Join([]string{revisionID, entityKey, featureID}, string(keyDelimiter))
+	return strings.Join([]string{revisionID, entityKey, featureID}, keyDelimiter)
 }
diff --git a/pkg/oomstore/join.go b/pkg/oomstore/join.go
index 718ad838..b804d677 100644
--- a/pkg/oomstore/join.go
+++ b/pkg/oomstore/join.go
@@ -19,11 +19,6 @@ import (
 // ChannelJoin gets point-in-time correct feature values for each entity row.
 // Currently, this API only supports batch features.
 func (s *OomStore) ChannelJoin(ctx context.Context, opt types.ChannelJoinOpt) (*types.JoinResult, error) {
-	data := make(chan []interface{})
-	defer close(data)
-	emptyResult := &types.JoinResult{
-		Data: data,
-	}
 	features, err := s.metadata.ListFeature(ctx, metadata.ListFeatureOpt{
 		FeatureFullNames: &opt.FeatureFullNames,
 	})
@@ -35,7 +30,7 @@ func (s *OomStore) ChannelJoin(ctx context.Context, opt types.ChannelJoinOpt) (*
 		return f.Group.Category == types.CategoryBatch
 	})
 	if len(features) == 0 {
-		return emptyResult, nil
+		return nil, nil
 	}
 
 	entity, err := getSharedEntity(features)
diff --git a/pkg/oomstore/online_query.go b/pkg/oomstore/online_query.go
index 62081524..354efd2a 100644
--- a/pkg/oomstore/online_query.go
+++ b/pkg/oomstore/online_query.go
@@ -10,27 +10,28 @@ import (
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
-// OnlineGet gets online features of a particular entity instance.
+// Get online features of a particular entity instance.
 func (s *OomStore) OnlineGet(ctx context.Context, opt types.OnlineGetOpt) (*types.FeatureValues, error) {
-	rs := types.FeatureValues{
-		EntityKey:        opt.EntityKey,
-		FeatureFullNames: opt.FeatureFullNames,
-		FeatureValueMap:  make(map[string]interface{}),
-	}
 	features := s.metadata.CacheListFeature(ctx, metadata.ListFeatureOpt{
 		FeatureFullNames: &opt.FeatureFullNames,
 	}).Filter(func(f *types.Feature) bool {
 		return f.Group.OnlineRevisionID != nil
 	})
 	if len(features) == 0 {
-		return &rs, nil
+		return nil, nil
 	}
 
 	entity, err := getSharedEntity(features)
 	if err != nil {
 		return nil, err
 	}
-	rs.EntityName = entity.Name
+
+	rs := types.FeatureValues{
+		EntityName:       entity.Name,
+		EntityKey:        opt.EntityKey,
+		FeatureFullNames: opt.FeatureFullNames,
+		FeatureValueMap:  make(map[string]interface{}),
+	}
 
 	featureMap := groupFeaturesByRevisionID(features)
 
@@ -54,9 +55,8 @@ func (s *OomStore) OnlineGet(ctx context.Context, opt types.OnlineGetOpt) (*type
 	return &rs, nil
 }
 
-// OnlineMultiGet gets online features of multiple entity instances.
+// Get online features of multiple entity instances.
 func (s *OomStore) OnlineMultiGet(ctx context.Context, opt types.OnlineMultiGetOpt) (map[string]*types.FeatureValues, error) {
-	result := make(map[string]*types.FeatureValues)
 	features := s.metadata.CacheListFeature(ctx, metadata.ListFeatureOpt{
 		FeatureFullNames: &opt.FeatureFullNames,
 	})
@@ -65,7 +65,7 @@ func (s *OomStore) OnlineMultiGet(ctx context.Context, opt types.OnlineMultiGetO
 		return f.OnlineRevisionID() != nil
 	})
 	if len(features) == 0 {
-		return result, nil
+		return nil, nil
 	}
 
 	entity, err := getSharedEntity(features)
@@ -83,6 +83,7 @@ func (s *OomStore) OnlineMultiGet(ctx context.Context, opt types.OnlineMultiGetO
 		return nil, err
 	}
 
+	result := make(map[string]*types.FeatureValues)
 	for _, entityKey := range opt.EntityKeys {
 		result[entityKey] = &types.FeatureValues{
 			EntityName:       entity.Name,
diff --git a/pkg/oomstore/types/config.go b/pkg/oomstore/types/config.go
index 016c7898..67f8059f 100644
--- a/pkg/oomstore/types/config.go
+++ b/pkg/oomstore/types/config.go
@@ -17,7 +17,6 @@ const (
 	BackendCassandra BackendType = "cassandra"
 	BackendBigQuery  BackendType = "bigquery"
 	BackendRedshift  BackendType = "redshift"
-	BackendTiDB      BackendType = "tidb"
 	BackendTiKV      BackendType = "tikv"
 )
 
@@ -149,7 +148,7 @@ func (cfg *MetadataStoreConfig) Validate() error {
 		n++
 	}
 	if cfg.TiDB != nil {
-		cfg.Backend = BackendTiDB
+		cfg.Backend = BackendMySQL
 		n++
 	}
 	if n != 1 {
@@ -185,7 +184,7 @@ func (cfg *OnlineStoreConfig) Validate() error {
 		n++
 	}
 	if cfg.TiDB != nil {
-		cfg.Backend = BackendTiDB
+		cfg.Backend = BackendMySQL
 		n++
 	}
 	if cfg.TiKV != nil {
@@ -225,7 +224,7 @@ func (cfg *OfflineStoreConfig) Validate() error {
 		n++
 	}
 	if cfg.TiDB != nil {
-		cfg.Backend = BackendTiDB
+		cfg.Backend = BackendMySQL
 		n++
 	}
 	if n != 1 {


Test patch:
diff --git a/internal/database/dbutil/schema_test.go b/internal/database/dbutil/schema_test.go
index 62813484..37846274 100644
--- a/internal/database/dbutil/schema_test.go
+++ b/internal/database/dbutil/schema_test.go
@@ -1,6 +1,7 @@
 package dbutil_test
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
@@ -86,8 +87,8 @@ CREATE TABLE "user" (
 
 	for _, c := range cases {
 		t.Run(c.description, func(t *testing.T) {
-			schema := dbutil.BuildTableSchema(c.tableName, c.entity, false, c.features, c.backend)
-			require.Equal(t, c.want, schema)
+			schema := dbutil.BuildTableSchema(c.tableName, c.entity, false, c.features, nil, c.backend)
+			require.Equal(t, strings.TrimSpace(c.want), schema)
 		})
 	}
 }
diff --git a/internal/database/offline/test_impl/create_table.go b/internal/database/offline/test_impl/create_table.go
index 7efa5bec..33573366 100644
--- a/internal/database/offline/test_impl/create_table.go
+++ b/internal/database/offline/test_impl/create_table.go
@@ -36,21 +36,21 @@ func TestCreateTable(t *testing.T, prepareStore PrepareStoreFn, destroyStore Des
 		expectedError error
 	}{
 		{
-			describtion: "with unix milli",
+			describtion: "cdc table (with unix_milli)",
 			opt: offline.CreateTableOpt{
-				TableName:      tableName,
-				Entity:         &entity,
-				Features:       features,
-				WithUnixMillis: true,
+				TableName: tableName,
+				Entity:    &entity,
+				Features:  features,
+				IsCDC:     true,
 			},
 		},
 		{
-			describtion: "without milli",
+			describtion: "snapshot table (without_milli)",
 			opt: offline.CreateTableOpt{
-				TableName:      tableName,
-				Entity:         &entity,
-				Features:       features,
-				WithUnixMillis: false,
+				TableName: tableName,
+				Entity:    &entity,
+				Features:  features,
+				IsCDC:     false,
 			},
 		},
 	}
diff --git a/internal/database/offline/test_impl/join.go b/internal/database/offline/test_impl/join.go
index 9282cfd2..b1203465 100644
--- a/internal/database/offline/test_impl/join.go
+++ b/internal/database/offline/test_impl/join.go
@@ -76,7 +76,7 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroySto
 			opt: offline.JoinOpt{
 				FeatureMap: make(map[string]types.FeatureList),
 			},
-			expected: prepareEmptyResult(),
+			expected: nil,
 		},
 		{
 			description: "no entity rows",
@@ -85,7 +85,7 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroySto
 				FeatureMap: oneGroupFeatureMap,
 				EntityRows: prepareEntityRows(true, false),
 			},
-			expected: prepareEmptyResult(),
+			expected: nil,
 		},
 		{
 			description: "one batch feature group",
@@ -135,12 +135,15 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroySto
 		t.Run(tc.description, func(t *testing.T) {
 			actual, err := store.Join(context.Background(), tc.opt)
 			require.NoError(t, err)
-
-			assert.ElementsMatch(t, tc.expected.Header, actual.Header)
-			expectedValues := extractValues(tc.expected.Data)
-			actualValues := extractValues(actual.Data)
-			for i := range expectedValues {
-				assert.ElementsMatch(t, expectedValues[i], actualValues[i])
+			if tc.expected == nil {
+				assert.Equal(t, tc.expected, actual)
+			} else {
+				assert.ElementsMatch(t, tc.expected.Header, actual.Header)
+				expectedValues := extractValues(tc.expected.Data)
+				actualValues := extractValues(actual.Data)
+				for i := range expectedValues {
+					assert.ElementsMatch(t, expectedValues[i], actualValues[i])
+				}
 			}
 		})
 	}
@@ -293,13 +296,6 @@ func prepareResult(oneGroup bool, withValue bool, values []map[string]interface{
 	}
 }
 
-func prepareEmptyResult() *types.JoinResult {
-	data := make(chan []interface{})
-	defer close(data)
-	return &types.JoinResult{
-		Data: data,
-	}
-}
 func prepareStreamingResultValues() []map[string]interface{} {
 	return []map[string]interface{}{
 		{
diff --git a/internal/database/offline/test_impl/util.go b/internal/database/offline/test_impl/util.go
index bc468741..8a18f29f 100644
--- a/internal/database/offline/test_impl/util.go
+++ b/internal/database/offline/test_impl/util.go
@@ -37,6 +37,7 @@ func buildTestSnapshotTable(
 		Header:            header,
 		Source:            source,
 		Revision:          &revision,
+		NoPK:              true,
 	}
 	_, err := store.Import(ctx, opt)
 	require.NoError(t, err)
diff --git a/internal/database/online/dynamodb/store_test.go b/internal/database/online/dynamodb/store_test.go
index f70a2ee5..bf4a46fc 100644
--- a/internal/database/online/dynamodb/store_test.go
+++ b/internal/database/online/dynamodb/store_test.go
@@ -68,7 +68,3 @@ func TestPurgeNotRemovesOtherRevisions(t *testing.T) {
 func TestPrepareStreamTable(t *testing.T) {
 	test_impl.TestPrepareStreamTable(t, prepareStore, destroyStore(t))
 }
-
-func TestPush(t *testing.T) {
-	test_impl.TestPush(t, prepareStore, destroyStore(t))
-}
diff --git a/pkg/oomstore/join_test.go b/pkg/oomstore/join_test.go
index 51b6b2bb..bfb88a94 100644
--- a/pkg/oomstore/join_test.go
+++ b/pkg/oomstore/join_test.go
@@ -40,6 +40,7 @@ func TestChannelJoin(t *testing.T) {
 		},
 	}
 
+	var emptyResult *types.JoinResult
 	validResult := prepareResult()
 	entityRows := make(chan types.EntityRow)
 
@@ -61,7 +62,7 @@ func TestChannelJoin(t *testing.T) {
 			},
 			features:      streamFeatures,
 			expectedError: nil,
-			expected:      prepareEmptyResult(),
+			expected:      emptyResult,
 		},
 		{
 			description: "inconsistent features, return nil",
@@ -85,9 +86,9 @@ func TestChannelJoin(t *testing.T) {
 				"device_basic":    {consistentFeatures[0]},
 				"device_advanced": {consistentFeatures[1]},
 			},
-			joined:        prepareEmptyResult(),
+			joined:        nil,
 			expectedError: nil,
-			expected:      prepareEmptyResult(),
+			expected:      emptyResult,
 		},
 		{
 			description: "consistent entity type, succeed",
@@ -121,6 +122,8 @@ func TestChannelJoin(t *testing.T) {
 			actual, err := store.ChannelJoin(context.Background(), tc.opt)
 			if tc.expectedError != nil {
 				assert.Error(t, err, tc.expectedError)
+			} else if tc.expected == nil {
+				assert.Equal(t, tc.expected, actual)
 			} else {
 				assert.NoError(t, err)
 				assert.Equal(t, tc.expected.Header, actual.Header)
@@ -148,19 +151,8 @@ func prepareResult() *types.JoinResult {
 
 func extractValues(stream <-chan []interface{}) [][]interface{} {
 	values := make([][]interface{}, 0)
-	if stream == nil {
-		return values
-	}
 	for item := range stream {
 		values = append(values, item)
 	}
 	return values
 }
-
-func prepareEmptyResult() *types.JoinResult {
-	data := make(chan []interface{})
-	defer close(data)
-	return &types.JoinResult{
-		Data: data,
-	}
-}
diff --git a/pkg/oomstore/online_query_test.go b/pkg/oomstore/online_query_test.go
index 180c9b8f..21a12332 100644
--- a/pkg/oomstore/online_query_test.go
+++ b/pkg/oomstore/online_query_test.go
@@ -43,11 +43,7 @@ func TestOnlineGet(t *testing.T) {
 			},
 			features:      unavailableFeatures,
 			expectedError: nil,
-			expected: &types.FeatureValues{
-				EntityKey:        "1234",
-				FeatureFullNames: consistentFeatures.FullNames(),
-				FeatureValueMap:  map[string]interface{}{},
-			},
+			expected:      nil,
 		},
 		{
 			description: "inconsistent entity type, fail",
@@ -131,7 +127,7 @@ func TestOnlineMultiGet(t *testing.T) {
 			},
 			features:      unavailableFeatures,
 			expectedError: nil,
-			expected:      map[string]*types.FeatureValues{},
+			expected:      nil,
 		},
 		{
 			description: "inconsistent entity type, fail",