Repository: oom-ai/oomstore
Pull request #972 (2022-01-14 03:21:55)

Jinghan/refactor Export return type
This PR changes the return type of offline method `Export` to 
```
Export(ctx context.Context, opt ExportOpt) (*types.ExportResult, error)
type ExportResult struct {
	Header []string
	Data   <-chan ExportRecord
	error  <-chan error
}
```

So that we generate `Header` and `Data` in the same place.

related to #962 


Patch:
diff --git a/internal/database/offline/bigquery/export.go b/internal/database/offline/bigquery/export.go
index f10fa480..8f1a0c4f 100644
--- a/internal/database/offline/bigquery/export.go
+++ b/internal/database/offline/bigquery/export.go
@@ -67,7 +67,7 @@ func bigqueryQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt off
 	return stream, errs
 }
 
-func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
+func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (*types.ExportResult, error) {
 	//TODO implement me
 	panic("implement me")
 }
diff --git a/internal/database/offline/mock_offline/store.go b/internal/database/offline/mock_offline/store.go
index 580d286d..760b4781 100644
--- a/internal/database/offline/mock_offline/store.go
+++ b/internal/database/offline/mock_offline/store.go
@@ -65,11 +65,11 @@ func (mr *MockStoreMockRecorder) CreateTable(ctx, opt interface{}) *gomock.Call
 }
 
 // Export mocks base method.
-func (m *MockStore) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
+func (m *MockStore) Export(ctx context.Context, opt offline.ExportOpt) (*types.ExportResult, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "Export", ctx, opt)
-	ret0, _ := ret[0].(<-chan types.ExportRecord)
-	ret1, _ := ret[1].(<-chan error)
+	ret0, _ := ret[0].(*types.ExportResult)
+	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
diff --git a/internal/database/offline/mysql/store.go b/internal/database/offline/mysql/store.go
index ffb69792..03dd2688 100644
--- a/internal/database/offline/mysql/store.go
+++ b/internal/database/offline/mysql/store.go
@@ -41,7 +41,7 @@ func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt)
 	return sqlutil.ExportOneGroup(ctx, db.DB, opt, Backend)
 }
 
-func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
+func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (*types.ExportResult, error) {
 	return sqlutil.Export(ctx, db.DB, opt, Backend)
 }
 
diff --git a/internal/database/offline/postgres/store.go b/internal/database/offline/postgres/store.go
index 6ae9ffb9..59b6a03a 100644
--- a/internal/database/offline/postgres/store.go
+++ b/internal/database/offline/postgres/store.go
@@ -38,7 +38,7 @@ func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt)
 	return sqlutil.ExportOneGroup(ctx, db.DB, opt, Backend)
 }
 
-func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
+func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (*types.ExportResult, error) {
 	return sqlutil.Export(ctx, db.DB, opt, Backend)
 
 }
diff --git a/internal/database/offline/redshift/store.go b/internal/database/offline/redshift/store.go
index 66f3f3b6..453960fb 100644
--- a/internal/database/offline/redshift/store.go
+++ b/internal/database/offline/redshift/store.go
@@ -41,7 +41,7 @@ func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt)
 	return sqlutil.ExportOneGroup(ctx, db.DB, opt, Backend)
 }
 
-func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
+func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (*types.ExportResult, error) {
 	return sqlutil.Export(ctx, db.DB, opt, Backend)
 }
 
diff --git a/internal/database/offline/snowflake/store.go b/internal/database/offline/snowflake/store.go
index b89f054c..0e0aebcc 100644
--- a/internal/database/offline/snowflake/store.go
+++ b/internal/database/offline/snowflake/store.go
@@ -55,7 +55,7 @@ func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt)
 	return sqlutil.ExportOneGroup(ctx, db.DB, opt, Backend)
 }
 
-func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
+func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (*types.ExportResult, error) {
 	return sqlutil.Export(ctx, db.DB, opt, Backend)
 }
 
diff --git a/internal/database/offline/sqlite/store.go b/internal/database/offline/sqlite/store.go
index 6d55a7fa..68a23063 100644
--- a/internal/database/offline/sqlite/store.go
+++ b/internal/database/offline/sqlite/store.go
@@ -41,7 +41,7 @@ func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt)
 	return sqlutil.ExportOneGroup(ctx, db.DB, opt, Backend)
 }
 
-func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
+func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (*types.ExportResult, error) {
 	return sqlutil.Export(ctx, db.DB, opt, Backend)
 }
 
diff --git a/internal/database/offline/sqlutil/export.go b/internal/database/offline/sqlutil/export.go
index 7d5b44af..686f7c90 100644
--- a/internal/database/offline/sqlutil/export.go
+++ b/internal/database/offline/sqlutil/export.go
@@ -55,7 +55,7 @@ func DoExportOneGroup(ctx context.Context, dbOpt dbutil.DBOpt, opt DoExportOneGr
 	return opt.QueryResults(ctx, dbOpt, opt.ExportOneGroupOpt, query, args)
 }
 
