Repository: oom-ai/oomstore
Pull request #876 (2021-12-31 05:28:17)

feature: improve BuildCreateSchema util function 


Patch:
diff --git a/internal/database/dbutil/schema.go b/internal/database/dbutil/schema.go
index 88346620..cae4bbe1 100644
--- a/internal/database/dbutil/schema.go
+++ b/internal/database/dbutil/schema.go
@@ -9,49 +9,33 @@ import (
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
-const createSchema = `
+const tableSchemaTmpl = `
 CREATE TABLE {{ .TableName }} (
 	{{ entity .Entity .Backend }},
-	{{ columnJoin .Columns .Backend }}
-)`
-
-const snowflakeCreateSchema = `
-CREATE TABLE "{{ .TableName }}" (
-	{{ entity .Entity .Backend }},
-	{{ columnJoin .Columns .Backend }}
+	{{ fields .Fields .Backend }}
 )`
 
 // TODO: Add back `PRIMARY KEY` back when we have functions for
 // creating cdc table
 var (
-	createSchemaFuncs = template.FuncMap{
+	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`, entity.Name)
-			case types.BackendPostgres, types.BackendSnowflake:
-				return fmt.Sprintf(`"%s" VARCHAR(%d)`, entity.Name, entity.Length)
-			case types.BackendMySQL:
-				return fmt.Sprintf("`%s` VARCHAR(%d)", entity.Name, entity.Length)
+				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`, entity.Name)
+				return fmt.Sprintf(`%s STRING`, entityName)
 			default:
-				return fmt.Sprintf("%s VARCHAR(%d)", entity.Name, entity.Length)
+				panic(fmt.Sprintf("unsupported backend type %s", backend))
 			}
 		},
-		"columnJoin": func(columns []Column, backend types.BackendType) string {
-			var format string
-			switch backend {
-			case types.BackendPostgres, types.BackendSnowflake, types.BackendCassandra:
-				format = `"%s" %s`
-			case types.BackendMySQL:
-				format = "`%s` %s"
-			default:
-				format = "%s %s"
-			}
+		"fields": func(columns []Column, backend types.BackendType) string {
 			rs := make([]string, 0, len(columns))
 			for _, column := range columns {
-				rs = append(rs, fmt.Sprintf(format, column.Name, column.DbType))
+				rs = append(rs, fmt.Sprintf("%s %s", QuoteFn(backend)(column.Name), column.DbType))
 			}
 			return strings.Join(rs, ",\n\t")
 		},
@@ -74,53 +58,61 @@ func (c ColumnList) Names() []string {
 	return names
 }
 
-type CreateSchema struct {
+type TableSchemaTmplOpts struct {
 	TableName string
 	Entity    types.Entity
-	Columns   []Column
-
-	Backend types.BackendType
+	Fields    []Column
+	Backend   types.BackendType
 }
 
-func BuildCreateSchema(tableName string, entity *types.Entity, features types.FeatureList, backend types.BackendType) (string, error) {
-	var text string
-	switch backend {
-	case types.BackendSnowflake:
-		text = snowflakeCreateSchema
-	default:
-		text = createSchema
-	}
+func BuildTableSchema(
+	tableName string,
+	entity *types.Entity,
+	withUnixMillis bool,
+	features types.FeatureList,
+	backend types.BackendType,
+) string {
 	buf := bytes.NewBuffer(nil)
-	schema, err := newSchema(tableName, *entity, features, backend)
-	if err != nil {
-		return "", err
+	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)
 	}
-
-	t := template.Must(template.New("schema").Funcs(createSchemaFuncs).Parse(text))
-	if err = t.Execute(buf, schema); err != nil {
-		return "", err
-	}
-	return buf.String(), nil
+	return buf.String()
 }
 
-func newSchema(tableName string, entity types.Entity, features types.FeatureList, backend types.BackendType) (CreateSchema, error) {
-	columns := make([]Column, 0, len(features))
+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)
+		}
+		fields = append(fields, Column{
+			Name:   "unix_milli",
+			DbType: dbType,
+		})
+	}
 	for _, feature := range features {
 		dbType, err := DBValueType(backend, feature.ValueType)
 		if err != nil {
-			return CreateSchema{}, err
+			panic(err)
 		}
-
-		columns = append(columns, Column{
+		fields = append(fields, Column{
 			Name:   feature.Name,
 			DbType: dbType,
 		})
 	}
 
-	return CreateSchema{
-		TableName: tableName,
+	return &TableSchemaTmplOpts{
+		TableName: QuoteFn(backend)(tableName),
 		Entity:    entity,
-		Columns:   columns,
+		Fields:    fields,
 		Backend:   backend,
-	}, nil
+	}
 }
diff --git a/internal/database/offline/sqlutil/import.go b/internal/database/offline/sqlutil/import.go
index 16718bd1..abd7f9e3 100644
--- a/internal/database/offline/sqlutil/import.go
+++ b/internal/database/offline/sqlutil/import.go
@@ -16,11 +16,8 @@ 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
-		schema, err := dbutil.BuildCreateSchema(opt.SnapshotTableName, opt.Entity, opt.Features, backendType)
-		if err != nil {
-			return err
-		}
-		_, err = tx.ExecContext(ctx, schema)
+		schema := dbutil.BuildTableSchema(opt.SnapshotTableName, opt.Entity, false, opt.Features, backendType)
+		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return err
 		}
diff --git a/internal/database/offline/sqlutil/snapshot.go b/internal/database/offline/sqlutil/snapshot.go
index 25fba9fd..728f447c 100644
--- a/internal/database/offline/sqlutil/snapshot.go
+++ b/internal/database/offline/sqlutil/snapshot.go
@@ -69,12 +69,8 @@ func Snapshot(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.SnapshotOpt)
 		currCdcTableName = fmt.Sprintf("%s.%s", *dbOpt.DatasetID, currCdcTableName)
 	}
 
-	schema, err := dbutil.BuildCreateSchema(currSnapshotTableName, opt.Group.Entity, opt.Features, dbOpt.Backend)
-	if err != nil {
-		return err
-	}
-
-	if err = dbOpt.ExecContext(ctx, schema, nil); err != nil {
+	schema := dbutil.BuildTableSchema(currSnapshotTableName, opt.Group.Entity, false, opt.Features, dbOpt.Backend)
+	if err := dbOpt.ExecContext(ctx, schema, nil); err != nil {
 		return err
 	}
 	query, err := buildSnapshotQuery(snapshotQueryParams{
diff --git a/internal/database/online/cassandra/import.go b/internal/database/online/cassandra/import.go
index 5fa1a59d..2d3a63e5 100644
--- a/internal/database/online/cassandra/import.go
+++ b/internal/database/online/cassandra/import.go
@@ -16,10 +16,7 @@ func (db *DB) Import(ctx context.Context, opt online.ImportOpt) error {
 	columns := append([]string{opt.Entity.Name}, opt.FeatureList.Names()...)
 	tableName := sqlutil.OnlineBatchTableName(opt.Revision.ID)
 
-	table, err := dbutil.BuildCreateSchema(tableName, opt.Entity, opt.FeatureList, Backend)
-	if err != nil {
-		return err
-	}
+	table := dbutil.BuildTableSchema(tableName, opt.Entity, false, opt.FeatureList, Backend)
 
 	// create table
 	if err := db.Query(table).Exec(); err != nil {
@@ -38,7 +35,7 @@ func (db *DB) Import(ctx context.Context, opt online.ImportOpt) error {
 		if batch.Size() != BatchSize {
 			batch.Query(insertStmt, record...)
 		} else {
-			if err = db.ExecuteBatch(batch); err != nil {
+			if err := db.ExecuteBatch(batch); err != nil {
 				return err
 			}
 			batch = db.NewBatch(gocql.LoggedBatch)
diff --git a/internal/database/online/sqlutil/import.go b/internal/database/online/sqlutil/import.go
index 9590fcbc..b48e4faa 100644
--- a/internal/database/online/sqlutil/import.go
+++ b/internal/database/online/sqlutil/import.go
@@ -17,11 +17,8 @@ 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, err := dbutil.BuildCreateSchema(tableName, opt.Entity, opt.FeatureList, backend)
-		if err != nil {
-			return err
-		}
-		_, err = tx.ExecContext(ctx, schema)
+		schema := dbutil.BuildTableSchema(tableName, opt.Entity, false, opt.FeatureList, backend)
+		_, err := tx.ExecContext(ctx, schema)
 		if err != nil {
 			return err
 		}


Test patch:
diff --git a/internal/database/dbutil/schema_test.go b/internal/database/dbutil/schema_test.go
index a855beb9..62813484 100644
--- a/internal/database/dbutil/schema_test.go
+++ b/internal/database/dbutil/schema_test.go
@@ -41,7 +41,7 @@ func TestBuildSchema(t *testing.T) {
 			features: features,
 
 			want: `
