nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: oom-ai/oomstore
Pull request #1171 (2022-02-10 07:42:54)

Jinghan/modify field of `ImportOpt`
This PR modifies `offline.ImportOpt`'s field: `NoPK` -> `Category`, make it easier to understand.


Patch:
diff --git a/internal/database/offline/sqlutil/import.go b/internal/database/offline/sqlutil/import.go
index b5f6b915..b1d98a8b 100644
--- a/internal/database/offline/sqlutil/import.go
+++ b/internal/database/offline/sqlutil/import.go
@@ -18,9 +18,10 @@ 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
-		pkFields := []string{opt.EntityName}
-		if opt.NoPK {
-			pkFields = nil
+		var pkFields []string
+		if opt.Category == types.CategoryBatch {
+			// streaming snapshot table doesn't have primary key
+			pkFields = []string{opt.EntityName}
 		}
 		schema := dbutil.BuildTableSchema(opt.SnapshotTableName, opt.EntityName, false, opt.Features, pkFields, backend)
 		_, err := tx.ExecContext(ctx, schema)
diff --git a/internal/database/offline/types.go b/internal/database/offline/types.go
index 9b04c439..cf8baaff 100644
--- a/internal/database/offline/types.go
+++ b/internal/database/offline/types.go
@@ -38,7 +38,7 @@ type ImportOpt struct {
 	Features          types.FeatureList
 	Revision          *int64
 	Source            *CSVSource
-	NoPK              bool // TODO: to import cdc data temporarily for testing
+	Category          types.Category
 }
 
 type PushOpt struct {
diff --git a/pkg/oomstore/import_batch.go b/pkg/oomstore/import_batch.go
index 7cba3ea4..07679e6b 100644
--- a/pkg/oomstore/import_batch.go
+++ b/pkg/oomstore/import_batch.go
@@ -45,6 +45,7 @@ func (s *OomStore) csvReaderImportBatch(ctx context.Context, opt *importOpt, dat
 		Revision:          opt.Revision,
 		SnapshotTableName: snapshotTableName,
 		Source:            source,
+		Category:          types.CategoryBatch,
 	})
 	if err != nil {
 		return 0, err


Test patch:
diff --git a/internal/database/offline/test_impl/import.go b/internal/database/offline/test_impl/import.go
index a01d729e..8ee33325 100644
--- a/internal/database/offline/test_impl/import.go
+++ b/internal/database/offline/test_impl/import.go
@@ -39,6 +39,7 @@ func TestImport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 `)),
 			Delimiter: ",",
 		},
+		Category: types.CategoryBatch,
 	}
 
 	t.Run("normal import call", func(t *testing.T) {
diff --git a/internal/database/offline/test_impl/util.go b/internal/database/offline/test_impl/util.go
index d2747613..8430ff28 100644
--- a/internal/database/offline/test_impl/util.go
+++ b/internal/database/offline/test_impl/util.go
@@ -34,7 +34,7 @@ func buildTestSnapshotTable(
 		Header:            header,
 		Source:            source,
 		Revision:          &revision,
-		NoPK:              true,
+		Category:          types.CategoryStream,
 	}
 	_, err := store.Import(ctx, opt)
 	require.NoError(t, err)