Repository: oom-ai/oomstore
Pull request #890 (2022-01-04 06:02:18)

fix: open tidb

codecov[bot]: # [Codecov](https://codecov.io/gh/oom-ai/oomstore/pull/890?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai) Report
> Merging [#890](https://codecov.io/gh/oom-ai/oomstore/pull/890?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai) (35d401e) into [main](https://codecov.io/gh/oom-ai/oomstore/commit/77cd80d135c3df8456fa4a2059c9ae64c52cc5e2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai) (77cd80d) will **increase** coverage by `0.07%`.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/oom-ai/oomstore/pull/890/graphs/tree.svg?width=650&height=150&src=pr&token=C59L7LTRM4&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai)](https://codecov.io/gh/oom-ai/oomstore/pull/890?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai)

```diff
@@            Coverage Diff             @@
##             main     #890      +/-   ##
==========================================
+ Coverage   51.28%   51.35%   +0.07%     
==========================================
  Files          84       84              
  Lines        3436     3441       +5     
==========================================
+ Hits         1762     1767       +5     
  Misses       1461     1461              
  Partials      213      213              
```


| [Impacted Files](https://codecov.io/gh/oom-ai/oomstore/pull/890?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai) | Coverage Δ | |
|---|---|---|
| [pkg/oomstore/online\_query.go](https://codecov.io/gh/oom-ai/oomstore/pull/890/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai#diff-cGtnL29vbXN0b3JlL29ubGluZV9xdWVyeS5nbw==) | `84.03% <0.00%> (ø)` | |
| [pkg/oomstore/join.go](https://codecov.io/gh/oom-ai/oomstore/pull/890/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai#diff-cGtnL29vbXN0b3JlL2pvaW4uZ28=) | `38.21% <0.00%> (+2.03%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/oom-ai/oomstore/pull/890?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/oom-ai/oomstore/pull/890?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai). Last update [77cd80d...35d401e](https://codecov.io/gh/oom-ai/oomstore/pull/890?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=oom-ai).

YikSanChan: @wfxr because tidb and mysql share the same backend type. Having a separate backend type will add quite some code.
wfxr: @YikSanChan I mean just add a dedicate backend enum:
```
	switch opt.Backend {
	case types.BackendPostgres:
		return onlinePG.Open(opt.Postgres)
	case types.BackendRedis:
		return onlineRedis.Open(opt.Redis), nil
	case types.BackendMySQL, types.BackendTiDB:
		return onlineMySQL.Open(opt.TiDB)
	case types.BackendDynamoDB:
		return onlineDynamoDB.Open(opt.DynamoDB)
	case types.BackendCassandra:
		return onlineCassandra.Open(opt.Cassandra)
	case types.BackendSQLite:
		return onlineSQLite.Open(opt.SQLite)
	case types.BackendTiKV:
		return onlineTiKV.Open(opt.TiKV)
	default:
		return nil, fmt.Errorf("unsupported backend: %s", opt.Backend)
	}
```

Patch:
diff --git a/internal/database/database.go b/internal/database/database.go
index 0a9db934..30266104 100644
--- a/internal/database/database.go
+++ b/internal/database/database.go
@@ -37,6 +37,8 @@ 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:
@@ -56,6 +58,8 @@ 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:
@@ -69,6 +73,8 @@ 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:
@@ -82,6 +88,8 @@ 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/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/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/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/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/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/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 67f8059f..016c7898 100644
--- a/pkg/oomstore/types/config.go
+++ b/pkg/oomstore/types/config.go
@@ -17,6 +17,7 @@ const (
 	BackendCassandra BackendType = "cassandra"
 	BackendBigQuery  BackendType = "bigquery"
 	BackendRedshift  BackendType = "redshift"
+	BackendTiDB      BackendType = "tidb"
 	BackendTiKV      BackendType = "tikv"
 )
 
@@ -148,7 +149,7 @@ func (cfg *MetadataStoreConfig) Validate() error {
 		n++
 	}
 	if cfg.TiDB != nil {
-		cfg.Backend = BackendMySQL
+		cfg.Backend = BackendTiDB
 		n++
 	}
 	if n != 1 {
@@ -184,7 +185,7 @@ func (cfg *OnlineStoreConfig) Validate() error {
 		n++
 	}
 	if cfg.TiDB != nil {
-		cfg.Backend = BackendMySQL
+		cfg.Backend = BackendTiDB
 		n++
 	}
 	if cfg.TiKV != nil {
@@ -224,7 +225,7 @@ func (cfg *OfflineStoreConfig) Validate() error {
 		n++
 	}
 	if cfg.TiDB != nil {
-		cfg.Backend = BackendMySQL
+		cfg.Backend = BackendTiDB
 		n++
 	}
 	if n != 1 {


Test patch:
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/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",