perf: Implement single storage and retrieval of schemas

.github/workflows/benchmark.yaml

@@ -4,7 +4,7 @@ on: # yamllint disable-line rule:truthy
   push:
     branches:
       - "main"
-  # NOTE! Do NOT add any other "on", because this workflow has permission to write to the repo!
+# NOTE! Do NOT add any other "on", because this workflow has permission to write to the repo!
 
 permissions:
   # permission to update benchmark contents in gh-pages branch

internal/caveats/run_test.go

@@ -467,10 +467,10 @@ func TestRunCaveatExpressions(t *testing.T) {
 					third
 				}
 				`, nil, req)
-			headRevision, err := ds.HeadRevision(t.Context())
+			headRevision, _, err := ds.HeadRevision(t.Context())
 			req.NoError(err)
 
-			reader := ds.SnapshotReader(headRevision)
+			reader := ds.SnapshotReader(headRevision, datastore.NoSchemaHashForTesting)
 
 			for _, debugOption := range []RunCaveatExpressionDebugOption{
 				RunCaveatExpressionNoDebugging,
@@ -519,10 +519,10 @@ func TestRunCaveatWithMissingMap(t *testing.T) {
 				}
 				`, nil, req)
 
-	headRevision, err := ds.HeadRevision(t.Context())
+	headRevision, _, err := ds.HeadRevision(t.Context())
 	req.NoError(err)
 
-	reader := ds.SnapshotReader(headRevision)
+	reader := ds.SnapshotReader(headRevision, datastore.NoSchemaHashForTesting)
 
 	result, err := RunSingleCaveatExpression(
 		t.Context(),
@@ -549,10 +549,10 @@ func TestRunCaveatWithEmptyMap(t *testing.T) {
 				}
 				`, nil, req)
 
-	headRevision, err := ds.HeadRevision(t.Context())
+	headRevision, _, err := ds.HeadRevision(t.Context())
 	req.NoError(err)
 
-	reader := ds.SnapshotReader(headRevision)
+	reader := ds.SnapshotReader(headRevision, datastore.NoSchemaHashForTesting)
 
 	_, err = RunSingleCaveatExpression(
 		t.Context(),
@@ -585,10 +585,10 @@ func TestRunCaveatMultipleTimes(t *testing.T) {
 				}
 				`, nil, req)
 
-	headRevision, err := ds.HeadRevision(t.Context())
+	headRevision, _, err := ds.HeadRevision(t.Context())
 	req.NoError(err)
 
-	reader := ds.SnapshotReader(headRevision)
+	reader := ds.SnapshotReader(headRevision, datastore.NoSchemaHashForTesting)
 	runner := NewCaveatRunner(types.Default.TypeSet)
 
 	// Run the first caveat.
@@ -646,10 +646,10 @@ func TestRunCaveatWithMissingDefinition(t *testing.T) {
 		}
 		`, nil, req)
 
-	headRevision, err := ds.HeadRevision(t.Context())
+	headRevision, _, err := ds.HeadRevision(t.Context())
 	req.NoError(err)
 
-	reader := ds.SnapshotReader(headRevision)
+	reader := ds.SnapshotReader(headRevision, datastore.NoSchemaHashForTesting)
 
 	// Try to run a caveat that doesn't exist
 	_, err = RunSingleCaveatExpression(
@@ -679,10 +679,10 @@ func TestCaveatRunnerPopulateCaveatDefinitionsForExpr(t *testing.T) {
 		}
 		`, nil, req)
 
-	headRevision, err := ds.HeadRevision(t.Context())
+	headRevision, _, err := ds.HeadRevision(t.Context())
 	req.NoError(err)
 
-	reader := ds.SnapshotReader(headRevision)
+	reader := ds.SnapshotReader(headRevision, datastore.NoSchemaHashForTesting)
 	runner := NewCaveatRunner(types.Default.TypeSet)
 
 	// Test populating definitions for complex expression
@@ -721,10 +721,10 @@ func TestCaveatRunnerEmptyExpression(t *testing.T) {
 		}
 		`, nil, req)
 
-	headRevision, err := ds.HeadRevision(t.Context())
+	headRevision, _, err := ds.HeadRevision(t.Context())
 	req.NoError(err)
 
-	reader := ds.SnapshotReader(headRevision)
+	reader := ds.SnapshotReader(headRevision, datastore.NoSchemaHashForTesting)
 	runner := NewCaveatRunner(types.Default.TypeSet)
 
 	// Test with an expression that has no caveats (empty operation)
@@ -799,10 +799,10 @@ func TestUnknownCaveatOperation(t *testing.T) {
 		}
 		`, nil, req)
 
