Repository: oom-ai/oomstore
Pull request #893 (2022-01-04 07:33:48)

Jinghan/fix integration test for MySQL backend
This PR fixes integration tests for backends combination `MySQL, MySQL, MySQL`:
- import: serialize import feature value (for boolean feature value, convert from string to int)
- export: deserialize export data (for boolean value type, convert from `int` to `bool`)

related to #894 


Patch:
diff --git a/internal/database/dbutil/data_loader.go b/internal/database/dbutil/data_loader.go
index 0ab18e10..0c11f416 100644
--- a/internal/database/dbutil/data_loader.go
+++ b/internal/database/dbutil/data_loader.go
@@ -12,16 +12,16 @@ import (
 )
 
 // Currying
-func LoadDataFromSource(backendType types.BackendType, batchSize int) func(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSource, tableName string, header []string) error {
-	return func(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSource, tableName string, header []string) error {
-		return loadDataFromSource(tx, ctx, source, tableName, header, backendType, batchSize)
+func LoadDataFromSource(backendType types.BackendType, batchSize int) func(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSource, tableName string, header []string, features types.FeatureList) error {
+	return func(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSource, tableName string, header []string, features types.FeatureList) error {
+		return loadDataFromSource(tx, ctx, source, tableName, header, features, backendType, batchSize)
 	}
 }
 
-func loadDataFromSource(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSource, tableName string, header []string, backendType types.BackendType, batchSize int) error {
+func loadDataFromSource(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSource, tableName string, header []string, features types.FeatureList, backendType types.BackendType, batchSize int) error {
 	records := make([]interface{}, 0, batchSize)
 	for {
-		record, err := ReadLine(source.Reader, source.Delimiter)
+		record, err := ReadLine(source.Reader, source.Delimiter, features, backendType)
 		if err == io.EOF {
 			break
 		}
@@ -45,10 +45,36 @@ func loadDataFromSource(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSou
 	return nil
 }
 
-func ReadLine(reader *bufio.Reader, delimiter string) ([]string, error) {
+func ReadLine(reader *bufio.Reader, delimiter string, features types.FeatureList, backend types.BackendType) ([]interface{}, error) {
 	row, err := reader.ReadString('\n')
 	if err != nil {
 		return nil, err
 	}
-	return strings.Split(strings.Trim(row, "\n"), delimiter), nil
+	rowSlice := strings.Split(strings.Trim(row, "\n"), delimiter)
+	line := make([]interface{}, 0, len(rowSlice))
+	for i, ele := range rowSlice {
+		if i == 0 || len(features) == 0 {
+			// entity_key doesn't need to change type
+			line = append(line, ele)
+		} else {
+			line = append(line, castElement(ele, features[i-1].ValueType, backend))
+		}
+	}
+	return line, nil
+}
+
+func castElement(s string, valueType types.ValueType, backend types.BackendType) interface{} {
+	if backend != types.BackendMySQL {
+		return s
+	}
+	if valueType != types.Bool {
+		return s
+	}
+	if s == "true" || s == "TRUE" {
+		return 1
+	}
+	if s == "false" || s == "FALSE" {
+		return 0
+	}
+	return s
 }
diff --git a/internal/database/dbutil/schema.go b/internal/database/dbutil/schema.go
index 82ffe28d..cae4bbe1 100644
--- a/internal/database/dbutil/schema.go
+++ b/internal/database/dbutil/schema.go
@@ -1,12 +1,47 @@
 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
@@ -23,96 +58,61 @@ 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,
-	isCDC bool,
+	withUnixMillis bool,
 	features types.FeatureList,
-	pkFields []string,
 	backend types.BackendType,
 ) string {
-	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))
+	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)
 	}
-	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)
+	return buf.String()
 }
 
-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))
+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)
 		}
-		rs = append(rs, c)
+		fields = append(fields, Column{
+			Name:   "unix_milli",
+			DbType: dbType,
+		})
 	}
