From d3c78be6f832f0893e3f500dca1591a1ed546f91 Mon Sep 17 00:00:00 2001
From: Jinghan Ying <jinghancc@gmail.com>
Date: Thu, 30 Dec 2021 10:50:00 +0800
Subject: [PATCH] feat(database/offline): implement offline store method
 Snapshot

---
 internal/database/offline/bigquery/store.go  | 11 +++++++++++
 internal/database/offline/mysql/store.go     |  8 ++++++++
 internal/database/offline/postgres/store.go  |  8 ++++++++
 internal/database/offline/redshift/store.go  |  8 ++++++++
 internal/database/offline/snowflake/store.go |  8 ++++++++
 internal/database/offline/sqlite/store.go    |  8 ++++++++
 internal/database/offline/store.go           |  1 +
 internal/database/offline/types.go           |  7 +++++++
 8 files changed, 59 insertions(+)

diff --git a/internal/database/offline/bigquery/store.go b/internal/database/offline/bigquery/store.go
index ee5023007..7b86bf605 100644
--- 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
index cff8171ba..824cda4ec 100644
--- 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
index 2378c3396..1104a43c8 100644
--- 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
index 2d2eb66c4..ec0581bc7 100644
--- 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
index b7a1ba22b..6b9e78c46 100644
--- 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
index d333c4f68..5ba1d3266 100644
--- 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
index b09dc275c..305ccc466 100644
--- 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
index f1450f5f5..15cb333c8 100644
--- 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
+}