Repository: oom-ai/oomstore
Pull request #880 (2021-12-31 06:28:21)

feat: add offline store CreateTable api


Patch:
diff --git a/internal/database/offline/bigquery/store.go b/internal/database/offline/bigquery/store.go
index d4939576..a02a7314 100644
--- a/internal/database/offline/bigquery/store.go
+++ b/internal/database/offline/bigquery/store.go
@@ -81,3 +81,11 @@ func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
 	}
 	return sqlutil.Snapshot(ctx, dbOpt, opt)
 }
+
+func (db *DB) CreateTable(ctx context.Context, opt offline.CreateTableOpt) error {
+	schema := dbutil.BuildTableSchema(opt.TableName, opt.Entity, opt.WithUnixMillis, opt.Features, Backend)
+	if _, err := db.Client.Query(schema).Run(ctx); err != nil {
+		return err
+	}
+	return nil
+}
diff --git a/internal/database/offline/mock_offline/store.go b/internal/database/offline/mock_offline/store.go
index 4ffdf912..ebac294d 100644
--- a/internal/database/offline/mock_offline/store.go
+++ b/internal/database/offline/mock_offline/store.go
@@ -50,6 +50,20 @@ func (mr *MockStoreMockRecorder) Close() *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockStore)(nil).Close))
 }
 
+// CreateTable mocks base method.
+func (m *MockStore) CreateTable(ctx context.Context, opt offline.CreateTableOpt) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CreateTable", ctx, opt)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// CreateTable indicates an expected call of CreateTable.
+func (mr *MockStoreMockRecorder) CreateTable(ctx, opt interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateTable", reflect.TypeOf((*MockStore)(nil).CreateTable), ctx, opt)
+}
+
 // Export mocks base method.
 func (m *MockStore) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
 	m.ctrl.T.Helper()
diff --git a/internal/database/offline/mysql/store.go b/internal/database/offline/mysql/store.go
index 58509f09..2f50dd2f 100644
--- a/internal/database/offline/mysql/store.go
+++ b/internal/database/offline/mysql/store.go
@@ -54,3 +54,7 @@ func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
 	dbOpt := dbutil.DBOpt{Backend: Backend, SqlxDB: db.DB}
 	return sqlutil.Snapshot(ctx, dbOpt, opt)
 }
+
+func (db *DB) CreateTable(ctx context.Context, opt offline.CreateTableOpt) error {
+	return sqlutil.CreateTable(ctx, db.DB, opt, Backend)
+}
diff --git a/internal/database/offline/postgres/store.go b/internal/database/offline/postgres/store.go
index 79c98e89..06a47f70 100644
--- a/internal/database/offline/postgres/store.go
+++ b/internal/database/offline/postgres/store.go
@@ -53,3 +53,7 @@ func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
 	dbOpt := dbutil.DBOpt{Backend: Backend, SqlxDB: db.DB}
 	return sqlutil.Snapshot(ctx, dbOpt, opt)
 }
+
+func (db *DB) CreateTable(ctx context.Context, opt offline.CreateTableOpt) error {
+	return sqlutil.CreateTable(ctx, db.DB, opt, Backend)
+}
diff --git a/internal/database/offline/redshift/store.go b/internal/database/offline/redshift/store.go
index 233784e0..4b959ceb 100644
--- a/internal/database/offline/redshift/store.go
+++ b/internal/database/offline/redshift/store.go
@@ -54,3 +54,7 @@ func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
 	dbOpt := dbutil.DBOpt{Backend: Backend, SqlxDB: db.DB}
 	return sqlutil.Snapshot(ctx, dbOpt, opt)
 }
+
+func (db *DB) CreateTable(ctx context.Context, opt offline.CreateTableOpt) error {
+	return sqlutil.CreateTable(ctx, db.DB, opt, Backend)
+}
diff --git a/internal/database/offline/snowflake/store.go b/internal/database/offline/snowflake/store.go
index 0cdce748..e3e9fbd6 100644
--- a/internal/database/offline/snowflake/store.go
+++ b/internal/database/offline/snowflake/store.go
@@ -69,3 +69,7 @@ func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
 	dbOpt := dbutil.DBOpt{Backend: Backend, SqlxDB: db.DB}
 	return sqlutil.Snapshot(ctx, dbOpt, opt)
 }