-
-	// 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})
+	for _, feature := range features {
+		dbType, err := DBValueType(backend, feature.ValueType)
+		if err != nil {
+			panic(err)
 		}
+		fields = append(fields, Column{
+			Name:   feature.Name,
+			DbType: dbType,
+		})
 	}
 
-	// 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 &TableSchemaTmplOpts{
+		TableName: QuoteFn(backend)(tableName),
+		Entity:    entity,
+		Fields:    fields,
+		Backend:   backend,
 	}
-	return
 }
diff --git a/internal/database/offline/bigquery/store.go b/internal/database/offline/bigquery/store.go
index d0d59f81..a02a7314 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.IsCDC, opt.Features, nil, Backend)
+	schema := dbutil.BuildTableSchema(opt.TableName, opt.Entity, opt.WithUnixMillis, opt.Features, Backend)
 	if _, err := db.Client.Query(schema).Run(ctx); err != nil {
 		return err
 	}
diff --git a/internal/database/offline/postgres/util.go b/internal/database/offline/postgres/util.go
index cc34149b..ed11894b 100644
--- a/internal/database/offline/postgres/util.go
+++ b/internal/database/offline/postgres/util.go
@@ -4,13 +4,15 @@ import (
 	"context"
 	"io"
 
+	"github.com/oom-ai/oomstore/pkg/oomstore/types"
+
 	"github.com/jmoiron/sqlx"
 	"github.com/lib/pq"
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
 	"github.com/oom-ai/oomstore/internal/database/offline"
 )
 
