Add target CREATEDB preflight check for snapshot DB creation

pkg/stream/config.go

@@ -204,6 +204,26 @@ func (c *Config) isInjectorEnabled() bool {
 	return c.Processor.Injector != nil && c.Processor.Injector.URL != ""
 }
 
+func (c *Config) SnapshotTargetPostgresURL() string {
+	if c.Listener.Postgres == nil ||
+		c.Listener.Postgres.Snapshot == nil ||
+		c.Listener.Postgres.Snapshot.Schema == nil ||
+		c.Listener.Postgres.Snapshot.Schema.DumpRestore == nil {
+		return ""
+	}
+	return c.Listener.Postgres.Snapshot.Schema.DumpRestore.TargetPGURL
+}
+
+func (c *Config) SnapshotCreateTargetDB() bool {
+	if c.Listener.Postgres == nil ||
+		c.Listener.Postgres.Snapshot == nil ||
+		c.Listener.Postgres.Snapshot.Schema == nil ||
+		c.Listener.Postgres.Snapshot.Schema.DumpRestore == nil {
+		return false
+	}
+	return c.Listener.Postgres.Snapshot.Schema.DumpRestore.CreateTargetDB
+}
+
 // restoreConflictTargetsBeforeData reports whether the schema snapshot must
 // restore primary keys, unique constraints and unique indexes before the data
 // snapshot runs. This is required when the postgres batch writer emits

pkg/stream/preflight/access.go

@@ -32,6 +32,14 @@ func (c *SourceSequenceSelectPrivilegesCheck) Name() string {
 	return "source_sequence_select_privileges"
 }
 
+type TargetCreateDBPrivilegeCheck struct {
+	Target postgres.AcquireFunc
+}
+
+func (c *TargetCreateDBPrivilegeCheck) Name() string {
+	return "target_createdb_privilege"
+}
+
 const sourceTableSelectPrivilegesQuery = `
 SELECT
   current_user,
@@ -94,6 +102,14 @@ WHERE c.relkind = 'S'
 ORDER BY table_schema, table_name, sequence_schema, sequence_name
 `
 
