1 files changed, 35 insertions, 24 deletions

diff --git a/syz-cluster/pkg/db/spanner.go b/syz-cluster/pkg/db/spanner.go
index c1484df74..148a58e35 100644
--- a/syz-cluster/pkg/db/spanner.go
+++ b/syz-cluster/pkg/db/spanner.go
@@ -72,25 +72,30 @@ func CreateSpannerInstance(ctx context.Context, uri ParsedURI) error {
 	return err
 }
 
-func CreateSpannerDB(ctx context.Context, uri ParsedURI) error {
+func CreateSpannerDB(ctx context.Context, uri ParsedURI, ddl []string) error {
 	client, err := database.NewDatabaseAdminClient(ctx)
 	if err != nil {
 		return err
 	}
 	defer client.Close()
-	_, err = client.GetDatabase(ctx, &databasepb.GetDatabaseRequest{Name: uri.Full})
-	if err != nil && spanner.ErrCode(err) == codes.NotFound {
-		op, err := client.CreateDatabase(ctx, &databasepb.CreateDatabaseRequest{
-			Parent:          uri.InstancePrefix,
-			CreateStatement: `CREATE DATABASE ` + uri.Database,
-			ExtraStatements: []string{},
-		})
-		if err != nil {
+	if ddl == nil {
+		_, err = client.GetDatabase(ctx, &databasepb.GetDatabaseRequest{Name: uri.Full})
+		if err != nil && spanner.ErrCode(err) != codes.NotFound {
 			return err
 		}
-		_, err = op.Wait(ctx)
+		if err == nil {
+			return nil
+		}
+	}
+	op, err := client.CreateDatabase(ctx, &databasepb.CreateDatabaseRequest{
+		Parent:          uri.InstancePrefix,
+		CreateStatement: `CREATE DATABASE ` + uri.Database,
+		ExtraStatements: ddl,
+	})
+	if err != nil {
 		return err
 	}
+	_, err = op.Wait(ctx)
 	return err
 }
 
@@ -128,31 +133,37 @@ func getMigrateInstance(uri string) (*migrate.Migrate, error) {
 }
 
 func NewTransientDB(t *testing.T) (*spanner.Client, context.Context) {
-	setupSpannerEmulator(t)
-	uri, err := ParseURI("projects/my-project/instances/test-instance/databases/" +
-		fmt.Sprintf("db%v", time.Now().UnixNano()))
-	if err != nil {
-		t.Fatal(err)
-	}
+	uri := "projects/my-project/instances/test-instance/databases/" +
+		fmt.Sprintf("db%v", time.Now().UnixNano())
+	NewTestDB(t, uri, nil)
 	ctx := t.Context()
-	err = CreateSpannerInstance(ctx, uri)
+	client, err := spanner.NewClient(ctx, uri)
 	if err != nil {
 		t.Fatal(err)
 	}
-	err = CreateSpannerDB(ctx, uri)
+	t.Cleanup(client.Close)
+	err = RunMigrations(uri)
 	if err != nil {
 		t.Fatal(err)
 	}
-	client, err := spanner.NewClient(ctx, uri.Full)
+	return client, ctx
+}
+
+func NewTestDB(t *testing.T, uri string, ddl []string) {
+	setupSpannerEmulator(t)
+	// Don't bother destroying instances/databases.
+	// We create isolated per-test databases, and the emulator is all in-memory.
+	// So when the emulator is killed with the test binary, everything is gone.
+	parsedURI, err := ParseURI(uri)
 	if err != nil {
 		t.Fatal(err)
 	}
-	t.Cleanup(client.Close)
-	err = RunMigrations(uri.Full)
-	if err != nil {
-		t.Fatal(err)
+	if err := CreateSpannerInstance(t.Context(), parsedURI); err != nil {
+		t.Fatalf("failed CreateSpannerInstance: %v", err)
+	}
+	if err := CreateSpannerDB(t.Context(), parsedURI, ddl); err != nil {
+		t.Fatalf("failed CreateSpannerDB: %v", err)
 	}
-	return client, ctx
 }
 
 var (