Merge remote-tracking branch 'sashayakovtseva/feature/ydb-insert-from…

…-as-table' into feature/build-ydb
sashayakovtseva · Mar 19, 2024 · 4962f73 · 4962f73
2 parents 6061ddc + 78ea937
commit 4962f73
Show file tree

Hide file tree

Showing 8 changed files with 337 additions and 250 deletions.
diff --git a/internal/datastore/benchmark/driver_bench_test.go b/internal/datastore/benchmark/driver_bench_test.go
@@ -1,5 +1,4 @@
 //go:build ci && docker
-// +build ci,docker
 
 package benchmark
 
@@ -18,6 +17,7 @@ import (
 	"github.com/authzed/spicedb/internal/datastore/mysql"
 	"github.com/authzed/spicedb/internal/datastore/postgres"
 	"github.com/authzed/spicedb/internal/datastore/spanner"
+	"github.com/authzed/spicedb/internal/datastore/ydb"
 	"github.com/authzed/spicedb/internal/testfixtures"
 	testdatastore "github.com/authzed/spicedb/internal/testserver/datastore"
 	"github.com/authzed/spicedb/internal/testserver/datastore/config"
@@ -47,6 +47,8 @@ var drivers = []struct {
 	{postgres.Engine, "", nil},
 	{crdb.Engine, "-overlap-static", []dsconfig.ConfigOption{dsconfig.WithOverlapStrategy("static")}},
 	{crdb.Engine, "-overlap-insecure", []dsconfig.ConfigOption{dsconfig.WithOverlapStrategy("insecure")}},
+	{ydb.Engine, "-uniqueness-check", []dsconfig.ConfigOption{dsconfig.WithYDBEnableUniquenessCheck(true)}},
+	{ydb.Engine, "-no-uniqueness-check", []dsconfig.ConfigOption{dsconfig.WithYDBEnableUniquenessCheck(false)}},
 	{mysql.Engine, "", nil},
 }
 

diff --git a/internal/datastore/ydb/options.go b/internal/datastore/ydb/options.go
@@ -21,8 +21,9 @@ type ydbConfig struct {
 
 	bulkLoadBatchSize int
 
-	gcEnabled             bool
+	enableGC              bool
 	enablePrometheusStats bool
+	enableUniquenessCheck bool
 }
 
 var defaultConfig = ydbConfig{
@@ -34,7 +35,8 @@ var defaultConfig = ydbConfig{
 	gcInterval:                  3 * time.Minute,
 	gcMaxOperationTime:          time.Minute,
 	bulkLoadBatchSize:           1000,
-	gcEnabled:                   true,
+	enableGC:                    true,
+	enableUniquenessCheck:       true,
 }
 
 // Option provides the facility to configure how clients within the YDB
@@ -80,11 +82,11 @@ func GCInterval(interval time.Duration) Option {
 	return func(o *ydbConfig) { o.gcInterval = interval }
 }
 
-// GCEnabled indicates whether garbage collection is enabled.
+// WithEnableGC indicates whether garbage collection is enabled.
 //
 // GC is enabled by default.
-func GCEnabled(isGCEnabled bool) Option {
-	return func(o *ydbConfig) { o.gcEnabled = isGCEnabled }
+func WithEnableGC(isGCEnabled bool) Option {
+	return func(o *ydbConfig) { o.enableGC = isGCEnabled }
 }
 
 // GCMaxOperationTime is the maximum operation time of a garbage collection pass before it times out.
@@ -143,6 +145,13 @@ func WatchBufferWriteTimeout(watchBufferWriteTimeout time.Duration) Option {
 // clients being used by the datastore are enabled.
 //
 // Prometheus metrics are disabled by default.
-func WithEnablePrometheusStats(enablePrometheusStats bool) Option {
-	return func(o *ydbConfig) { o.enablePrometheusStats = enablePrometheusStats }
+func WithEnablePrometheusStats(v bool) Option {
+	return func(o *ydbConfig) { o.enablePrometheusStats = v }
+}
+
+// WithEnableUniquenessCheck marks whether relation tuples will be checked against
+// unique index during CREATE operation. YDB doesn't support unique secondary indexes,
+// and since this check is quite expensive one may turn it off.
+func WithEnableUniquenessCheck(v bool) Option {
+	return func(o *ydbConfig) { o.enableUniquenessCheck = v }
 }
diff --git a/internal/datastore/ydb/query.go b/internal/datastore/ydb/query.go
@@ -281,7 +281,7 @@ func queryTuples(
 			var structuredCtx map[string]any
 			if caveatCtx != nil {
 				if err := json.Unmarshal(*caveatCtx, &structuredCtx); err != nil {
-					return nil, fmt.Errorf("failed to unmarhsla relation tuple caveat context: %w", err)
+					return nil, fmt.Errorf("failed to unmarshal relation tuple caveat context: %w", err)
 				}
 			}
 

diff --git a/internal/datastore/ydb/readwrite.go b/internal/datastore/ydb/readwrite.go
@@ -27,8 +27,9 @@ var _ datastore.ReadWriteTransaction = (*ydbReadWriter)(nil)
 
 type ydbReadWriter struct {
 	*ydbReader
-	bulkLoadBatchSize int
-	newRevision       revisions.TimestampRevision
+	bulkLoadBatchSize     int
+	newRevision           revisions.TimestampRevision
+	enableUniquenessCheck bool
 }
 
 // WriteCaveats stores the provided caveats.
@@ -91,7 +92,7 @@ func (rw *ydbReadWriter) WriteRelationships(ctx context.Context, mutations []*co
 	}
 
 	// Perform SELECT queries first as a part of uniqueness check.
-	if len(insertionTuples) > 0 {
+	if len(insertionTuples) > 0 && rw.enableUniquenessCheck {
 		dups, err := rw.selectTuples(ctx, insertionTuples)
 		if err != nil {
 			return fmt.Errorf("failed to ensure CREATE tuples uniqueness: %w", err)
@@ -284,12 +285,14 @@ func (rw *ydbReadWriter) BulkLoad(ctx context.Context, iter datastore.BulkWriteR
 		}
 
 		if len(insertionTuples) > 0 {
-			dups, err := rw.selectTuples(ctx, insertionTuples)
-			if err != nil {
-				return 0, fmt.Errorf("failed to ensure CREATE tuples uniqueness: %w", err)
-			}
-			if len(dups) > 0 {
-				return 0, datastoreCommon.NewCreateRelationshipExistsError(dups[0])
+			if rw.enableUniquenessCheck {
+				dups, err := rw.selectTuples(ctx, insertionTuples)
+				if err != nil {
+					return 0, fmt.Errorf("failed to ensure CREATE tuples uniqueness: %w", err)
+				}
+				if len(dups) > 0 {
+					return 0, datastoreCommon.NewCreateRelationshipExistsError(dups[0])
+				}
 			}
 
 			if err := executeQuery(ctx, rw.tablePathPrefix, rw.executor, insertBuilder); err != nil {