-func loadDataFromSource(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSource, tableName string, header []string) error {
+func loadDataFromSource(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSource, tableName string, header []string, features types.FeatureList) error {
 	stmt, err := tx.PreparexContext(ctx, pq.CopyIn(tableName, header...))
 	if err != nil {
 		return err
@@ -18,7 +20,7 @@ func loadDataFromSource(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSou
 	defer stmt.Close()
 
 	for {
-		record, err := dbutil.ReadLine(source.Reader, source.Delimiter)
+		record, err := dbutil.ReadLine(source.Reader, source.Delimiter, features, Backend)
 		if err == io.EOF {
 			break
 		}
@@ -28,11 +30,7 @@ func loadDataFromSource(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSou
 		if len(record) != len(header) {
 			continue
 		}
-		args := make([]interface{}, 0, len(record))
-		for _, v := range record {
-			args = append(args, v)
-		}
-		if _, err := stmt.ExecContext(ctx, args...); err != nil {
+		if _, err := stmt.ExecContext(ctx, record...); err != nil {
 			return err
 		}
 	}
diff --git a/internal/database/offline/sqlutil/create_table.go b/internal/database/offline/sqlutil/create_table.go
index e472a54e..0edc2b79 100644
--- a/internal/database/offline/sqlutil/create_table.go
+++ b/internal/database/offline/sqlutil/create_table.go
@@ -10,24 +10,9 @@ import (
 )
 
 func CreateTable(ctx context.Context, db *sqlx.DB, opt offline.CreateTableOpt, backend types.BackendType) error {
-	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
-		}
+	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/sqlutil/export.go b/internal/database/offline/sqlutil/export.go
index 79bd35d3..49f7009d 100644
--- a/internal/database/offline/sqlutil/export.go
+++ b/internal/database/offline/sqlutil/export.go
@@ -49,19 +49,12 @@ func Export(ctx context.Context, db *sqlx.DB, opt offline.ExportOpt, backendType
 				if record[i+1] == nil {
 					continue
 				}
-				if backendType == types.BackendSnowflake {
-
-					v, err := deserializeByTagForSnowflake(record[i+1], f.ValueType)
-					if err != nil {
-						errs <- fmt.Errorf("failed at deserializeByTag, err=%v", err)
-						return
-					}
-					record[i+1] = v
-				} else {
-					if f.ValueType == types.String {
-						record[i+1] = cast.ToString(record[i+1])
-					}
+				v, err := DeserializeByValueType(record[i+1], f.ValueType, backendType)
+				if err != nil {
+					errs <- err
+					return
 				}
+				record[i+1] = v
 			}
 			stream <- record
 		}
diff --git a/internal/database/offline/sqlutil/import.go b/internal/database/offline/sqlutil/import.go
index 99761f53..7ab8e17d 100644
--- a/internal/database/offline/sqlutil/import.go
+++ b/internal/database/offline/sqlutil/import.go
@@ -10,24 +10,20 @@ import (
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
-type LoadData func(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSource, tableName string, header []string) error
+type LoadData func(tx *sqlx.Tx, ctx context.Context, source *offline.CSVSource, tableName string, header []string, features types.FeatureList) error
 
 func Import(ctx context.Context, db *sqlx.DB, opt offline.ImportOpt, loadData LoadData, backendType types.BackendType) (int64, error) {
 	var revision int64
 	err := dbutil.WithTransaction(db, ctx, func(ctx context.Context, tx *sqlx.Tx) error {
 		// create the data table
-		pkFields := []string{opt.Entity.Name}
-		if opt.NoPK {
-			pkFields = nil
-		}
-		schema := dbutil.BuildTableSchema(opt.SnapshotTableName, opt.Entity, false, opt.Features, pkFields, backendType)
+		schema := dbutil.BuildTableSchema(opt.SnapshotTableName, opt.Entity, false, opt.Features, backendType)
 		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return err
 		}
 
 		// populate the data table
-		err = loadData(tx, ctx, opt.Source, opt.SnapshotTableName, opt.Header)
+		err = loadData(tx, ctx, opt.Source, opt.SnapshotTableName, opt.Header, opt.Features)
 		if err != nil {
 			return err
 		}
diff --git a/internal/database/offline/sqlutil/snapshot.go b/internal/database/offline/sqlutil/snapshot.go
index 8be3fd22..728f447c 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, []string{opt.Group.Entity.Name}, dbOpt.Backend)
+	schema := dbutil.BuildTableSchema(currSnapshotTableName, opt.Group.Entity, false, opt.Features, 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 6f2f4966..c0882aef 100644
--- a/internal/database/offline/types.go
+++ b/internal/database/offline/types.go
@@ -45,7 +45,6 @@ type ImportOpt struct {
 	Revision          *int64
 	SnapshotTableName string
 	Source            *CSVSource
-	NoPK              bool // TODO: to import cdc data temporarily for testing
 }
 
 type CSVSource struct {
@@ -61,8 +60,8 @@ type SnapshotOpt struct {
 }
 
 type CreateTableOpt struct {
-	TableName string
-	Entity    *types.Entity
-	Features  types.FeatureList
-	IsCDC     bool
+	TableName      string
+	Entity         *types.Entity
+	Features       types.FeatureList
+	WithUnixMillis bool
 }
diff --git a/internal/database/online/cassandra/import.go b/internal/database/online/cassandra/import.go
index 1f20861e..3a1fe5a5 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, []string{opt.Entity.Name}, Backend)
+	table := dbutil.BuildTableSchema(tableName, opt.Entity, false, opt.Features, Backend)
 
 	// create table
 	if err := db.Query(table).Exec(); err != nil {
diff --git a/internal/database/online/sqlutil/import.go b/internal/database/online/sqlutil/import.go
index ecbd3b7a..80075eac 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, []string{opt.Entity.Name}, backend)
+		schema := dbutil.BuildTableSchema(tableName, opt.Entity, false, opt.Features, backend)
 		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return err
diff --git a/pkg/oomstore/import.go b/pkg/oomstore/import.go
index 51956571..1d35787b 100644
--- a/pkg/oomstore/import.go
+++ b/pkg/oomstore/import.go
@@ -7,6 +7,8 @@ import (
 	"os"
 	"time"
 
+	"github.com/spf13/cast"
+
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
 	"github.com/oom-ai/oomstore/internal/database/metadata"
 	"github.com/oom-ai/oomstore/internal/database/offline"
@@ -44,15 +46,16 @@ func (s *OomStore) Import(ctx context.Context, opt types.ImportOpt) (int, error)
 func (s *OomStore) csvReaderImport(ctx context.Context, opt *importOpt, dataSource *types.CsvReaderDataSource) (int, error) {
 	//make sure csv data source has all defined columns
 	reader := bufio.NewReader(dataSource.Reader)
-	header, err := dbutil.ReadLine(reader, dataSource.Delimiter)
+	// read header does not need pass down features
+	header, err := dbutil.ReadLine(reader, dataSource.Delimiter, nil, "")
 	if err != nil {
 		return 0, err
 	}
-	if hasDup(header) {
+	if hasDup(cast.ToStringSlice(header)) {
 		return 0, fmt.Errorf("csv data source has duplicated columns: %v", header)
 	}
 	columnNames := append([]string{opt.entity.Name}, opt.features.Names()...)
-	if !stringSliceEqual(header, columnNames) {
+	if !stringSliceEqual(cast.ToStringSlice(header), columnNames) {
 		return 0, fmt.Errorf("csv header of the data source %v doesn't match the feature group schema %v", header, columnNames)
 	}
 
@@ -70,7 +73,7 @@ func (s *OomStore) csvReaderImport(ctx context.Context, opt *importOpt, dataSour
 	revision, err := s.offline.Import(ctx, offline.ImportOpt{
 		Entity:            opt.entity,
 		Features:          opt.features,
-		Header:            header,
+		Header:            cast.ToStringSlice(header),
 		Revision:          opt.Revision,
 		SnapshotTableName: snapshotTableName,
 		Source: &offline.CSVSource{


Test patch:
diff --git a/internal/database/dbutil/schema_test.go b/internal/database/dbutil/schema_test.go
index 37846274..62813484 100644
--- a/internal/database/dbutil/schema_test.go
+++ b/internal/database/dbutil/schema_test.go
@@ -1,7 +1,6 @@
 package dbutil_test
 
 import (
-	"strings"
 	"testing"
 
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
@@ -87,8 +86,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, nil, c.backend)
-			require.Equal(t, strings.TrimSpace(c.want), schema)
+			schema := dbutil.BuildTableSchema(c.tableName, c.entity, false, c.features, c.backend)
+			require.Equal(t, c.want, schema)
 		})
 	}
 }
diff --git a/internal/database/offline/test_impl/create_table.go b/internal/database/offline/test_impl/create_table.go
index 33573366..7efa5bec 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: "cdc table (with unix_milli)",
+			describtion: "with unix milli",
 			opt: offline.CreateTableOpt{
-				TableName: tableName,
-				Entity:    &entity,
-				Features:  features,
-				IsCDC:     true,
+				TableName:      tableName,
+				Entity:         &entity,
+				Features:       features,
+				WithUnixMillis: true,
 			},
 		},
 		{
-			describtion: "snapshot table (without_milli)",
+			describtion: "without milli",
 			opt: offline.CreateTableOpt{
-				TableName: tableName,
-				Entity:    &entity,
-				Features:  features,
-				IsCDC:     false,
+				TableName:      tableName,
+				Entity:         &entity,
+				Features:       features,
+				WithUnixMillis: false,
 			},
 		},
 	}
diff --git a/internal/database/offline/test_impl/util.go b/internal/database/offline/test_impl/util.go
index 8a18f29f..bc468741 100644
--- a/internal/database/offline/test_impl/util.go
+++ b/internal/database/offline/test_impl/util.go
@@ -37,7 +37,6 @@ func buildTestSnapshotTable(
 		Header:            header,
 		Source:            source,
 		Revision:          &revision,
-		NoPK:              true,
 	}
 	_, err := store.Import(ctx, opt)
 	require.NoError(t, err)