Repository: oom-ai/oomstore
Pull request #1175 (2022-02-10 09:25:37)

Jinghan/aggregate params of `BuildTableSchema`
This PR modifies helper function `BuildTableSchema`: aggregate its parameters to `BuildTableSchemaParams`.


Patch:
diff --git a/internal/database/dbutil/schema.go b/internal/database/dbutil/schema.go
index 8ff1bebf..ce9124e9 100644
--- a/internal/database/dbutil/schema.go
+++ b/internal/database/dbutil/schema.go
@@ -23,16 +23,18 @@ func (c ColumnList) Names() []string {
 	return names
 }
 
-func BuildTableSchema(
-	tableName string,
-	entityName string,
-	hasUnixMilli bool,
-	features types.FeatureList,
-	pkFields []string,
-	backend types.BackendType,
-) string {
-	columns := parseColumns(entityName, hasUnixMilli, features, backend)
-	return createTableDDL(tableName, columns, pkFields, backend)
+type BuildTableSchemaParams struct {
+	TableName    string
+	EntityName   string
+	HasUnixMilli bool
+	Features     types.FeatureList
+	PrimaryKeys  []string
+	Backend      types.BackendType
+}
+
+func BuildTableSchema(params BuildTableSchemaParams) string {
+	columns := parseColumns(params.EntityName, params.HasUnixMilli, params.Features, params.Backend)
+	return createTableDDL(params.TableName, columns, params.PrimaryKeys, params.Backend)
 }
 
 func createTableDDL(tableName string, columns ColumnList, pkFields []string, backend types.BackendType) string {
diff --git a/internal/database/offline/sqlutil/create_table.go b/internal/database/offline/sqlutil/create_table.go
index 20358bc2..3496246f 100644
--- a/internal/database/offline/sqlutil/create_table.go
+++ b/internal/database/offline/sqlutil/create_table.go
@@ -13,23 +13,33 @@ func CreateTable(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.CreateTabl
 	if !supportIndex(dbOpt.Backend) {
 		return createTableNoIndex(ctx, dbOpt, opt)
 	}
+	params := dbutil.BuildTableSchemaParams{
+		TableName:  opt.TableName,
+		EntityName: opt.EntityName,
+		Features:   opt.Features,
+		Backend:    dbOpt.Backend,
+	}
 	switch opt.TableType {
 	case types.TableBatchSnapshot:
 		// Create primary key (entity_key) on batch snapshot table
-		pkFields := []string{opt.EntityName}
-		schema := dbutil.BuildTableSchema(opt.TableName, opt.EntityName, false, opt.Features, pkFields, dbOpt.Backend)
+		params.PrimaryKeys = []string{opt.EntityName}
+		params.HasUnixMilli = false
+		schema := dbutil.BuildTableSchema(params)
 		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.EntityName}
-		schema := dbutil.BuildTableSchema(opt.TableName, opt.EntityName, true, opt.Features, pkFields, dbOpt.Backend)
+		params.PrimaryKeys = []string{opt.EntityName}
+		params.HasUnixMilli = true
+		schema := dbutil.BuildTableSchema(params)
 		if err := dbOpt.ExecContext(ctx, schema, nil); err != nil {
 			return err
 		}
 	case types.TableStreamCdc:
-		schema := dbutil.BuildTableSchema(opt.TableName, opt.EntityName, true, opt.Features, nil, dbOpt.Backend)
+		params.PrimaryKeys = nil
+		params.HasUnixMilli = true
+		schema := dbutil.BuildTableSchema(params)
 		if err := dbOpt.ExecContext(ctx, schema, nil); err != nil {
 			return err
 		}
@@ -58,9 +68,16 @@ func createTableNoIndex(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.Cre
 
 	tableName := opt.TableName
 	if dbOpt.Backend == types.BackendBigQuery {
-		tableName = fmt.Sprintf("%s.%s", *dbOpt.DatasetID, tableName)
+		tableName = fmt.Sprintf("%s.%s", *dbOpt.DatasetID, opt.TableName)
 	}
-	schema := dbutil.BuildTableSchema(tableName, opt.EntityName, hasUnixMilli, opt.Features, nil, dbOpt.Backend)
+	schema := dbutil.BuildTableSchema(dbutil.BuildTableSchemaParams{
+		TableName:    tableName,
+		EntityName:   opt.EntityName,
+		HasUnixMilli: hasUnixMilli,
+		Features:     opt.Features,
+		PrimaryKeys:  nil,
+		Backend:      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 b1d98a8b..f4ee2363 100644
--- a/internal/database/offline/sqlutil/import.go
+++ b/internal/database/offline/sqlutil/import.go
@@ -16,14 +16,19 @@ type LoadData func(tx *sqlx.Tx, ctx context.Context, opt dbutil.LoadDataFromSour
 
 func Import(ctx context.Context, db *sqlx.DB, opt offline.ImportOpt, loadData LoadData, backend types.BackendType) (int64, error) {
 	var revision int64
+	params := dbutil.BuildTableSchemaParams{
+		TableName:    opt.SnapshotTableName,
+		EntityName:   opt.EntityName,
+		HasUnixMilli: false,
+		Features:     opt.Features,
+		Backend:      backend,
+	}
 	err := dbutil.WithTransaction(db, ctx, func(ctx context.Context, tx *sqlx.Tx) error {
-		// create the data table
-		var pkFields []string
+		// create the data table: streaming snapshot table doesn't have primary key
 		if opt.Category == types.CategoryBatch {
-			// streaming snapshot table doesn't have primary key
-			pkFields = []string{opt.EntityName}
+			params.PrimaryKeys = []string{opt.EntityName}
 		}
-		schema := dbutil.BuildTableSchema(opt.SnapshotTableName, opt.EntityName, false, opt.Features, pkFields, backend)
+		schema := dbutil.BuildTableSchema(params)
 		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return errdefs.WithStack(err)
diff --git a/internal/database/online/sqlutil/create_table.go b/internal/database/online/sqlutil/create_table.go
index d891495d..858b0a73 100644
--- a/internal/database/online/sqlutil/create_table.go
+++ b/internal/database/online/sqlutil/create_table.go
@@ -15,7 +15,14 @@ func CreateTable(ctx context.Context, dbOpt dbutil.DBOpt, opt online.CreateTable
 		return err
 	}
 	// Step 2: create new table
-	schema := dbutil.BuildTableSchema(opt.TableName, opt.EntityName, false, opt.Features, []string{opt.EntityName}, dbOpt.Backend)
+	schema := dbutil.BuildTableSchema(dbutil.BuildTableSchemaParams{
+		TableName:    opt.TableName,
+		EntityName:   opt.EntityName,
+		HasUnixMilli: false,
+		Features:     opt.Features,
+		PrimaryKeys:  []string{opt.EntityName},
+		Backend:      dbOpt.Backend,
+	})
 	err := dbOpt.ExecContext(ctx, schema, nil)
 	if err != nil {
 		return errdefs.WithStack(err)
diff --git a/internal/database/online/sqlutil/import.go b/internal/database/online/sqlutil/import.go
index e011cd16..c29245e3 100644
--- a/internal/database/online/sqlutil/import.go
+++ b/internal/database/online/sqlutil/import.go
@@ -26,7 +26,14 @@ func Import(ctx context.Context, db *sqlx.DB, opt online.ImportOpt, backend type
 
 	return dbutil.WithTransaction(db, ctx, func(ctx context.Context, tx *sqlx.Tx) error {
 		// create the data table
-		schema := dbutil.BuildTableSchema(tableName, entity.Name, false, opt.Features, []string{entity.Name}, backend)
+		schema := dbutil.BuildTableSchema(dbutil.BuildTableSchemaParams{
+			TableName:    tableName,
+			EntityName:   entity.Name,
+			HasUnixMilli: false,
+			Features:     opt.Features,
+			PrimaryKeys:  []string{entity.Name},
+			Backend:      backend,
+		})
 		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return errdefs.WithStack(err)


Test patch:
diff --git a/internal/database/dbutil/schema_test.go b/internal/database/dbutil/schema_test.go
index e11ea001..5f487fa2 100644
--- a/internal/database/dbutil/schema_test.go
+++ b/internal/database/dbutil/schema_test.go
@@ -81,7 +81,14 @@ CREATE TABLE "user" (
 
 	for _, c := range cases {
 		t.Run(c.description, func(t *testing.T) {
-			schema := dbutil.BuildTableSchema(c.tableName, c.entity.Name, false, c.features, nil, c.backend)
+			schema := dbutil.BuildTableSchema(dbutil.BuildTableSchemaParams{
+				TableName:    c.tableName,
+				EntityName:   c.entity.Name,
+				HasUnixMilli: false,
+				Features:     c.features,
+				PrimaryKeys:  nil,
+				Backend:      c.backend,
+			})
 			require.Equal(t, strings.TrimSpace(c.want), schema)
 		})
 	}