Repository: oom-ai/oomstore
Pull request #1076 (2022-01-25 07:17:41)

Jinghan/simplify offline method options
This PR refactors `ImportOpt`, `CreateTableOpt`, `JoinOpt`: change `Entity` to `EntityName`, since in offline store, we only need `EntityName`.



Patch:
diff --git a/internal/database/dbutil/data_loader.go b/internal/database/dbutil/data_loader.go
index 62f752e0..7f27b67c 100644
--- a/internal/database/dbutil/data_loader.go
+++ b/internal/database/dbutil/data_loader.go
@@ -13,12 +13,12 @@ import (
 )
 
 type LoadDataFromSourceOpt struct {
-	Source    *offline.CSVSource
-	Entity    *types.Entity
-	TableName string
-	Header    []string
-	Features  types.FeatureList
-	Backend   types.BackendType
+	Source     *offline.CSVSource
+	EntityName string
+	TableName  string
+	Header     []string
+	Features   types.FeatureList
+	Backend    types.BackendType
 }
 
 // Currying
@@ -32,10 +32,10 @@ func loadDataFromSource(tx *sqlx.Tx, ctx context.Context, opt LoadDataFromSource
 	records := make([]interface{}, 0, batchSize)
 	for {
 		record, err := ReadLine(ReadLineOpt{
-			Source:   opt.Source,
-			Entity:   opt.Entity,
-			Header:   opt.Header,
-			Features: opt.Features,
+			Source:     opt.Source,
+			EntityName: opt.EntityName,
+			Header:     opt.Header,
+			Features:   opt.Features,
 		})
 		if errdefs.Cause(err) == io.EOF {
 			break
@@ -61,10 +61,10 @@ func loadDataFromSource(tx *sqlx.Tx, ctx context.Context, opt LoadDataFromSource
 }
 
 type ReadLineOpt struct {
-	Source   *offline.CSVSource
-	Entity   *types.Entity
-	Header   []string
-	Features types.FeatureList
+	Source     *offline.CSVSource
+	EntityName string
+	Header     []string
+	Features   types.FeatureList
 }
 
 func ReadLine(opt ReadLineOpt) ([]interface{}, error) {
@@ -75,7 +75,7 @@ func ReadLine(opt ReadLineOpt) ([]interface{}, error) {
 	rowSlice := strings.Split(strings.Trim(row, "\n"), opt.Source.Delimiter)
 	line := make([]interface{}, 0, len(rowSlice))
 	for i, ele := range rowSlice {
-		if len(opt.Header) == 0 || len(opt.Features) == 0 || opt.Header[i] == opt.Entity.Name {
+		if len(opt.Header) == 0 || len(opt.Features) == 0 || opt.Header[i] == opt.EntityName {
 			// entity_key doesn't need to change type
 			line = append(line, ele)
 		} else if opt.Header[i] == "unix_milli" {
diff --git a/internal/database/dbutil/schema.go b/internal/database/dbutil/schema.go
index d51c747e..8ff1bebf 100644
--- a/internal/database/dbutil/schema.go
+++ b/internal/database/dbutil/schema.go
@@ -25,13 +25,13 @@ func (c ColumnList) Names() []string {
 
 func BuildTableSchema(
 	tableName string,
-	entity *types.Entity,
+	entityName string,
 	hasUnixMilli bool,
 	features types.FeatureList,
 	pkFields []string,
 	backend types.BackendType,
 ) string {
-	columns := parseColumns(entity, hasUnixMilli, features, backend)
+	columns := parseColumns(entityName, hasUnixMilli, features, backend)
 	return createTableDDL(tableName, columns, pkFields, backend)
 }
 
@@ -73,10 +73,10 @@ func BuildIndexDDL(tableName string, indexName string, fields []string, backend
 	return fmt.Sprintf("CREATE INDEX %s ON %s (%s)", qt(index), qt(tableName), qt(fields...))
 }
 
-func parseColumns(entity *types.Entity, hasUnixMilli bool, features types.FeatureList, backend types.BackendType) (rs []Column) {
+func parseColumns(entityName string, hasUnixMilli bool, features types.FeatureList, backend types.BackendType) (rs []Column) {
 	// entity column
 	{
-		c := Column{Name: entity.Name, ValueType: types.String}
+		c := Column{Name: entityName, ValueType: types.String}
 		switch backend {
 		case types.BackendCassandra, types.BackendSQLite, types.BackendPostgres, types.BackendRedshift, types.BackendSnowflake:
 			c.DbType = "TEXT"
diff --git a/internal/database/offline/bigquery/import.go b/internal/database/offline/bigquery/import.go
index 94847bf7..69324c1e 100644
--- a/internal/database/offline/bigquery/import.go
+++ b/internal/database/offline/bigquery/import.go
@@ -17,9 +17,9 @@ func (db *DB) Import(ctx context.Context, opt offline.ImportOpt) (int64, error)
 	// Step 1: define table schema
 	schema := make(bigquery.Schema, 0, len(opt.Features)+1)
 	for _, h := range opt.Header {
-		if h == opt.Entity.Name {
+		if h == opt.EntityName {
 			schema = append(schema, &bigquery.FieldSchema{
-				Name:        opt.Entity.Name,
+				Name:        opt.EntityName,
 				Type:        bigquery.StringFieldType,
 				Description: "entity key",
 			})
diff --git a/internal/database/offline/postgres/util.go b/internal/database/offline/postgres/util.go
index 7b1d36bf..ecb4c41b 100644
--- a/internal/database/offline/postgres/util.go
+++ b/internal/database/offline/postgres/util.go
@@ -20,10 +20,10 @@ func loadDataFromSource(tx *sqlx.Tx, ctx context.Context, opt dbutil.LoadDataFro
 
 	for {
 		record, err := dbutil.ReadLine(dbutil.ReadLineOpt{
-			Source:   opt.Source,
-			Entity:   opt.Entity,
-			Header:   opt.Header,
-			Features: opt.Features,
+			Source:     opt.Source,
+			EntityName: opt.EntityName,
+			Header:     opt.Header,
+			Features:   opt.Features,
 		})
 		if errdefs.Cause(err) == io.EOF {
 			break
diff --git a/internal/database/offline/sqlutil/create_table.go b/internal/database/offline/sqlutil/create_table.go
index b91f24ed..78e272d2 100644
--- a/internal/database/offline/sqlutil/create_table.go
+++ b/internal/database/offline/sqlutil/create_table.go
@@ -16,25 +16,25 @@ func CreateTable(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.CreateTabl
 	switch opt.TableType {
 	case types.TableBatchSnapshot:
 		// Create primary key (entity_key) on batch snapshot table
-		pkFields := []string{opt.Entity.Name}
-		schema := dbutil.BuildTableSchema(opt.TableName, opt.Entity, false, opt.Features, pkFields, dbOpt.Backend)
+		pkFields := []string{opt.EntityName}
+		schema := dbutil.BuildTableSchema(opt.TableName, opt.EntityName, false, opt.Features, pkFields, dbOpt.Backend)
 		if err := dbOpt.ExecContext(ctx, schema, nil); err != nil {
 			return err
 		}
 	case types.TableStreamSnapshot:
 		// Create primary key (entity_key) on stream snapshot table
-		pkFields := []string{opt.Entity.Name}
-		schema := dbutil.BuildTableSchema(opt.TableName, opt.Entity, true, opt.Features, pkFields, dbOpt.Backend)
+		pkFields := []string{opt.EntityName}
+		schema := dbutil.BuildTableSchema(opt.TableName, opt.EntityName, true, opt.Features, pkFields, dbOpt.Backend)
 		if err := dbOpt.ExecContext(ctx, schema, nil); err != nil {
 			return err
 		}
 	case types.TableStreamCdc:
-		schema := dbutil.BuildTableSchema(opt.TableName, opt.Entity, true, opt.Features, nil, dbOpt.Backend)
+		schema := dbutil.BuildTableSchema(opt.TableName, opt.EntityName, true, opt.Features, nil, dbOpt.Backend)
 		if err := dbOpt.ExecContext(ctx, schema, nil); err != nil {
 			return err
 		}
 		// Create index (entity_key, unix_milli) on stream cdc table
-		indexFields := []string{opt.Entity.Name, "unix_milli"}
+		indexFields := []string{opt.EntityName, "unix_milli"}
 		indexDDL := dbutil.BuildIndexDDL(opt.TableName, "idx", indexFields, dbOpt.Backend)
 		if err := dbOpt.ExecContext(ctx, indexDDL, nil); err != nil {
 			return err
@@ -60,7 +60,7 @@ func CreateTableNoIndex(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.Cre
 	if dbOpt.Backend == types.BackendBigQuery {
 		tableName = fmt.Sprintf("%s.%s", *dbOpt.DatasetID, tableName)
 	}
-	schema := dbutil.BuildTableSchema(tableName, opt.Entity, hasUnixMilli, opt.Features, nil, dbOpt.Backend)
+	schema := dbutil.BuildTableSchema(tableName, opt.EntityName, hasUnixMilli, opt.Features, nil, dbOpt.Backend)
 	if err := dbOpt.ExecContext(ctx, schema, nil); err != nil {
 		return err
 	}
diff --git a/internal/database/offline/sqlutil/import.go b/internal/database/offline/sqlutil/import.go
index 26a66cd4..b5f6b915 100644
--- a/internal/database/offline/sqlutil/import.go
+++ b/internal/database/offline/sqlutil/import.go
@@ -18,11 +18,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
-		pkFields := []string{opt.Entity.Name}
+		pkFields := []string{opt.EntityName}
 		if opt.NoPK {
 			pkFields = nil
 		}
-		schema := dbutil.BuildTableSchema(opt.SnapshotTableName, opt.Entity, false, opt.Features, pkFields, backend)
+		schema := dbutil.BuildTableSchema(opt.SnapshotTableName, opt.EntityName, false, opt.Features, pkFields, backend)
 		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return errdefs.WithStack(err)
@@ -30,12 +30,12 @@ func Import(ctx context.Context, db *sqlx.DB, opt offline.ImportOpt, loadData Lo
 
 		// populate the data table
 		err = loadData(tx, ctx, dbutil.LoadDataFromSourceOpt{
-			Source:    opt.Source,
-			Entity:    opt.Entity,
-			TableName: opt.SnapshotTableName,
-			Header:    opt.Header,
-			Features:  opt.Features,
-			Backend:   backend,
+			Source:     opt.Source,
+			EntityName: opt.EntityName,
+			TableName:  opt.SnapshotTableName,
+			Header:     opt.Header,
+			Features:   opt.Features,
+			Backend:    backend,
 		})
 		if err != nil {
 			return err
diff --git a/internal/database/offline/sqlutil/join.go b/internal/database/offline/sqlutil/join.go
index d5665a83..9af0f008 100644
--- a/internal/database/offline/sqlutil/join.go
+++ b/internal/database/offline/sqlutil/join.go
@@ -62,7 +62,7 @@ func DoJoin(ctx context.Context, dbOpt dbutil.DBOpt, opt DoJoinOpt) (*types.Join
 	}
 
 	// Step 1: prepare temporary table entity_rows
-	entityRowsTableName, err := PrepareEntityRowsTable(ctx, dbOpt, opt.Entity, opt.EntityRows, opt.ValueNames)
+	entityRowsTableName, err := prepareEntityRowsTable(ctx, dbOpt, opt.EntityRows, opt.ValueNames)
 	if err != nil {
 		return nil, err
 	}
@@ -83,10 +83,10 @@ func DoJoin(ctx context.Context, dbOpt dbutil.DBOpt, opt DoJoinOpt) (*types.Join
 		} else {
 			category = types.CategoryStream
 		}
-		joinedTables, err := JoinOneGroup(ctx, dbOpt, offline.JoinOneGroupOpt{
+		joinedTables, err := joinOneGroup(ctx, dbOpt, offline.JoinOneGroupOpt{
 			GroupName:           groupName,
 			Category:            category,
-			Entity:              opt.Entity,
+			EntityName:          opt.EntityName,
 			Features:            featureList,
 			RevisionRanges:      revisionRanges,
 			EntityRowsTableName: entityRowsTableName,
@@ -115,13 +115,13 @@ func DoJoin(ctx context.Context, dbOpt dbutil.DBOpt, opt DoJoinOpt) (*types.Join
 	})
 }
 
-func JoinOneGroup(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.JoinOneGroupOpt) ([]string, error) {
+func joinOneGroup(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.JoinOneGroupOpt) ([]string, error) {
 	if len(opt.Features) == 0 {
 		return nil, nil
 	}
 
 	// Step 1: create temporary joined table
-	snapshotJoinedTableName, err := PrepareJoinedTable(ctx, dbOpt, opt.Features, opt.Entity, opt.GroupName, opt.ValueNames)
+	snapshotJoinedTableName, err := prepareJoinedTable(ctx, dbOpt, opt.Features, opt.GroupName, opt.ValueNames)
 	if err != nil {
 		return nil, err
 	}
@@ -131,7 +131,7 @@ func JoinOneGroup(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.JoinOneGr
 	for _, r := range opt.RevisionRanges {
 		query, err := buildJoinQuery(joinQueryParams{
 			TableName:           snapshotJoinedTableName,
-			EntityName:          opt.Entity.Name,
+			EntityName:          opt.EntityName,
 			EntityKey:           "entity_key",
 			UnixMilli:           "unix_milli",
 			Columns:             columns,
@@ -152,7 +152,7 @@ func JoinOneGroup(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.JoinOneGr
 	}
 
 	// Step 3: for streaming features, keep joining with cdc_table
-	cdcJoinedTableName, err := PrepareJoinedTable(ctx, dbOpt, opt.Features, opt.Entity, opt.GroupName, opt.ValueNames)
+	cdcJoinedTableName, err := prepareJoinedTable(ctx, dbOpt, opt.Features, opt.GroupName, opt.ValueNames)
 	if err != nil {
 		return nil, err
 	}
@@ -160,7 +160,7 @@ func JoinOneGroup(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.JoinOneGr
 		query, err := buildCdcJoinQuery(cdcJoinQueryParams{
 			TableName:           cdcJoinedTableName,
 			EntityKey:           "entity_key",
-			EntityName:          opt.Entity.Name,
+			EntityName:          opt.EntityName,
 			UnixMilli:           "unix_milli",
 			ValueNames:          opt.ValueNames,
 			FeatureNames:        opt.Features.Names(),
diff --git a/internal/database/offline/sqlutil/join_helper.go b/internal/database/offline/sqlutil/join_helper.go
index f9daad57..9d419088 100644
--- a/internal/database/offline/sqlutil/join_helper.go
+++ b/internal/database/offline/sqlutil/join_helper.go
@@ -69,11 +69,10 @@ LEFT JOIN {{ $t2 }}
 ON {{ $t1 }}.{{ qt $.UnixMilli }} = {{ $t2 }}.{{ qt $.UnixMilli }} AND {{ $t1 }}.{{ qt $.EntityKey }} = {{ $t2 }}.{{ qt $.EntityKey }}
 {{end}}`
 
-func PrepareJoinedTable(
+func prepareJoinedTable(
 	ctx context.Context,
 	dbOpt dbutil.DBOpt,
 	features types.FeatureList,
-	entity types.Entity,
 	groupName string,
 	valueNames []string,
 ) (string, error) {
@@ -121,9 +120,8 @@ func PrepareJoinedTable(
 	return tableName, nil
 }
 
-func PrepareEntityRowsTable(ctx context.Context,
+func prepareEntityRowsTable(ctx context.Context,
 	dbOpt dbutil.DBOpt,
-	entity types.Entity,
 	entityRows <-chan types.EntityRow,
 	valueNames []string,
 ) (string, error) {
diff --git a/internal/database/offline/sqlutil/snapshot.go b/internal/database/offline/sqlutil/snapshot.go
index 7f70963b..d464d981 100644
--- a/internal/database/offline/sqlutil/snapshot.go
+++ b/internal/database/offline/sqlutil/snapshot.go
@@ -71,7 +71,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, true, opt.Features, []string{opt.Group.Entity.Name}, dbOpt.Backend)
+	schema := dbutil.BuildTableSchema(currSnapshotTableName, opt.Group.Entity.Name, true, 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 8d378f4a..514e83e8 100644
--- a/internal/database/offline/types.go
+++ b/internal/database/offline/types.go
@@ -23,7 +23,7 @@ type RevisionRange struct {
 }
 
 type JoinOpt struct {
-	Entity           types.Entity
+	EntityName       string
 	EntityRows       <-chan types.EntityRow
 	GroupNames       []string
 	FeatureMap       map[string]types.FeatureList
@@ -36,17 +36,17 @@ type JoinOneGroupOpt struct {
 	Category            types.Category
 	Features            types.FeatureList
 	RevisionRanges      []*RevisionRange
-	Entity              types.Entity
+	EntityName          string
 	EntityRowsTableName string
 	ValueNames          []string
 }
 
 type ImportOpt struct {
-	Entity            *types.Entity
-	Features          types.FeatureList
+	EntityName        string
+	SnapshotTableName string
 	Header            []string
+	Features          types.FeatureList
 	Revision          *int64
-	SnapshotTableName string
 	Source            *CSVSource
 	NoPK              bool // TODO: to import cdc data temporarily for testing
 }
@@ -65,17 +65,17 @@ type CSVSource struct {
 }
 
 type SnapshotOpt struct {
-	Group        *types.Group
+	Group        types.Group
 	Features     types.FeatureList
 	Revision     int64
 	PrevRevision int64
 }
 
 type CreateTableOpt struct {
-	TableName string
-	Entity    *types.Entity
-	Features  types.FeatureList
-	TableType types.TableType
+	TableName  string
+	EntityName string
+	Features   types.FeatureList
+	TableType  types.TableType
 }
 
 type TableSchemaOpt struct {
diff --git a/internal/database/online/cassandra/import.go b/internal/database/online/cassandra/import.go
index 478e9381..3f40cbd0 100644
--- a/internal/database/online/cassandra/import.go
+++ b/internal/database/online/cassandra/import.go
@@ -29,7 +29,7 @@ func (db *DB) Import(ctx context.Context, opt online.ImportOpt) error {
 	// Step 1: create online table
 	entity := opt.Group.Entity
 	columns := append([]string{entity.Name}, opt.Features.Names()...)
-	schema := dbutil.BuildTableSchema(tableName, entity, false, opt.Features, []string{entity.Name}, Backend)
+	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)
 	}
diff --git a/internal/database/online/sqlutil/import.go b/internal/database/online/sqlutil/import.go
index 95c56adf..26dc2010 100644
--- a/internal/database/online/sqlutil/import.go
+++ b/internal/database/online/sqlutil/import.go
@@ -24,7 +24,7 @@ func Import(ctx context.Context, db *sqlx.DB, opt online.ImportOpt, backend type
 	columns := append([]string{entity.Name}, opt.Features.Names()...)
 	err := dbutil.WithTransaction(db, ctx, func(ctx context.Context, tx *sqlx.Tx) error {
 		// create the data table
-		schema := dbutil.BuildTableSchema(tableName, entity, false, opt.Features, []string{entity.Name}, backend)
+		schema := dbutil.BuildTableSchema(tableName, entity.Name, false, opt.Features, []string{entity.Name}, backend)
 		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return errdefs.WithStack(err)
diff --git a/pkg/oomstore/import.go b/pkg/oomstore/import.go
index d000a844..38a82312 100644
--- a/pkg/oomstore/import.go
+++ b/pkg/oomstore/import.go
@@ -51,7 +51,7 @@ func (s *OomStore) csvReaderImport(ctx context.Context, opt *importOpt, dataSour
 		Delimiter: dataSource.Delimiter,
 	}
 	// read header
-	columnNames := append([]string{opt.entity.Name}, opt.features.Names()...)
+	columnNames := append([]string{opt.entityName}, opt.features.Names()...)
 	header, err := readHeader(source, columnNames)
 	if err != nil {
 		return 0, err
@@ -69,7 +69,7 @@ func (s *OomStore) csvReaderImport(ctx context.Context, opt *importOpt, dataSour
 	}
 
 	revision, err := s.offline.Import(ctx, offline.ImportOpt{
-		Entity:            opt.entity,
+		EntityName:        opt.entityName,
 		Features:          opt.features,
 		Header:            cast.ToStringSlice(header),
 		Revision:          opt.Revision,
@@ -199,10 +199,10 @@ func (s *OomStore) parseImportOpt(ctx context.Context, opt types.ImportOpt) (*im
 		return nil, err
 	}
 	return &importOpt{
-		ImportOpt: &opt,
-		entity:    entity,
-		group:     group,
-		features:  features,
+		ImportOpt:  &opt,
+		entityName: entity.Name,
+		group:      group,
+		features:   features,
 	}, nil
 }
 
@@ -232,7 +232,7 @@ func (s *OomStore) getGroupInfo(ctx context.Context, groupName string) (*types.E
 
 type importOpt struct {
 	*types.ImportOpt
-	entity   *types.Entity
-	group    *types.Group
-	features types.FeatureList
+	entityName string
+	group      *types.Group
+	features   types.FeatureList
 }
diff --git a/pkg/oomstore/import_stream.go b/pkg/oomstore/import_stream.go
index 1ea6a554..b6335459 100644
--- a/pkg/oomstore/import_stream.go
+++ b/pkg/oomstore/import_stream.go
@@ -66,10 +66,10 @@ func (s *OomStore) csvReaderImportStream(ctx context.Context, opt types.ImportSt
 	records := make([]types.StreamRecord, 0, ImportStreamBatchSize)
 	for {
 		line, err := dbutil.ReadLine(dbutil.ReadLineOpt{
-			Source:   source,
-			Entity:   entity,
-			Header:   header,
-			Features: features,
+			Source:     source,
+			EntityName: entity.Name,
+			Header:     header,
+			Features:   features,
 		})
 		if errdefs.Cause(err) == io.EOF {
 			break
diff --git a/pkg/oomstore/join.go b/pkg/oomstore/join.go
index 7f157431..acef96bf 100644
--- a/pkg/oomstore/join.go
+++ b/pkg/oomstore/join.go
@@ -62,7 +62,7 @@ func (s *OomStore) ChannelJoin(ctx context.Context, opt types.ChannelJoinOpt) (*
 	}
 
 	return s.offline.Join(ctx, offline.JoinOpt{
-		Entity:           *entity,
+		EntityName:       entity.Name,
 		EntityRows:       opt.EntityRows,
 		GroupNames:       groupNames,
 		FeatureMap:       featureMap,
diff --git a/pkg/oomstore/push_processor.go b/pkg/oomstore/push_processor.go
index 01c9370e..6e668ff3 100644
--- a/pkg/oomstore/push_processor.go
+++ b/pkg/oomstore/push_processor.go
@@ -153,10 +153,10 @@ func (s *OomStore) newRevisionForStream(ctx context.Context, groupID int, revisi
 	entity := features[0].Entity()
 
 	if err := s.offline.CreateTable(ctx, offline.CreateTableOpt{
-		TableName: dbutil.OfflineStreamCdcTableName(groupID, revision),
-		Entity:    entity,
-		Features:  features,
-		TableType: types.TableStreamCdc,
+		TableName:  dbutil.OfflineStreamCdcTableName(groupID, revision),
+		EntityName: entity.Name,
+		Features:   features,
+		TableType:  types.TableStreamCdc,
 	}); err != nil {
 		return err
 	}
diff --git a/pkg/oomstore/revision.go b/pkg/oomstore/revision.go
index 023ae5c0..2ee4656a 100644
--- a/pkg/oomstore/revision.go
+++ b/pkg/oomstore/revision.go
@@ -77,10 +77,10 @@ func (s *OomStore) createSnapshotAndCdcTable(ctx context.Context, revision *type
 	}
 
 	if err = s.offline.CreateTable(ctx, offline.CreateTableOpt{
-		TableName: snapshotTable,
-		Entity:    revision.Group.Entity,
-		Features:  features,
-		TableType: types.TableStreamSnapshot,
+		TableName:  snapshotTable,
+		EntityName: revision.Group.Entity.Name,
+		Features:   features,
+		TableType:  types.TableStreamSnapshot,
 	}); err != nil {
 		return err
 	}
@@ -89,10 +89,10 @@ func (s *OomStore) createSnapshotAndCdcTable(ctx context.Context, revision *type
 	if revision.Group.Category == types.CategoryStream {
 		tableName := dbutil.OfflineStreamCdcTableName(revision.GroupID, revision.Revision)
 		if err = s.offline.CreateTable(ctx, offline.CreateTableOpt{
-			TableName: tableName,
-			Entity:    revision.Group.Entity,
-			Features:  features,
-			TableType: types.TableStreamCdc,
+			TableName:  tableName,
+			EntityName: revision.Group.Entity.Name,
+			Features:   features,
+			TableType:  types.TableStreamCdc,
 		}); err != nil {
 			return err
 		}
@@ -127,10 +127,10 @@ func (s *OomStore) createFirstSnapshotTable(ctx context.Context, revision *types
 		return err
 	}
 	if err = s.offline.CreateTable(ctx, offline.CreateTableOpt{
-		TableName: snapshotTable,
-		Entity:    revision.Group.Entity,
-		Features:  features,
-		TableType: types.TableStreamSnapshot,
+		TableName:  snapshotTable,
+		EntityName: revision.Group.Entity.Name,
+		Features:   features,
+		TableType:  types.TableStreamSnapshot,
 	}); err != nil {
 		return err
 	}
diff --git a/pkg/oomstore/snapshot.go b/pkg/oomstore/snapshot.go
index 566a5a50..c8b519ab 100644
--- a/pkg/oomstore/snapshot.go
+++ b/pkg/oomstore/snapshot.go
@@ -50,7 +50,7 @@ func (s *OomStore) Snapshot(ctx context.Context, groupName string) error {
 		}
 		tableName := dbutil.OfflineStreamSnapshotTableName(group.ID, revision.Revision)
 		if err = s.offline.Snapshot(ctx, offline.SnapshotOpt{
-			Group:        group,
+			Group:        *group,
 			Features:     features,
 			Revision:     revisions[i].Revision,
 			PrevRevision: revisions[i-1].Revision,


Test patch:
diff --git a/internal/database/dbutil/schema_test.go b/internal/database/dbutil/schema_test.go
index 36bb6270..e11ea001 100644
--- a/internal/database/dbutil/schema_test.go
+++ b/internal/database/dbutil/schema_test.go
@@ -81,7 +81,7 @@ 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)
+			schema := dbutil.BuildTableSchema(c.tableName, c.entity.Name, 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 77eb4222..2c77b395 100644
--- a/internal/database/offline/test_impl/create_table.go
+++ b/internal/database/offline/test_impl/create_table.go
@@ -35,28 +35,28 @@ func TestCreateTable(t *testing.T, prepareStore PrepareStoreFn, destroyStore Des
 		{
 			describtion: "cdc table (with unix_milli)",
 			opt: offline.CreateTableOpt{
-				TableName: tableName + "_0",
-				Entity:    &entity,
-				Features:  features,
-				TableType: types.TableStreamCdc,
+				TableName:  tableName + "_0",
+				EntityName: entity.Name,
+				Features:   features,
+				TableType:  types.TableStreamCdc,
 			},
 		},
 		{
 			describtion: "stream snapshot table (with unix_milli)",
 			opt: offline.CreateTableOpt{
-				TableName: tableName + "_1",
-				Entity:    &entity,
-				Features:  features,
-				TableType: types.TableStreamSnapshot,
+				TableName:  tableName + "_1",
+				EntityName: entity.Name,
+				Features:   features,
+				TableType:  types.TableStreamSnapshot,
 			},
 		},
 		{
 			describtion: "batch snapshot table (without unix_milli)",
 			opt: offline.CreateTableOpt{
-				TableName: tableName + "_2",
-				Entity:    &entity,
-				Features:  features,
-				TableType: types.TableBatchSnapshot,
+				TableName:  tableName + "_2",
+				EntityName: entity.Name,
+				Features:   features,
+				TableType:  types.TableBatchSnapshot,
 			},
 		},
 	}
diff --git a/internal/database/offline/test_impl/import.go b/internal/database/offline/test_impl/import.go
index 5c171ecb..fa33b3d8 100644
--- a/internal/database/offline/test_impl/import.go
+++ b/internal/database/offline/test_impl/import.go
@@ -28,7 +28,7 @@ func TestImport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 	snapshotTable := "offline_1_1"
 
 	opt := offline.ImportOpt{
-		Entity:            &entity,
+		EntityName:        entity.Name,
 		SnapshotTableName: snapshotTable,
 		Header:            []string{"device", "model", "price"},
 		Source: &offline.CSVSource{
diff --git a/internal/database/offline/test_impl/join.go b/internal/database/offline/test_impl/join.go
index 37a595b1..00b609c6 100644
--- a/internal/database/offline/test_impl/join.go
+++ b/internal/database/offline/test_impl/join.go
@@ -71,7 +71,6 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroySto
 		{
 			description: "no features",
 			opt: offline.JoinOpt{
-
 				FeatureMap: make(map[string]types.FeatureList),
 			},
 			expected: prepareEmptyResult(),
@@ -79,7 +78,7 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroySto
 		{
 			description: "no entity rows",
 			opt: offline.JoinOpt{
-				Entity:     *entity,
+				EntityName: entity.Name,
 				GroupNames: oneGroupGroupNames,
 				FeatureMap: oneGroupFeatureMap,
 				EntityRows: prepareEntityRows(true, false),
@@ -89,7 +88,7 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroySto
 		{
 			description: "one batch feature group",
 			opt: offline.JoinOpt{
-				Entity:           *entity,
+				EntityName:       entity.Name,
 				EntityRows:       prepareEntityRows(false, false),
 				GroupNames:       oneGroupGroupNames,
 				FeatureMap:       oneGroupFeatureMap,
@@ -100,7 +99,7 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroySto
 		{
 			description: "two batch feature groups",
 			opt: offline.JoinOpt{
-				Entity:           *entity,
+				EntityName:       entity.Name,
 				EntityRows:       prepareEntityRows(false, false),
 				GroupNames:       twoGroupGroupNames,
 				FeatureMap:       twoGroupFeatureMap,
@@ -111,7 +110,7 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroySto
 		{
 			description: "two batch feature groups, with extra values",
 			opt: offline.JoinOpt{
-				Entity:           *entity,
+				EntityName:       entity.Name,
 				EntityRows:       prepareEntityRows(false, true),
 				GroupNames:       twoGroupGroupNames,
 				FeatureMap:       twoGroupFeatureMap,
@@ -123,7 +122,7 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroySto
 		{
 			description: "one streaming feature group, one batch group",
 			opt: offline.JoinOpt{
-				Entity:           *entity,
+				EntityName:       entity.Name,
 				EntityRows:       prepareEntityRows(false, true),
 				GroupNames:       twoGroupGroupNames,
 				FeatureMap:       twoGroupFeatureMap,
diff --git a/internal/database/offline/test_impl/push.go b/internal/database/offline/test_impl/push.go
index cc6720e1..56b2268e 100644
--- a/internal/database/offline/test_impl/push.go
+++ b/internal/database/offline/test_impl/push.go
@@ -72,10 +72,10 @@ func TestPush(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroySto
 
 	t.Run("push when cdc table exists", func(t *testing.T) {
 		err := store.CreateTable(ctx, offline.CreateTableOpt{
-			TableName: dbutil.OfflineStreamCdcTableName(group.ID, revision),
-			Entity:    &entity,
-			Features:  features,
-			TableType: types.TableStreamCdc,
+			TableName:  dbutil.OfflineStreamCdcTableName(group.ID, revision),
+			EntityName: entity.Name,
+			Features:   features,
+			TableType:  types.TableStreamCdc,
 		})
 		require.NoError(t, err)
 
diff --git a/internal/database/offline/test_impl/snapshot.go b/internal/database/offline/test_impl/snapshot.go
index 4d056371..5b0cddd7 100644
--- a/internal/database/offline/test_impl/snapshot.go
+++ b/internal/database/offline/test_impl/snapshot.go
@@ -47,7 +47,7 @@ func TestSnapshot(t *testing.T, prepareStore PrepareStoreFn, destroyStore Destro
 	})
 
 	err := store.Snapshot(ctx, offline.SnapshotOpt{
-		Group:        group,
+		Group:        *group,
 		Features:     features,
 		Revision:     2,
 		PrevRevision: 1,
diff --git a/internal/database/offline/test_impl/util.go b/internal/database/offline/test_impl/util.go
index 47396ed8..d2747613 100644
--- a/internal/database/offline/test_impl/util.go
+++ b/internal/database/offline/test_impl/util.go
@@ -28,7 +28,7 @@ func buildTestSnapshotTable(
 		header = append(header, f.Name)
 	}
 	opt := offline.ImportOpt{
-		Entity:            entity,
+		EntityName:        entity.Name,
 		SnapshotTableName: snapshotTable,
 		Features:          features,
 		Header:            header,