-	headRevision, err := ds.HeadRevision(t.Context())
+	headRevision, _, err := ds.HeadRevision(t.Context())
 	req.NoError(err)
 
-	reader := ds.SnapshotReader(headRevision)
+	reader := ds.SnapshotReader(headRevision, datastore.NoSchemaHashForTesting)
 	runner := NewCaveatRunner(types.Default.TypeSet)
 
 	// Create an expression with an unknown operation

internal/datastore/benchmark/driver_bench_test.go

@@ -95,7 +95,7 @@ func BenchmarkDatastoreDriver(b *testing.B) {
 			// Sleep to give the datastore time to stabilize after all the writes
 			time.Sleep(1 * time.Second)
 
-			headRev, err := ds.HeadRevision(ctx)
+			headRev, _, err := ds.HeadRevision(ctx)
 			require.NoError(b, err)
 
 			b.StartTimer()
@@ -104,7 +104,7 @@ func BenchmarkDatastoreDriver(b *testing.B) {
 				b.Run("SnapshotRead", func(b *testing.B) {
 					for n := 0; n < b.N; n++ {
 						randDocNum := rand.Intn(numDocuments) //nolint:gosec
-						iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+						iter, err := ds.SnapshotReader(headRev, datastore.NoSchemaHashForTesting).QueryRelationships(ctx, datastore.RelationshipsFilter{
 							OptionalResourceType:     testfixtures.DocumentNS.Name,
 							OptionalResourceIds:      []string{strconv.Itoa(randDocNum)},
 							OptionalResourceRelation: "viewer",
@@ -120,7 +120,7 @@ func BenchmarkDatastoreDriver(b *testing.B) {
 				})
 				b.Run("SnapshotReadOnlyNamespace", func(b *testing.B) {
 					for n := 0; n < b.N; n++ {
-						iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+						iter, err := ds.SnapshotReader(headRev, datastore.NoSchemaHashForTesting).QueryRelationships(ctx, datastore.RelationshipsFilter{
 							OptionalResourceType: testfixtures.DocumentNS.Name,
 						}, options.WithQueryShape(queryshape.FindResourceOfType))
 						require.NoError(b, err)
@@ -136,7 +136,7 @@ func BenchmarkDatastoreDriver(b *testing.B) {
 						order := order
 						b.Run(orderName, func(b *testing.B) {
 							for n := 0; n < b.N; n++ {
-								iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+								iter, err := ds.SnapshotReader(headRev, datastore.NoSchemaHashForTesting).QueryRelationships(ctx, datastore.RelationshipsFilter{
 									OptionalResourceType: testfixtures.DocumentNS.Name,
 								}, options.WithSort(order), options.WithQueryShape(queryshape.FindResourceOfType))
 								require.NoError(b, err)
@@ -154,7 +154,7 @@ func BenchmarkDatastoreDriver(b *testing.B) {
 						order := order
 						b.Run(orderName, func(b *testing.B) {
 							for n := 0; n < b.N; n++ {
-								iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+								iter, err := ds.SnapshotReader(headRev, datastore.NoSchemaHashForTesting).QueryRelationships(ctx, datastore.RelationshipsFilter{
 									OptionalResourceType:     testfixtures.DocumentNS.Name,
 									OptionalResourceRelation: "viewer",
 								}, options.WithSort(order), options.WithQueryShape(queryshape.Varying))
@@ -174,7 +174,7 @@ func BenchmarkDatastoreDriver(b *testing.B) {
 						b.Run(orderName, func(b *testing.B) {
 							for n := 0; n < b.N; n++ {
 								randDocNum := rand.Intn(numDocuments) //nolint:gosec
-								iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+								iter, err := ds.SnapshotReader(headRev, datastore.NoSchemaHashForTesting).QueryRelationships(ctx, datastore.RelationshipsFilter{
 									OptionalResourceType:     testfixtures.DocumentNS.Name,
 									OptionalResourceIds:      []string{strconv.Itoa(randDocNum)},
 									OptionalResourceRelation: "viewer",
@@ -191,7 +191,7 @@ func BenchmarkDatastoreDriver(b *testing.B) {
 				})
 				b.Run("SnapshotReverseRead", func(b *testing.B) {
 					for n := 0; n < b.N; n++ {
-						iter, err := ds.SnapshotReader(headRev).ReverseQueryRelationships(ctx, datastore.SubjectsFilter{
+						iter, err := ds.SnapshotReader(headRev, datastore.NoSchemaHashForTesting).ReverseQueryRelationships(ctx, datastore.SubjectsFilter{
 							SubjectType: testfixtures.UserNS.Name,
 						}, options.WithSortForReverse(options.ByResource), options.WithQueryShapeForReverse(queryshape.Varying))
 						require.NoError(b, err)

internal/datastore/common/chunkbytes.go

@@ -81,48 +81,70 @@
 	AliveValue T
 }
 
+// WithExecutor returns a copy of the config with the specified executor.
+func (c SQLByteChunkerConfig[T]) WithExecutor(executor ChunkedBytesExecutor) SQLByteChunkerConfig[T] {
+	c.Executor = executor
+	return c
+}
+
+// WithTableName returns a copy of the config with the specified table name.
+func (c SQLByteChunkerConfig[T]) WithTableName(tableName string) SQLByteChunkerConfig[T] {
+	c.TableName = tableName
+	return c
+}
+
 // SQLByteChunker provides methods for reading and writing byte data
 // that is chunked across multiple rows in a SQL table.
 type SQLByteChunker[T any] struct {
 	config SQLByteChunkerConfig[T]
 }
 
-// MustNewSQLByteChunker creates a new SQLByteChunker with the specified configuration.
-// Panics if the configuration is invalid.
-func MustNewSQLByteChunker[T any](config SQLByteChunkerConfig[T]) *SQLByteChunker[T] {
+// NewSQLByteChunker creates a new SQLByteChunker with the specified configuration.
+// Returns an error if the configuration is invalid.
+func NewSQLByteChunker[T any](config SQLByteChunkerConfig[T]) (*SQLByteChunker[T], error) {
 	if config.MaxChunkSize <= 0 {
-		panic("maxChunkSize must be greater than 0")
+		return nil, errors.New("maxChunkSize must be greater than 0")
 	}
 	if config.TableName == "" {
-		panic("tableName cannot be empty")
+		return nil, errors.New("tableName cannot be empty")
 	}
 	if config.NameColumn == "" {
-		panic("nameColumn cannot be empty")
+		return nil, errors.New("nameColumn cannot be empty")
 	}
 	if config.ChunkIndexColumn == "" {
-		panic("chunkIndexColumn cannot be empty")
+		return nil, errors.New("chunkIndexColumn cannot be empty")
 	}
 	if config.ChunkDataColumn == "" {
-		panic("chunkDataColumn cannot be empty")
+		return nil, errors.New("chunkDataColumn cannot be empty")
 	}
 	if config.PlaceholderFormat == nil {
-		panic("placeholderFormat cannot be nil")
+		return nil, errors.New("placeholderFormat cannot be nil")
 	}
 	if config.Executor == nil {
-		panic("executor cannot be nil")
+		return nil, errors.New("executor cannot be nil")
 	}
 	if config.WriteMode == WriteModeInsertWithTombstones {
 		if config.CreatedAtColumn == "" {
-			panic("createdAtColumn is required when using WriteModeInsertWithTombstones")
+			return nil, errors.New("createdAtColumn is required when using WriteModeInsertWithTombstones")
 		}
 		if config.DeletedAtColumn == "" {
-			panic("deletedAtColumn is required when using WriteModeInsertWithTombstones")
+			return nil, errors.New("deletedAtColumn is required when using WriteModeInsertWithTombstones")
 		}
 	}
 
 	return &SQLByteChunker[T]{
 		config: config,
+	}, nil
+}
+
+// MustNewSQLByteChunker creates a new SQLByteChunker with the specified configuration.
+// Panics if the configuration is invalid.
+func MustNewSQLByteChunker[T any](config SQLByteChunkerConfig[T]) *SQLByteChunker[T] {
+	chunker, err := NewSQLByteChunker(config)
+	if err != nil {
+		panic(err)
 	}
+	return chunker
 }
 
 // WriteChunkedBytes writes chunked byte data to the database within a transaction.

internal/datastore/common/chunkbytes_test.go

@@ -57,13 +57,17 @@ type fakeExecutor struct {
 	readResult  map[int][]byte
 	readErr     error
 	transaction *fakeTransaction
+	onRead      func() // Optional callback invoked on each read
 }
 
 func (m *fakeExecutor) BeginTransaction(ctx context.Context) (ChunkedBytesTransaction, error) {
 	return m.transaction, nil
 }
 
 func (m *fakeExecutor) ExecuteRead(ctx context.Context, builder sq.SelectBuilder) (map[int][]byte, error) {
+	if m.onRead != nil {
+		m.onRead()
+	}
 	if m.readErr != nil {
 		return nil, m.readErr
 	}