diff --git a/pkg/stream/config.go b/pkg/stream/config.go index 2041c713..e18443a1 100644 --- a/pkg/stream/config.go +++ b/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 diff --git a/pkg/stream/preflight/access.go b/pkg/stream/preflight/access.go index ec685d50..63e660e9 100644 --- a/pkg/stream/preflight/access.go +++ b/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, + ) +} diff --git a/pkg/stream/preflight/access_test.go b/pkg/stream/preflight/access_test.go index 78b4b75f..d1d635ef 100644 --- a/pkg/stream/preflight/access_test.go +++ b/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{ diff --git a/pkg/stream/preflight/builder.go b/pkg/stream/preflight/builder.go index 3863d4f1..4818e7dc 100644 --- a/pkg/stream/preflight/builder.go +++ b/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" @@ -95,13 +96,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, @@ -110,7 +111,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,