+
+func (db *DB) CreateTable(ctx context.Context, opt offline.CreateTableOpt) error {
+	return sqlutil.CreateTable(ctx, db.DB, opt, Backend)
+}
diff --git a/internal/database/offline/sqlite/store.go b/internal/database/offline/sqlite/store.go
index 7b526327..de78990e 100644
--- a/internal/database/offline/sqlite/store.go
+++ b/internal/database/offline/sqlite/store.go
@@ -63,3 +63,7 @@ func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
 	dbOpt := dbutil.DBOpt{Backend: Backend, SqlxDB: db.DB}
 	return sqlutil.Snapshot(ctx, dbOpt, opt)
 }
+
+func (db *DB) CreateTable(ctx context.Context, opt offline.CreateTableOpt) error {
+	return sqlutil.CreateTable(ctx, db.DB, opt, Backend)
+}
diff --git a/internal/database/offline/sqlutil/create_table.go b/internal/database/offline/sqlutil/create_table.go
new file mode 100644
index 00000000..0edc2b79
--- /dev/null
+++ b/internal/database/offline/sqlutil/create_table.go
@@ -0,0 +1,18 @@
+package sqlutil
+
+import (
+	"context"
+
+	"github.com/jmoiron/sqlx"
+	"github.com/oom-ai/oomstore/internal/database/dbutil"
+	"github.com/oom-ai/oomstore/internal/database/offline"
+	"github.com/oom-ai/oomstore/pkg/oomstore/types"
+)
+
+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
+	}
+	return nil
+}
diff --git a/internal/database/offline/store.go b/internal/database/offline/store.go
index 305ccc46..7ae7c755 100644
--- a/internal/database/offline/store.go
+++ b/internal/database/offline/store.go
@@ -12,6 +12,7 @@ type Store interface {
 	Export(ctx context.Context, opt ExportOpt) (<-chan types.ExportRecord, <-chan error)
 	Import(ctx context.Context, opt ImportOpt) (int64, error)
 
+	CreateTable(ctx context.Context, opt CreateTableOpt) error
 	TableSchema(ctx context.Context, tableName string) (*types.DataTableSchema, error)
 	Snapshot(ctx context.Context, opt SnapshotOpt) error
 
diff --git a/internal/database/offline/types.go b/internal/database/offline/types.go
index 16bac9e7..c0882aef 100644
--- a/internal/database/offline/types.go
+++ b/internal/database/offline/types.go
@@ -58,3 +58,10 @@ type SnapshotOpt struct {
 	Revision     int64
 	PrevRevision int64
 }
+
+type CreateTableOpt struct {
+	TableName      string
+	Entity         *types.Entity
+	Features       types.FeatureList
+	WithUnixMillis bool
+}


Test patch:
diff --git a/internal/database/offline/test_impl/create_table.go b/internal/database/offline/test_impl/create_table.go
new file mode 100644
index 00000000..7efa5bec
--- /dev/null
+++ b/internal/database/offline/test_impl/create_table.go
@@ -0,0 +1,64 @@
+package test_impl
+
+import (
+	"testing"
+
+	"github.com/oom-ai/oomstore/internal/database/offline"
+	"github.com/oom-ai/oomstore/pkg/oomstore/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCreateTable(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyStoreFn) {
+	t.Cleanup(destroyStore)
+
+	ctx, store := prepareStore(t)
+	defer store.Close()
+
+	tableName := "new_table"
+	entity := types.Entity{
+		Name:   "device",
+		Length: 16,
+	}
+	features := []*types.Feature{
+		{
+			Name:      "price",
+			ValueType: types.Int64,
+		},
+		{
+			Name:      "model",
+			ValueType: types.String,
+		},
+	}
+
+	testCases := []struct {
+		describtion   string
+		opt           offline.CreateTableOpt
+		expectedError error
+	}{
+		{
+			describtion: "with unix milli",
+			opt: offline.CreateTableOpt{
+				TableName:      tableName,
+				Entity:         &entity,
+				Features:       features,
+				WithUnixMillis: true,
+			},
+		},
+		{
+			describtion: "without milli",
+			opt: offline.CreateTableOpt{
+				TableName:      tableName,
+				Entity:         &entity,
+				Features:       features,
+				WithUnixMillis: false,
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.describtion, func(t *testing.T) {
+			err := store.CreateTable(ctx, testCase.opt)
+			require.NoError(t, err)
+		})
+	}
+}