feat(database/offline): implement offline store method Snapshot

oom-ai · Dec 30, 2021 · d3c78be · d3c78be
1 parent eaf4448
commit d3c78be
Show file tree

Hide file tree

Showing 8 changed files with 59 additions and 0 deletions.
diff --git a/internal/database/offline/bigquery/store.go b/internal/database/offline/bigquery/store.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/oom-ai/oomstore/internal/database/offline/sqlutil"
+
 	"github.com/spf13/cast"
 	"google.golang.org/api/iterator"
 
@@ -68,3 +70,12 @@ func (db *DB) TableSchema(ctx context.Context, tableName string) (*types.DataTab
 
 	return &schema, nil
 }
+
+func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
+	dbOpt := dbutil.DBOpt{
+		Backend:    types.BackendBigQuery,
+		BigQueryDB: db.Client,
+		DatasetID:  &db.datasetID,
+	}
+	return sqlutil.Snapshot(ctx, dbOpt, opt)
+}
diff --git a/internal/database/offline/mysql/store.go b/internal/database/offline/mysql/store.go
@@ -43,3 +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")
 }
+
+func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
+	dbOpt := dbutil.DBOpt{
+		Backend: types.BackendMySQL,
+		SqlxDB:  db.DB,
+	}
+	return sqlutil.Snapshot(ctx, dbOpt, opt)
+}
diff --git a/internal/database/offline/postgres/store.go b/internal/database/offline/postgres/store.go
@@ -46,3 +46,11 @@ func (db *DB) TableSchema(ctx context.Context, tableName string) (*types.DataTab
 	}
 	return sqlutil.SqlxTableSchema(ctx, db, BackendType, rows)
 }
+
+func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
+	dbOpt := dbutil.DBOpt{
+		Backend: types.BackendPostgres,
+		SqlxDB:  db.DB,
+	}
+	return sqlutil.Snapshot(ctx, dbOpt, opt)
+}
diff --git a/internal/database/offline/redshift/store.go b/internal/database/offline/redshift/store.go
@@ -48,3 +48,11 @@ func (db *DB) TableSchema(ctx context.Context, tableName string) (*types.DataTab
 	}
 	return sqlutil.SqlxTableSchema(ctx, db, BackendType, rows)
 }
+
+func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
+	dbOpt := dbutil.DBOpt{
+		Backend: types.BackendRedshift,
+		SqlxDB:  db.DB,
+	}
+	return sqlutil.Snapshot(ctx, dbOpt, opt)
+}
diff --git a/internal/database/offline/snowflake/store.go b/internal/database/offline/snowflake/store.go
@@ -58,3 +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")
 }
+
+func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
+	dbOpt := dbutil.DBOpt{
+		Backend: types.BackendSnowflake,
+		SqlxDB:  db.DB,
+	}
+	return sqlutil.Snapshot(ctx, dbOpt, opt)
+}
diff --git a/internal/database/offline/sqlite/store.go b/internal/database/offline/sqlite/store.go
@@ -44,3 +44,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")
 }
+
+func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
+	dbOpt := dbutil.DBOpt{
+		Backend: types.BackendSQLite,
+		SqlxDB:  db.DB,
+	}
+	return sqlutil.Snapshot(ctx, dbOpt, opt)
+}
diff --git a/internal/database/offline/store.go b/internal/database/offline/store.go
@@ -13,6 +13,7 @@ type Store interface {
 	Import(ctx context.Context, opt ImportOpt) (int64, error)
 
 	TableSchema(ctx context.Context, tableName string) (*types.DataTableSchema, error)
+	Snapshot(ctx context.Context, opt SnapshotOpt) error
 
 	Ping(ctx context.Context) error
 	io.Closer

diff --git a/internal/database/offline/types.go b/internal/database/offline/types.go
@@ -51,3 +51,10 @@ type CSVSource struct {
 	Reader    *bufio.Reader
 	Delimiter string
 }
+
+type SnapshotOpt struct {
+	Group          *types.Group
+	Features       types.FeatureList
+	RevisionID     int
+	PrevRevisionID int
+}