-CREATE TABLE user (
+CREATE TABLE "user" (
 	"user_id" VARCHAR(32),
 	"age" bigint,
 	"gender" text
@@ -59,7 +59,7 @@ CREATE TABLE user (
 			features: features,
 
 			want: "\n" +
-				"CREATE TABLE user (\n" +
+				"CREATE TABLE `user` (\n" +
 				"	`user_id` VARCHAR(32),\n" +
 				"	`age` bigint,\n" +
 				"	`gender` text\n)",
@@ -76,7 +76,7 @@ CREATE TABLE user (
 
 			wantErr: nil,
 			want: `
-CREATE TABLE user (
+CREATE TABLE "user" (
 	"user_id" TEXT,
 	"age" bigint,
 	"gender" text
@@ -86,8 +86,7 @@ CREATE TABLE user (
 
 	for _, c := range cases {
 		t.Run(c.description, func(t *testing.T) {
-			schema, err := dbutil.BuildCreateSchema(c.tableName, c.entity, c.features, c.backend)
-			require.Equal(t, c.wantErr, err)
+			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/import.go b/internal/database/offline/test_impl/import.go
index dab66df6..87a9fc8f 100644
--- a/internal/database/offline/test_impl/import.go
+++ b/internal/database/offline/test_impl/import.go
@@ -28,17 +28,7 @@ func TestImport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 	opt := offline.ImportOpt{
 		Entity:            &entity,
 		SnapshotTableName: snapshotTable,
-		Features: []*types.Feature{
-			{
-				Name:      "price",
-				ValueType: types.Int64,
-			},
-			{
-				Name:      "model",
-				ValueType: -1,
-			},
-		},
-		Header: []string{"device", "model", "price"},
+		Header:            []string{"device", "model", "price"},
 		Source: &offline.CSVSource{
 			Reader: bufio.NewReader(strings.NewReader(`1234,xiaomi,1899
 1235,apple,4999
@@ -49,11 +39,6 @@ func TestImport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 		},
 	}
 
-	t.Run("invalid value type", func(t *testing.T) {
-		_, err := store.Import(ctx, opt)
-		assert.Error(t, err)
-	})
-
 	t.Run("normal import call", func(t *testing.T) {
 		revision := int64(1234)
 		opt.Features = []*types.Feature{