Repository: oom-ai/oomstore
Pull request #871 (2021-12-30 09:48:07)

Offline: implement `TableSchema` method
This PR implements offline method `TableSchema` for all backends.


Patch:
diff --git a/internal/database/dbutil/type_mapping.go b/internal/database/dbutil/type_mapping.go
index cf8e14f8..a5f98848 100644
--- a/internal/database/dbutil/type_mapping.go
+++ b/internal/database/dbutil/type_mapping.go
@@ -57,11 +57,16 @@ func ValueType(backend types.BackendType, dbValueType string) (types.ValueType,
 	default:
 		return 0, errdefs.InvalidAttribute(fmt.Errorf("unsupported backend: %s", backend))
 	}
-
-	t, ok := mp[strings.ToLower(dbValueType)]
+	dbValueType = strings.ToLower(dbValueType)
+	i := strings.Index(dbValueType, "(")
+	if i != -1 {
+		dbValueType = dbValueType[0:i]
+	}
+	t, ok := mp[dbValueType]
 	if !ok {
 		return 0, errdefs.InvalidAttribute(fmt.Errorf("unsupported db value type: %s", dbValueType))
 	}
+
 	return t, nil
 }
 
@@ -235,6 +240,7 @@ var (
 		"bigint":   types.Int64,
 		"tinyint":  types.Int64,
 		"byteint":  types.Int64,
+		"number":   types.Int64,
 
 		"double":           types.Float64,
 		"double precision": types.Float64,
@@ -256,11 +262,32 @@ var (
 	}
 	sqliteTypeToValueType = map[string]types.ValueType{
 		"integer":   types.Int64,
-		"float":     types.Float64,
-		"blob":      types.Bytes,
-		"text":      types.String,
-		"timestamp": types.Time,
-		"datetime":  types.Time,
+		"int":       types.Int64,
+		"tinyint":   types.Int64,
+		"smallint":  types.Int64,
+		"bigint":    types.Int64,
+		"mediumint": types.Int64,
+
+		"float":            types.Float64,
+		"real":             types.Float64,
+		"double":           types.Float64,
+		"double precision": types.Float64,
+		"numeric":          types.Float64,
+		"decimal":          types.Float64,
+
+		"blob": types.Bytes,
+
+		"text":              types.String,
+		"clob":              types.String,
+		"character varying": types.String,
+		"varchar":           types.String,
+		"nvarchar":          types.String,
+		"nchar":             types.String,
+		"character":         types.String,
+		"native character":  types.String,
+
+		"boolean":  types.Bool,
+		"datetime": types.Time,
 	}
 	// Redshift data type: https://docs.aws.amazon.com/redshift/latest/dg/c_Supported_data_types.html
 	redshiftTypeToValueType = map[string]types.ValueType{
diff --git a/internal/database/offline/mysql/store.go b/internal/database/offline/mysql/store.go
index b0ed103f..58509f09 100644
--- a/internal/database/offline/mysql/store.go
+++ b/internal/database/offline/mysql/store.go
@@ -2,7 +2,6 @@ package mysql
 
 import (
 	"context"
-	"fmt"
 
 	"github.com/jmoiron/sqlx"
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
@@ -44,7 +43,11 @@ func (db *DB) Join(ctx context.Context, opt offline.JoinOpt) (*types.JoinResult,
 }
 
 func (db *DB) TableSchema(ctx context.Context, tableName string) (*types.DataTableSchema, error) {
-	return nil, fmt.Errorf("not implemented")
+	rows, err := db.QueryxContext(ctx, "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = ?", tableName)
+	if err != nil {
+		return nil, err
+	}
+	return sqlutil.SqlxTableSchema(ctx, db, types.BackendMySQL, rows)
 }
 
 func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
diff --git a/internal/database/offline/snowflake/store.go b/internal/database/offline/snowflake/store.go
index 9eb1ed23..0cdce748 100644
--- a/internal/database/offline/snowflake/store.go
+++ b/internal/database/offline/snowflake/store.go
@@ -2,7 +2,6 @@ package snowflake
 
 import (
 	"context"
-	"fmt"
 
 	"github.com/jmoiron/sqlx"
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
@@ -59,7 +58,11 @@ func (db *DB) Join(ctx context.Context, opt offline.JoinOpt) (*types.JoinResult,
 }
 
 func (db *DB) TableSchema(ctx context.Context, tableName string) (*types.DataTableSchema, error) {
-	return nil, fmt.Errorf("not implemented")
+	rows, err := db.QueryxContext(ctx, "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = ?", tableName)
+	if err != nil {
+		return nil, err
+	}
+	return sqlutil.SqlxTableSchema(ctx, db, types.BackendSnowflake, rows)
 }
 
 func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
diff --git a/internal/database/offline/sqlite/store.go b/internal/database/offline/sqlite/store.go
index b3910bf9..7b526327 100644
--- a/internal/database/offline/sqlite/store.go
+++ b/internal/database/offline/sqlite/store.go
@@ -2,7 +2,6 @@ package sqlite
 
 import (
 	"context"
-	"fmt"
 
 	"github.com/jmoiron/sqlx"
 
@@ -45,7 +44,19 @@ func (db *DB) Join(ctx context.Context, opt offline.JoinOpt) (*types.JoinResult,
 }
 
 func (db *DB) TableSchema(ctx context.Context, tableName string) (*types.DataTableSchema, error) {
-	return nil, fmt.Errorf("not implemented")
+	query := `
+		SELECT
+			p.name as column_name,
+			p.type AS data_type
+		FROM sqlite_master AS m
+		LEFT OUTER JOIN pragma_table_info((m.name)) AS p
+		WHERE m.type = 'table' AND m.name = ?
+`
+	rows, err := db.QueryxContext(ctx, query, tableName)
+	if err != nil {
+		return nil, err
+	}
+	return sqlutil.SqlxTableSchema(ctx, db, types.BackendSQLite, rows)
 }
 
 func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {


Test patch:
diff --git a/internal/database/offline/bigquery/store_test.go b/internal/database/offline/bigquery/store_test.go
index 7b4fa2a9..1d5299fb 100644
--- a/internal/database/offline/bigquery/store_test.go
+++ b/internal/database/offline/bigquery/store_test.go
@@ -96,7 +96,7 @@ func TestTableSchema(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		query := fmt.Sprintf("CREATE TABLE %s.user(`user` STRING, `age` BIGINT)", DATASET_ID)
+		query := fmt.Sprintf("CREATE TABLE %s.offline_batch_1_1(`user` STRING, `age` BIGINT)", DATASET_ID)
 		if _, err = db.Query(query).Read(ctx); err != nil {
 			t.Fatal(err)
 		}
diff --git a/internal/database/offline/mysql/store_test.go b/internal/database/offline/mysql/store_test.go
index c84a2e29..21207c1a 100644
--- a/internal/database/offline/mysql/store_test.go
+++ b/internal/database/offline/mysql/store_test.go
@@ -48,3 +48,17 @@ func TestJoin(t *testing.T) {
 func TestSnapshot(t *testing.T) {
 	test_impl.TestSnapshot(t, prepareStore, runtime_mysql.DestroyStore(DATABASE))
 }
+
+func TestTableSchema(t *testing.T) {
+	test_impl.TestTableSchema(t, prepareStore, runtime_mysql.DestroyStore(DATABASE), func(ctx context.Context) {
+		opt := runtime_mysql.GetOpt(DATABASE)
+		db, err := dbutil.OpenMysqlDB(opt.Host, opt.Port, opt.User, opt.Password, opt.Database)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer db.Close()
+		if _, err := db.ExecContext(ctx, "create table `offline_batch_1_1`(`user` varchar(16), `age` smallint)"); err != nil {
+			t.Fatal(err)
+		}
+	})
+}
diff --git a/internal/database/offline/postgres/store_test.go b/internal/database/offline/postgres/store_test.go
index a7364ba0..b1ed67ec 100644
--- a/internal/database/offline/postgres/store_test.go
+++ b/internal/database/offline/postgres/store_test.go
@@ -53,7 +53,7 @@ func TestTableSchema(t *testing.T) {
 			t.Fatal(err)
 		}
 		defer db.Close()
-		if _, err := db.ExecContext(ctx, `create table "user"("user" varchar(16), "age" smallint)`); err != nil {
+		if _, err := db.ExecContext(ctx, `create table "offline_batch_1_1"("user" varchar(16), "age" smallint)`); err != nil {
 			t.Fatal(err)
 		}
 	})
diff --git a/internal/database/offline/redshift/store_test.go b/internal/database/offline/redshift/store_test.go
index 3d6ed5f8..5b9d40ea 100644
--- a/internal/database/offline/redshift/store_test.go
+++ b/internal/database/offline/redshift/store_test.go
@@ -89,3 +89,16 @@ func TestImport(t *testing.T) {
 func TestJoin(t *testing.T) {
 	test_impl.TestJoin(t, prepareStore, destroyStore(DATABASE))
 }
+
+func TestTableSchema(t *testing.T) {
+	test_impl.TestTableSchema(t, prepareStore, destroyStore(DATABASE), func(ctx context.Context) {
+		db, err := redshift.Open(getOpt(DATABASE))
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer db.Close()
+		if _, err = db.ExecContext(ctx, `create table "offline_batch_1_1"("user" varchar(16), "age" smallint)`); err != nil {
+			t.Fatal(err)
+		}
+	})
+}
diff --git a/internal/database/offline/snowflake/store_test.go b/internal/database/offline/snowflake/store_test.go
index dace253e..2cb10797 100644
--- a/internal/database/offline/snowflake/store_test.go
+++ b/internal/database/offline/snowflake/store_test.go
@@ -7,6 +7,8 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
 	"github.com/oom-ai/oomstore/internal/database/offline"
 	"github.com/oom-ai/oomstore/internal/database/offline/snowflake"
@@ -71,3 +73,33 @@ func TestJoin(t *testing.T) {
 func TestSnapshot(t *testing.T) {
 	test_impl.TestSnapshot(t, prepareStore, destroyStore(DATABASE))
 }
+
+// We don't use test_impl.TestTableSchema because snowflake cannot be
+// accessed by two different sessions.
+func TestTableSchema(t *testing.T) {
+	ctx, store := prepareStore(t)
+	defer store.Close()
+	db := store.(*snowflake.DB)
+
+	if _, err := db.ExecContext(ctx, `create table "offline_batch_1_1"("user" varchar(16), "age" smallint)`); err != nil {
+		t.Fatal(err)
+	}
+
+	actual, err := store.TableSchema(ctx, "offline_batch_1_1")
+	require.NoError(t, err)
+	require.Equal(t, 2, len(actual.Fields))
+
+	expected := types.DataTableSchema{
+		Fields: []types.DataTableFieldSchema{
+			{
+				Name:      "user",
+				ValueType: types.String,
+			},
+			{
+				Name:      "age",
+				ValueType: types.Int64,
+			},
+		},
+	}
+	require.ElementsMatch(t, expected.Fields, actual.Fields)
+}
diff --git a/internal/database/offline/sqlite/store_test.go b/internal/database/offline/sqlite/store_test.go
index 0cf6c67f..d58a8b4a 100644
--- a/internal/database/offline/sqlite/store_test.go
+++ b/internal/database/offline/sqlite/store_test.go
@@ -11,8 +11,11 @@ import (
 	"github.com/oom-ai/oomstore/internal/database/offline/test_impl"
 )
 
+var DATABASE string
+
 func prepareStore(t *testing.T) (context.Context, offline.Store) {
-	file, err := os.CreateTemp(t.TempDir(), "")
+	DATABASE = t.TempDir() + "/test.db"
+	file, err := os.Create(DATABASE)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -26,7 +29,12 @@ func prepareStore(t *testing.T) (context.Context, offline.Store) {
 	return context.Background(), &sqlite.DB{conn}
 }
 
-func destroyStore() {}
+func destroyStore() {
+	if err := os.RemoveAll(DATABASE); err != nil {
+		panic(err)
+	}
+
+}
 
 func TestPing(t *testing.T) {
 	test_impl.TestPing(t, prepareStore, destroyStore)
@@ -47,3 +55,16 @@ func TestJoin(t *testing.T) {
 func TestSnapshot(t *testing.T) {
 	test_impl.TestSnapshot(t, prepareStore, destroyStore)
 }
+
+func TestTableSchema(t *testing.T) {
+	test_impl.TestTableSchema(t, prepareStore, destroyStore, func(ctx context.Context) {
+		db, err := dbutil.OpenSQLite(DATABASE)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer db.Close()
+		if _, err := db.ExecContext(ctx, "create table `offline_batch_1_1`(`user` varchar(16), `age` smallint)"); err != nil {
+			t.Fatal(err)
+		}
+	})
+}
diff --git a/internal/database/offline/test_impl/table_schema.go b/internal/database/offline/test_impl/table_schema.go
index 55ddd6de..1fba90a9 100644
--- a/internal/database/offline/test_impl/table_schema.go
+++ b/internal/database/offline/test_impl/table_schema.go
@@ -16,7 +16,7 @@ func TestTableSchema(t *testing.T, prepareStore PrepareStoreFn, destroyStore Des
 
 	prepareSampleTable(ctx)
 
-	actual, err := store.TableSchema(ctx, "user")
+	actual, err := store.TableSchema(ctx, "offline_batch_1_1")
 	require.NoError(t, err)
 	require.Equal(t, 2, len(actual.Fields))
 
diff --git a/internal/database/offline/tidb/store_test.go b/internal/database/offline/tidb/store_test.go
index 85022bb7..241fb4f7 100644
--- a/internal/database/offline/tidb/store_test.go
+++ b/internal/database/offline/tidb/store_test.go
@@ -44,3 +44,16 @@ func TestImport(t *testing.T) {
 func TestJoin(t *testing.T) {
 	test_impl.TestJoin(t, prepareStore, runtime_tidb.DestroyStore(DATABASE))
 }
+
+func TestTableSchema(t *testing.T) {
+	test_impl.TestTableSchema(t, prepareStore, runtime_tidb.DestroyStore(DATABASE), func(ctx context.Context) {
+		db, err := mysql.Open(runtime_tidb.GetOpt(DATABASE))
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer db.Close()
+		if _, err = db.ExecContext(ctx, "create table `offline_batch_1_1`(`user` varchar(16), `age` smallint)"); err != nil {
+			t.Fatal(err)
+		}
+	})
+}