-func Export(ctx context.Context, db *sqlx.DB, opt offline.ExportOpt, backend types.BackendType) (<-chan types.ExportRecord, <-chan error) {
+func Export(ctx context.Context, db *sqlx.DB, opt offline.ExportOpt, backend types.BackendType) (*types.ExportResult, error) {
 	dbOpt := dbutil.DBOpt{
 		Backend: backend,
 		SqlxDB:  db,
@@ -67,7 +67,7 @@ func Export(ctx context.Context, db *sqlx.DB, opt offline.ExportOpt, backend typ
 	return DoExport(ctx, dbOpt, doJoinOpt)
 }
 
-func DoExport(ctx context.Context, dbOpt dbutil.DBOpt, opt DoExportOpt) (<-chan types.ExportRecord, <-chan error) {
+func DoExport(ctx context.Context, dbOpt dbutil.DBOpt, opt DoExportOpt) (*types.ExportResult, error) {
 	var (
 		emptyStream = make(chan types.ExportRecord)
 		errs        = make(chan error, 1) // at most 1 error
@@ -97,8 +97,7 @@ func DoExport(ctx context.Context, dbOpt dbutil.DBOpt, opt DoExportOpt) (<-chan
 	// Step 1: prepare export_entity table, which contains all entity keys from source tables
 	tableName, err := prepareEntityTable(ctx, dbOpt, opt.ExportOpt, snapshotTables, cdcTables)
 	if err != nil {
-		errs <- errdefs.WithStack(err)
-		return emptyStream, errs
+		return nil, err
 	}
 
 	// Step 2: join export_entity table, snapshot tables and cdc tables
@@ -133,17 +132,16 @@ func DoExport(ctx context.Context, dbOpt dbutil.DBOpt, opt DoExportOpt) (<-chan
 		Backend:         dbOpt.Backend,
 	})
 	if err != nil {
-		errs <- errdefs.WithStack(err)
-		return emptyStream, errs
+		return nil, err
 	}
 	args := make([]interface{}, 0, len(opt.CdcTables)*2)
 	for i := 0; i < len(opt.CdcTables)*2; i++ {
 		args = append(args, opt.UnixMilli)
 	}
-	return queryExportResults(ctx, dbOpt, query, args, featureList)
+	return queryExportResults(ctx, dbOpt, opt.ExportOpt, query, args, featureList)
 }
 
-func queryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, query string, args []interface{}, featureList types.FeatureList) (<-chan types.ExportRecord, <-chan error) {
+func queryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.ExportOpt, query string, args []interface{}, features types.FeatureList) (*types.ExportResult, error) {
 	stream := make(chan types.ExportRecord)
 	errs := make(chan error, 1) // at most 1 error
 
@@ -169,7 +167,7 @@ func queryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, query string, a
 				return
 			}
 			record[0] = cast.ToString(record[0])
-			for i, f := range featureList {
+			for i, f := range features {
 				if record[i+1] == nil {
 					continue
 				}
@@ -183,8 +181,8 @@ func queryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, query string, a
 			stream <- record
 		}
 	}()
-
-	return stream, errs
+	header := append([]string{opt.EntityName}, features.FullNames()...)
+	return types.NewExportResult(header, stream, errs), nil
 }
 
 func sqlxQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.ExportOneGroupOpt, query string, args []interface{}) (<-chan types.ExportRecord, <-chan error) {
diff --git a/internal/database/offline/store.go b/internal/database/offline/store.go
index 2a376541..8b1eea5f 100644
--- a/internal/database/offline/store.go
+++ b/internal/database/offline/store.go
@@ -10,7 +10,7 @@ import (
 type Store interface {
 	Join(ctx context.Context, opt JoinOpt) (*types.JoinResult, error)
 	ExportOneGroup(ctx context.Context, opt ExportOneGroupOpt) (<-chan types.ExportRecord, <-chan error)
-	Export(ctx context.Context, opt ExportOpt) (<-chan types.ExportRecord, <-chan error)
+	Export(ctx context.Context, opt ExportOpt) (*types.ExportResult, error)
 	Import(ctx context.Context, opt ImportOpt) (int64, error)
 	Push(ctx context.Context, opt PushOpt) error
 
diff --git a/pkg/oomstore/export.go b/pkg/oomstore/export.go
index 7f9a251c..94ee16dd 100644
--- a/pkg/oomstore/export.go
+++ b/pkg/oomstore/export.go
@@ -181,7 +181,7 @@ func (s *OomStore) ChannelExport(ctx context.Context, opt types.ChannelExportOpt
 		}
 	}
 
-	stream, exportErr := s.offline.Export(ctx, offline.ExportOpt{
+	result, err := s.offline.Export(ctx, offline.ExportOpt{
 		SnapshotTables: snapshotTables,
 		CdcTables:      cdcTables,
 		Features:       featureMap,
@@ -189,9 +189,10 @@ func (s *OomStore) ChannelExport(ctx context.Context, opt types.ChannelExportOpt
 		EntityName:     features[0].Group.Entity.Name,
 		Limit:          opt.Limit,
 	})
-
-	header := append([]string{features[0].Group.Entity.Name}, features.Names()...)
-	return types.NewExportResult(header, stream, exportErr), nil
+	if err != nil {
+		return nil, err
+	}
+	return result, nil
 }
 
 // key: group_Id, value: slice of features


Test patch:
diff --git a/internal/database/offline/test_impl/export.go b/internal/database/offline/test_impl/export.go
index c7667482..e9e6c176 100644
--- a/internal/database/offline/test_impl/export.go
+++ b/internal/database/offline/test_impl/export.go
@@ -92,19 +92,20 @@ func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			actual, errs := store.Export(ctx, tc.opt)
+			result, err := store.Export(ctx, tc.opt)
 			values := make([][]interface{}, 0)
-			for row := range actual {
+			for row := range result.Data {
 				values = append(values, row)
 			}
 			if tc.expectedError != nil {
-				assert.EqualError(t, <-errs, tc.expectedError.Error())
+				assert.EqualError(t, err, tc.expectedError.Error())
 			} else {
 				fmt.Println("expected: ", tc.expected)
 				fmt.Println("actual: ", values)
 
 				assert.ElementsMatch(t, tc.expected, values)
-				assert.NoError(t, <-errs)
+				assert.NoError(t, err)
+				assert.NoError(t, result.CheckStreamError())
 			}
 		})
 	}