+const targetCreateDBPrivilegeQuery = `
+SELECT 
+	current_user,
+	rolcreatedb 
+FROM pg_roles
+WHERE rolname = current_user
+`
+
 func (c *SourceTableSelectPrivilegesCheck) Run(ctx context.Context) ([]Finding, error) {
 	conn, err := c.Source(ctx)
 	if err != nil {
@@ -161,6 +177,26 @@ func (c *SourceSequenceSelectPrivilegesCheck) Run(ctx context.Context) ([]Findin
 	return findings, nil
 }
 
+func (c *TargetCreateDBPrivilegeCheck) Run(ctx context.Context) ([]Finding, error) {
+	conn, err := c.Target(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("connecting to target: %w", err)
+	}
+	var role string
+	var hasCreateDB bool
+
+	if err := conn.QueryRow(ctx, []any{&role, &hasCreateDB}, targetCreateDBPrivilegeQuery); err != nil {
+		return nil, fmt.Errorf("querying target CREATEDB privilege: %w", err)
+	}
+
+	if hasCreateDB {
+		return nil, nil
+	}
+	return []Finding{
+		{Message: targetCreateDBPrivilegeMessage(role)},
+	}, nil
+}
+
 type sourceTableSelectPrivilegeRow struct {
 	Role      string
 	Schema    string
@@ -240,3 +276,11 @@ func sourceSequenceSelectPrivilegeMessage(row sourceSequenceSelectPrivilegeRow)
 		row.Role, row.SequenceSchema, row.Sequence, quotedSequence, quotedRole,
 	)
 }
+
+func targetCreateDBPrivilegeMessage(role string) string {
+	quotedRole := postgres.QuoteIdentifier(role)
+	return fmt.Sprintf(
+		"target role %q lacks CREATEDB; run ALTER ROLE %s CREATEDB",
+		role, quotedRole,
+	)
+}

pkg/stream/preflight/access_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/xataio/pgstream/internal/postgres"
 	"github.com/xataio/pgstream/internal/postgres/mocks"
+	pgdumprestore "github.com/xataio/pgstream/pkg/snapshot/generator/postgres/schema/pgdumprestore"
 	"github.com/xataio/pgstream/pkg/stream"
 	"github.com/xataio/pgstream/pkg/wal/listener/snapshot/adapter"
 	snapshotbuilder "github.com/xataio/pgstream/pkg/wal/listener/snapshot/builder"
@@ -402,6 +403,79 @@ func TestSourceSequenceSelectPrivilegesCheck_Name(t *testing.T) {
 	require.Equal(t, "source_sequence_select_privileges", (&SourceSequenceSelectPrivilegesCheck{}).Name())
 }
 
+func TestTargetCreateDBPrivilegeCheck_Run_HasCreateDB(t *testing.T) {
+	t.Parallel()
+
+	check := &TargetCreateDBPrivilegeCheck{
+		Target: targetWithCreateDBPrivilege(t, "pgstreamtarget", true),
+	}
+
+	findings, err := check.Run(context.Background())
+
+	require.NoError(t, err)
+	require.Empty(t, findings)
+}
+
+func TestTargetCreateDBPrivilegeCheck_Run_MissingCreateDBReturnsFinding(t *testing.T) {
+	t.Parallel()
+
+	check := &TargetCreateDBPrivilegeCheck{
+		Target: targetWithCreateDBPrivilege(t, "pgstreamtarget", false),
+	}
+
+	findings, err := check.Run(context.Background())
+
+	require.NoError(t, err)
+	require.Len(t, findings, 1)
+	require.Contains(t, findings[0].Message, `target role "pgstreamtarget"`)
+	require.Contains(t, findings[0].Message, "lacks CREATEDB")
+	require.Contains(t, findings[0].Message, `ALTER ROLE "pgstreamtarget" CREATEDB`)
+}
+
+func TestTargetCreateDBPrivilegeCheck_Run_TargetAcquireFails(t *testing.T) {
+	t.Parallel()
+
+	checkErr := errors.New("boom")
+	check := &TargetCreateDBPrivilegeCheck{
+		Target: func(context.Context) (postgres.Querier, error) {
+			return nil, checkErr
+		},
+	}
+
+	findings, err := check.Run(context.Background())
+
+	require.Nil(t, findings)
+	require.ErrorIs(t, err, checkErr)
+	require.ErrorContains(t, err, "connecting to target")
+}
+
+func TestTargetCreateDBPrivilegeCheck_Run_QueryFails(t *testing.T) {
+	t.Parallel()
+
+	queryErr := errors.New("query failed")
+	check := &TargetCreateDBPrivilegeCheck{
+		Target: func(context.Context) (postgres.Querier, error) {
+			return &mocks.Querier{
+				QueryRowFn: func(context.Context, []any, string, ...any) error {
+					return queryErr
+				},
+			}, nil
+		},
+	}
+
+	findings, err := check.Run(context.Background())
+
+	require.Nil(t, findings)
+	require.ErrorIs(t, err, queryErr)
+	require.ErrorContains(t, err, "querying target CREATEDB privilege")
+}
+
+func TestTargetCreateDBPrivilegeCheck_Name(t *testing.T) {
+	t.Parallel()
+
+	require.Equal(t, "target_createdb_privilege", (&TargetCreateDBPrivilegeCheck{}).Name())
+}
+
 func TestSourceTableSelectPrivilegeMessage(t *testing.T) {
 	t.Parallel()
 
@@ -446,6 +520,16 @@ func TestSourceSequenceSelectPrivilegeMessage(t *testing.T) {
 	require.Contains(t, msg, `GRANT SELECT ON SEQUENCE "public"."orders_id_seq" TO "pgstream_user"`)
 }
 
+func TestTargetCreateDBPrivilegeMessage(t *testing.T) {
+	t.Parallel()
+
+	msg := targetCreateDBPrivilegeMessage("pgstreamtarget")
+
+	require.Contains(t, msg, `target role "pgstreamtarget"`)
+	require.Contains(t, msg, "lacks CREATEDB")
+	require.Contains(t, msg, `ALTER ROLE "pgstreamtarget" CREATEDB`)
+}
+
 func TestBuildAccessChecks(t *testing.T) {
 	t.Parallel()
 
@@ -469,6 +553,25 @@ func TestBuildAccessChecks(t *testing.T) {
 			},
 			wantChecks: 2,
 		},
+		{
+			name: "create_target_db adds target createdb check",
+			cfg: &stream.Config{
+				Listener: stream.ListenerConfig{
+					Postgres: &stream.PostgresListenerConfig{
+						URL: "postgres://source",
+						Snapshot: &snapshotbuilder.SnapshotListenerConfig{
+							Schema: &snapshotbuilder.SchemaSnapshotConfig{
+								DumpRestore: &pgdumprestore.Config{
+									TargetPGURL:    "postgres://target",
+									CreateTargetDB: true,
+								},
+							},
+						},
+					},
+				},
+			},
+			wantChecks: 3,
+		},
 		{
 			name: "snapshot+filter Include unions through AccessTableSelection",
 			cfg: &stream.Config{
@@ -513,6 +616,10 @@ func TestBuildAccessChecks(t *testing.T) {
 			require.True(t, ok)
 			require.ElementsMatch(t, tc.wantInclude, sequenceCheck.Selection.Include(), "Include lists differ")
 			require.ElementsMatch(t, tc.wantExclude, sequenceCheck.Selection.Exclude(), "Exclude lists differ")
+			if tc.wantChecks > 2 {
+				_, ok = checks[2].(*TargetCreateDBPrivilegeCheck)
+				require.True(t, ok)
+			}
 		})
 	}
 }
@@ -536,6 +643,38 @@ func TestBuildChecks_SelectedAccessOnly(t *testing.T) {
 	require.Equal(t, "source_sequence_select_privileges", checks[1].Name())
 }
 
+func TestRemoveDatabaseFromConnectionString(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name       string
+		connection string
+		want       string
+	}{
+		{
+			name:       "database is removed",
+			connection: "postgres://pgstream:secret@localhost:5432/target_db?sslmode=disable",
+			want:       "postgres://pgstream:secret@localhost:5432/?sslmode=disable",
+		},
+		{
+			name:       "postgres database is preserved",
+			connection: "postgres://pgstream:secret@localhost:5432/postgres?sslmode=disable",
+			want:       "postgres://pgstream:secret@localhost:5432/postgres?sslmode=disable",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			got, err := removeDatabaseFromConnectionString(tc.connection)
+
+			require.NoError(t, err)
+			require.Equal(t, tc.want, got)
+		})
+	}
+}
+
 func sourceWithRows(t *testing.T, rows []sourceTableSelectPrivilegeRow) postgres.AcquireFunc {
 	t.Helper()
 	return sourceWithMockRows(privilegeRows(t, rows))
@@ -583,6 +722,24 @@ func sourceWithSequenceRows(t *testing.T, rows []sourceSequenceSelectPrivilegeRo
 	return sourceWithMockRows(sequencePrivilegeRows(t, rows))
 }
 
+func targetWithCreateDBPrivilege(t *testing.T, role string, hasCreateDB bool) postgres.AcquireFunc {
+	t.Helper()
+	return func(context.Context) (postgres.Querier, error) {
+		return &mocks.Querier{
+			QueryRowFn: func(_ context.Context, dest []any, _ string, _ ...any) error {
+				require.Len(t, dest, 2)
+				roleDest, ok := dest[0].(*string)
+				require.True(t, ok)
+				createDBDest, ok := dest[1].(*bool)
+				require.True(t, ok)
+				*roleDest = role
+				*createDBDest = hasCreateDB
+				return nil
+			},
+		}, nil
+	}
+}
+
 func sequencePrivilegeRows(t *testing.T, rows []sourceSequenceSelectPrivilegeRow) postgres.Rows {
 	t.Helper()
 	return &mocks.Rows{

pkg/stream/preflight/builder.go

@@ -4,6 +4,7 @@ package preflight
 
 import (
 	"context"
+	"strings"
 
 	"github.com/xataio/pgstream/internal/postgres"
 	"github.com/xataio/pgstream/pkg/stream"
@@ -117,13 +118,13 @@ func BuildReplicationChecks(cfg *stream.Config) ([]Check, CleanupFunc) {
 // BuildAccessChecks returns the access-preflight checks applicable to cfg,
 // plus a cleanup function that closes the shared source connection.
 func BuildAccessChecks(cfg *stream.Config) ([]Check, CleanupFunc) {
-	url := cfg.SourcePostgresURL()
-	if url == "" {
+	sourceURL := cfg.SourcePostgresURL()
+	if sourceURL == "" {
 		return nil, nil
 	}
-	src := postgres.NewLazyConn(url)
+	src := postgres.NewLazyConn(sourceURL)
 	selection := cfg.AccessTableSelection()
-	return []Check{
+	checks := []Check{
 		&SourceTableSelectPrivilegesCheck{
 			Source:    src.Acquire,
 			Selection: selection,
@@ -132,7 +133,51 @@ func BuildAccessChecks(cfg *stream.Config) ([]Check, CleanupFunc) {
 			Source:    src.Acquire,
 			Selection: selection,
 		},
-	}, src.Close
+	}
+
+	cleanups := []CleanupFunc{src.Close}
+
+	if cfg.SnapshotCreateTargetDB() {
+		if targetURL := cfg.SnapshotTargetPostgresURL(); targetURL != "" {
+			var err error
+			targetURL, err = removeDatabaseFromConnectionString(targetURL)
+			if err != nil {
+				checks = append(checks, &TargetCreateDBPrivilegeCheck{
+					Target: func(context.Context) (postgres.Querier, error) {
+						return nil, err
+					},
+				})
+				return checks, joinCleanups(cleanups)
+			}
+			target := postgres.NewLazyConn(targetURL)
+			checks = append(checks, &TargetCreateDBPrivilegeCheck{Target: target.Acquire})
+			cleanups = append(cleanups, target.Close)
+		}
+	}
+
+	return checks, func(ctx context.Context) error {
+		var firstErr error
+		for _, cleanup := range cleanups {
+			if err := cleanup(ctx); err != nil && firstErr == nil {
+				firstErr = err
+			}
+		}
+		return firstErr
+	}
+}
+
+// removeDatabaseFromConnectionString lets the CREATEDB preflight connect to
+// the target server before the configured target database exists.
+func removeDatabaseFromConnectionString(url string) (string, error) {
+	pgCfg, err := postgres.ParseConfig(url)
+	if err != nil {
+		return "", err
+	}
+	if pgCfg.Database == "" || pgCfg.Database == "postgres" {
+		return url, nil
+	}
+
+	return strings.ReplaceAll(url, "/"+pgCfg.Database, "/"), nil
 }
 
 // BuildSchemaChecks returns the schema-preflight checks applicable to cfg,