chore: small followups to the testcontainers PR

internal/datastore/postgres/pgbouncer_test.go

@@ -18,7 +18,7 @@ func pgbouncerTestVersion() string {
 	return version.LatestTestedPostgresVersion
 }
 
-var pgbouncerConfig = postgresTestConfig{"head", "", pgbouncerTestVersion(), true}
+var pgbouncerConfig = postgresTestConfig{"", pgbouncerTestVersion(), true}
 
 func TestPostgresWithPgBouncerDatastore(t *testing.T) {
 	testPostgresDatastore(t, pgbouncerConfig)

internal/datastore/postgres/postgres_shared_test.go

@@ -24,16 +24,13 @@ import (
 	"github.com/authzed/spicedb/internal/datastore/common"
 	pgcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
 	"github.com/authzed/spicedb/internal/datastore/postgres/schema"
-	pgversion "github.com/authzed/spicedb/internal/datastore/postgres/version"
 	"github.com/authzed/spicedb/internal/datastore/proxy"
 	"github.com/authzed/spicedb/internal/datastore/proxy/indexcheck"
 	"github.com/authzed/spicedb/internal/testfixtures"
 	testdatastore "github.com/authzed/spicedb/internal/testserver/datastore"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/datastore/options"
-	"github.com/authzed/spicedb/pkg/datastore/queryshape"
 	"github.com/authzed/spicedb/pkg/datastore/test"
-	"github.com/authzed/spicedb/pkg/migrate"
 	"github.com/authzed/spicedb/pkg/namespace"
 	"github.com/authzed/spicedb/pkg/tuple"
 )
@@ -47,10 +44,9 @@ const (
 var pgFactory = test.NewTesterFactory(&pgconn.PgError{Code: pgSerializationFailure})
 
 type postgresTestConfig struct {
-	targetMigration string
-	migrationPhase  string
-	pgVersion       string
-	pgbouncer       bool
+	migrationPhase string
+	pgVersion      string
+	pgbouncer      bool
 }
 
 // the global OTel tracer is used everywhere, so we synchronize tests over a global test tracer
@@ -72,8 +68,8 @@ func testPostgresDatastore(t *testing.T, config postgresTestConfig) {
 		pgbouncerStr = "pgbouncer-"
 	}
 
-	t.Run(fmt.Sprintf("%spostgres-%s-%s-%s-gc", pgbouncerStr, config.pgVersion, config.targetMigration, config.migrationPhase), func(t *testing.T) {
-		b := testdatastore.RunPostgresForTesting(t, config.targetMigration, config.pgVersion, config.pgbouncer)
+	t.Run(fmt.Sprintf("%spostgres-%s-%s-gc", pgbouncerStr, config.pgVersion, config.migrationPhase), func(t *testing.T) {
+		b := testdatastore.RunPostgresForTesting(t, config.pgVersion, config.pgbouncer)
 		ctx := t.Context()
 
 		// NOTE: gc tests take exclusive locks, so they are run under non-parallel.
@@ -109,9 +105,9 @@ func testPostgresDatastore(t *testing.T, config postgresTestConfig) {
 		))
 	})
 
-	t.Run(fmt.Sprintf("%spostgres-%s-%s-%s", pgbouncerStr, config.pgVersion, config.targetMigration, config.migrationPhase), func(t *testing.T) {
+	t.Run(fmt.Sprintf("%spostgres-%s-%s", pgbouncerStr, config.pgVersion, config.migrationPhase), func(t *testing.T) {
 		t.Parallel()
-		b := testdatastore.RunPostgresForTesting(t, config.targetMigration, config.pgVersion, config.pgbouncer)
+		b := testdatastore.RunPostgresForTesting(t, config.pgVersion, config.pgbouncer)
 		ctx := t.Context()
 
 		test.AllWithExceptions(t, pgFactory.NewTester(test.DatastoreTesterFunc(func(t testing.TB, revisionQuantization, _, gcWindow time.Duration, watchBufferLength uint16) (datastore.Datastore, error) {
@@ -320,7 +316,7 @@ func testPostgresDatastoreWithoutCommitTimestamps(t *testing.T, config postgresT
 		t.Parallel()
 
 		ctx := t.Context()
-		b := testdatastore.RunPostgresForTestingWithCommitTimestamps(t, "head", false, pgVersion, enablePgbouncer)
+		b := testdatastore.RunPostgresForTestingWithCommitTimestamps(t, false, pgVersion, enablePgbouncer)
 
 		// NOTE: watch API requires the commit timestamps, so we skip those tests here.
 		// NOTE: gc tests take exclusive locks, so they are run under non-parallel.
@@ -343,7 +339,7 @@ func testPostgresDatastoreWithoutCommitTimestamps(t *testing.T, config postgresT
 
 	t.Run(fmt.Sprintf("postgres-%s-gc", pgVersion), func(t *testing.T) {
 		ctx := t.Context()
-		b := testdatastore.RunPostgresForTestingWithCommitTimestamps(t, "head", false, pgVersion, enablePgbouncer)
+		b := testdatastore.RunPostgresForTestingWithCommitTimestamps(t, false, pgVersion, enablePgbouncer)
 		test.OnlyGCTests(t, test.DatastoreTesterFunc(func(t testing.TB, revisionQuantization, gcInterval, gcWindow time.Duration, watchBufferLength uint16) (datastore.Datastore, error) {
 			ds := b.NewDatastore(t, func(engine, uri string) datastore.Datastore {
 				ds, err := newPostgresDatastore(ctx, uri, primaryInstanceID,
@@ -1477,7 +1473,7 @@ func OTelTracingTest(t *testing.T, ds datastore.Datastore) {
 func WatchNotEnabledTest(t *testing.T, _ testdatastore.RunningEngineForTest, pgVersion string) {
 	require := require.New(t)
 
-	ds := testdatastore.RunPostgresForTestingWithCommitTimestamps(t, migrate.Head, false, pgVersion, false).NewDatastore(t, func(engine, uri string) datastore.Datastore {
+	ds := testdatastore.RunPostgresForTestingWithCommitTimestamps(t, false, pgVersion, false).NewDatastore(t, func(engine, uri string) datastore.Datastore {
 		ctx := t.Context()
 		ds, err := newPostgresDatastore(ctx, uri,
 			primaryInstanceID,
@@ -1504,49 +1500,11 @@ func WatchNotEnabledTest(t *testing.T, _ testdatastore.RunningEngineForTest, pgV
 	require.Contains(err.Error(), "track_commit_timestamp=on")
 }
 
-func BenchmarkPostgresQuery(b *testing.B) {
-	b.StopTimer()
-
-	ds := testdatastore.RunPostgresForTesting(b, migrate.Head, pgversion.MinimumSupportedPostgresVersion, false).NewDatastore(b, func(engine, uri string) datastore.Datastore {
-		ctx := b.Context()
-		ds, err := newPostgresDatastore(ctx, uri,
-			primaryInstanceID,
-			RevisionQuantization(0),
-			GCWindow(time.Millisecond*1),
-			GCInterval(veryLargeGCInterval),
-			WatchBufferLength(1),
-		)
-		require.NoError(b, err)
-		return ds
-	})
-	ds, revision := testfixtures.StandardDatastoreWithData(b, ds)
-	b.Cleanup(func() {
-		_ = ds.Close()
-	})
-
-	b.StartTimer()
-
-	b.Run("benchmark checks", func(b *testing.B) {
-		require := require.New(b)
-
-		for i := 0; i < b.N; i++ {
-			iter, err := ds.SnapshotReader(revision).QueryRelationships(b.Context(), datastore.RelationshipsFilter{
-				OptionalResourceType: testfixtures.DocumentNS.Name,
-			}, options.WithQueryShape(queryshape.FindResourceOfType))
-			require.NoError(err)
-			for rel, err := range iter {
-				require.NoError(err)
-				require.Equal(testfixtures.DocumentNS.Name, rel.Resource.ObjectType)
-			}
-		}
-	})
-}
-
 func datastoreWithInterceptorAndTestData(t *testing.T, interceptor pgcommon.QueryInterceptor, pgVersion string) datastore.Datastore {
 	t.Helper()
 	require := require.New(t)
 
-	ds := testdatastore.RunPostgresForTestingWithCommitTimestamps(t, migrate.Head, false, pgVersion, false).NewDatastore(t, func(engine, uri string) datastore.Datastore {
+	ds := testdatastore.RunPostgresForTestingWithCommitTimestamps(t, false, pgVersion, false).NewDatastore(t, func(engine, uri string) datastore.Datastore {
 		ctx := t.Context()
 		ds, err := newPostgresDatastore(ctx, uri,
 			primaryInstanceID,

internal/datastore/postgres/postgres_test.go

@@ -26,7 +26,7 @@ func postgresTestVersion() string {
 	return version.LatestTestedPostgresVersion
 }
 
-var postgresConfig = postgresTestConfig{"head", "", postgresTestVersion(), false}
+var postgresConfig = postgresTestConfig{"", postgresTestVersion(), false}
 
 func TestPostgresDatastore(t *testing.T) {
 	testPostgresDatastore(t, postgresConfig)
@@ -42,10 +42,10 @@ func TestPostgresDatastoreGC(t *testing.T) {
 	if config.pgbouncer {
 		pgbouncerStr = "pgbouncer-"
 	}
-	t.Run(fmt.Sprintf("%spostgres-gc-%s-%s-%s", pgbouncerStr, config.pgVersion, config.targetMigration, config.migrationPhase), func(t *testing.T) {
+	t.Run(fmt.Sprintf("%spostgres-gc-%s-%s", pgbouncerStr, config.pgVersion, config.migrationPhase), func(t *testing.T) {
 		t.Parallel()
 
-		b := testdatastore.RunPostgresForTesting(t, config.targetMigration, config.pgVersion, config.pgbouncer)
+		b := testdatastore.RunPostgresForTesting(t, config.pgVersion, config.pgbouncer)
 
 		t.Run("GarbageCollection", createDatastoreTest(
 			b,

internal/datastore/spanner/spanner_test.go

@@ -26,7 +26,7 @@ func TestSpannerDatastore(t *testing.T) {
 	// t.Parallel() //nolint:tparallel, the test sets environment variables (the emulator)
 
 	ctx := t.Context()
-	b := testdatastore.RunSpannerForTesting(t, "head")
+	b := testdatastore.RunSpannerForTesting(t)
 
 	// Transaction tests are excluded because, for reasons unknown, one cannot read its own write in one transaction in the Spanner emulator.
 	test.AllWithExceptions(t, spannerFactory.NewTester(test.DatastoreTesterFunc(func(_ testing.TB, revisionQuantization, _, _ time.Duration, watchBufferLength uint16) (datastore.Datastore, error) {

internal/dispatch/executor.go

@@ -629,10 +629,10 @@ func NewQueryContext(
 }
 
 // NewPlanContext builds a PlanContext proto from query.Context fields.
-func NewPlanContext(revision string, schemaHash datalayer.SchemaHash, caveatContext map[string]any, maxRecursionDepth int, datastoreLimit uint64) *v1.PlanContext {
+func NewPlanContext(revision string, schemaHash datalayer.SchemaHash, caveatContext map[string]any, maxRecursionDepth int32, datastoreLimit uint64) *v1.PlanContext {
 	pc := &v1.PlanContext{
 		Revision:               revision,
-		MaxRecursionDepth:      int32(maxRecursionDepth),
+		MaxRecursionDepth:      maxRecursionDepth,
 		OptionalDatastoreLimit: datastoreLimit,
 		SchemaHash:             []byte(schemaHash),
 	}

internal/services/v1/permissions_queryplan.go

@@ -4,6 +4,8 @@
 	"cmp"
 	"context"
 
+	"github.com/ccoveille/go-safecast/v2"
+
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
 	caveatsimpl "github.com/authzed/spicedb/internal/caveats"
@@ -99,7 +101,11 @@
 	countObserver := query.NewCountObserver()
 
 	// Create query context backed by a DispatchExecutor.
-	planContext := dispatch.NewPlanContext(atRevision.String(), schemaHash, caveatContext, int(ps.config.MaximumAPIDepth), 0)
+	maximumAPIDepth, err := safecast.Convert[int32](ps.config.MaximumAPIDepth)
+	if err != nil {
+		return nil, ps.rewriteError(ctx, err)
+	}
+	planContext := dispatch.NewPlanContext(atRevision.String(), schemaHash, caveatContext, maximumAPIDepth, 0)
 	qctx := dispatch.NewQueryContext(
 		ctx,
 		ps.dispatch,
@@ -239,7 +245,11 @@
 
 	countObserver := query.NewCountObserver()
 
-	planContext := dispatch.NewPlanContext(atRevision.String(), schemaHash, caveatContext, int(ps.config.MaximumAPIDepth), 0)
+	maximumAPIDepth, err := safecast.Convert[int32](ps.config.MaximumAPIDepth)
+	if err != nil {
+		return ps.rewriteError(ctx, err)
+	}
+	planContext := dispatch.NewPlanContext(atRevision.String(), schemaHash, caveatContext, maximumAPIDepth, 0)
 	qctx := dispatch.NewQueryContext(
 		ctx,
 		ps.dispatch,
@@ -374,7 +384,11 @@
 
 	countObserver := query.NewCountObserver()
 
-	planContext := dispatch.NewPlanContext(atRevision.String(), schemaHash, caveatContext, int(ps.config.MaximumAPIDepth), 0)
+	maximumAPIDepth, err := safecast.Convert[int32](ps.config.MaximumAPIDepth)
+	if err != nil {
+		return ps.rewriteError(ctx, err)
+	}
+	planContext := dispatch.NewPlanContext(atRevision.String(), schemaHash, caveatContext, maximumAPIDepth, 0)
 	qctx := dispatch.NewQueryContext(
 		ctx,
 		ps.dispatch,

internal/testserver/datastore/datastore.go

@@ -9,7 +9,6 @@ import (
 	crdbversion "github.com/authzed/spicedb/internal/datastore/crdb/version"
 	pgversion "github.com/authzed/spicedb/internal/datastore/postgres/version"
 	"github.com/authzed/spicedb/pkg/datastore"
-	"github.com/authzed/spicedb/pkg/migrate"
 )
 
 // InitFunc initializes a datastore instance from a uri that has been
@@ -50,11 +49,11 @@ func RunDatastoreEngine(t testing.TB, engine string, opts ...testcontainers.Cont
 		if ver == "" {
 			ver = pgversion.LatestTestedPostgresVersion
 		}
-		return RunPostgresForTesting(t, migrate.Head, ver, false, opts...)
+		return RunPostgresForTesting(t, ver, false, opts...)
 	case "mysql":
 		return RunMySQLForTesting(t, opts...)
 	case "spanner":
-		return RunSpannerForTesting(t, migrate.Head, opts...)
+		return RunSpannerForTesting(t, opts...)
 	default:
 		t.Fatalf("found missing engine for RunDatastoreEngine: %s", engine)
 		return nil

internal/testserver/datastore/postgres.go

@@ -33,22 +33,19 @@ const (
 )
 
 type postgresTester struct {
-	targetMigration string
-	pgbouncerProxy  *postgres.PostgresContainer
-	pgContainer     *postgres.PostgresContainer
+	pgbouncerProxy *postgres.PostgresContainer
+	pgContainer    *postgres.PostgresContainer
 }
 
 // RunPostgresForTesting returns a RunningEngineForTest for postgres
-func RunPostgresForTesting(t testing.TB, targetMigration string, pgVersion string, enablePgbouncer bool, opts ...testcontainers.ContainerCustomizer) RunningEngineForTest {
-	return RunPostgresForTestingWithCommitTimestamps(t, targetMigration, true, pgVersion, enablePgbouncer, opts...)
+func RunPostgresForTesting(t testing.TB, pgVersion string, enablePgbouncer bool, opts ...testcontainers.ContainerCustomizer) RunningEngineForTest {
+	return RunPostgresForTestingWithCommitTimestamps(t, true, pgVersion, enablePgbouncer, opts...)
 }
 
-func RunPostgresForTestingWithCommitTimestamps(t testing.TB, targetMigration string, withCommitTimestamps bool, pgVersion string, enablePgbouncer bool, opts ...testcontainers.ContainerCustomizer) RunningEngineForTest {
+func RunPostgresForTestingWithCommitTimestamps(t testing.TB, withCommitTimestamps bool, pgVersion string, enablePgbouncer bool, opts ...testcontainers.ContainerCustomizer) RunningEngineForTest {
 	t.Helper()
 
-	builder := &postgresTester{
-		targetMigration: targetMigration,
-	}
+	builder := &postgresTester{}
 
 	if enablePgbouncer {
 		// if we are running with pgbouncer enabled then set it up
@@ -131,7 +128,7 @@ func (b *postgresTester) NewDatastore(t testing.TB, initFunc InitFunc) datastore
 		defer func() {
 			migrationDriver.Close(ctx)
 		}()
-		if !assert.NoError(collect, pgmigrations.DatabaseMigrations.Run(ctx, migrationDriver, b.targetMigration, migrate.LiveRun)) {
+		if !assert.NoError(collect, pgmigrations.DatabaseMigrations.Run(ctx, migrationDriver, "head", migrate.LiveRun)) {
 			return
 		}
 		uri = connectStr

internal/testserver/datastore/spanner.go

@@ -11,6 +11,7 @@
 	adminpb "cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
 	instances "cloud.google.com/go/spanner/admin/instance/apiv1"
 	"cloud.google.com/go/spanner/admin/instance/apiv1/instancepb"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/wait"
@@ -22,13 +23,11 @@
 )
 
 type spannerTest struct {
-	hostname        string
-	port            string
-	targetMigration string
+	instancesClient *instances.InstanceAdminClient
 }
 
 // RunSpannerForTesting returns a RunningEngineForTest for spanner
-func RunSpannerForTesting(t testing.TB, targetMigration string, opts ...testcontainers.ContainerCustomizer) RunningEngineForTest {
+func RunSpannerForTesting(t testing.TB, opts ...testcontainers.ContainerCustomizer) RunningEngineForTest {
 	ctx := t.Context()
 
 	options := make([]testcontainers.ContainerCustomizer, 0, len(opts)+2)
@@ -53,13 +52,18 @@
 	// before any admin client is created below.
 	t.Setenv("SPANNER_EMULATOR_HOST", net.JoinHostPort(host, mappedPort.Port()))
 
+	builder := &spannerTest{}
+
 	// Wait until the emulator's admin API is responsive by creating an initial instance.
-	require.Eventually(t, func() bool {
+	require.EventuallyWithT(t, func(collect *assert.CollectT) {
 		instancesClient, err := instances.NewInstanceAdminClient(ctx)
-		if err != nil {
-			return false
+		if !assert.NoError(t, err) {
+			return
 		}
-		defer instancesClient.Close()
+		t.Cleanup(func() {
+			_ = instancesClient.Close()
+		})
+		builder.instancesClient = instancesClient
 
 		_, err = instancesClient.CreateInstance(ctx, &instancepb.CreateInstanceRequest{
 			Parent:     "projects/fake-project-id",
@@ -70,14 +74,9 @@
 				NodeCount:   1,
 			},
 		})
-		return err == nil
+		assert.NoError(t, err)
 	}, time.Minute, 500*time.Millisecond)
 
-	builder := &spannerTest{
-		hostname:        host,
-		port:            mappedPort.Port(),
-		targetMigration: targetMigration,
-	}
 	return builder
 }
 
@@ -92,11 +91,7 @@
 	ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
 	defer cancel()
 
-	instancesClient, err := instances.NewInstanceAdminClient(ctx)
-	require.NoError(t, err)
-	defer instancesClient.Close()
-
-	createInstanceOp, err := instancesClient.CreateInstance(ctx, &instancepb.CreateInstanceRequest{
+	createInstanceOp, err := b.instancesClient.CreateInstance(ctx, &instancepb.CreateInstanceRequest{
 		Parent:     "projects/fake-project-id",
 		InstanceId: newInstanceName,
 		Instance: &instancepb.Instance{
@@ -135,7 +130,7 @@
 		migrationDriver.Close(t.Context())
 	}()
 
-	err = migrations.SpannerMigrations.Run(t.Context(), migrationDriver, b.targetMigration, migrate.LiveRun)
+	err = migrations.SpannerMigrations.Run(t.Context(), migrationDriver, "head", migrate.LiveRun)
 	require.NoError(t, err)
 
 	return initFunc("spanner", db)

pkg/query/benchmarks/classic_check_benchmark_test.go

@@ -106,6 +106,7 @@ func requireMember(b *testing.B, resp *v1.DispatchCheckResponse, resourceID stri
 // convention of 50.
 func classicDepth(qs *bm.QuerySets) uint32 {
 	if qs.MaxRecursionDepth > 0 {
+		//nolint:gosec  // it'd be hard to overflow this reasonably
 		return uint32(qs.MaxRecursionDepth)
 	}
 	return classicDepthRemaining

pkg/query/canonicalize.go

@@ -58,7 +58,7 @@ func extractCaveats(outline Outline) (Outline, []*core.ContextualizedCaveat, err
 		caveat := outline.Args.Caveat
 		var err error
 		var childTree Outline
-		var childCaveats []*core.ContextualizedCaveat
+		childCaveats := make([]*core.ContextualizedCaveat, 0, 1)
 
 		if len(outline.SubOutlines) > 0 {
 			childTree, childCaveats, err = extractCaveats(outline.SubOutlines[0])