diff --git a/internal/postgres/pg_restore.go b/internal/postgres/pg_restore.go index 76735844..146bab49 100644 --- a/internal/postgres/pg_restore.go +++ b/internal/postgres/pg_restore.go @@ -94,7 +94,7 @@ func RunPGRestore(ctx context.Context, opts PGRestoreOptions, dump []byte) (stri // does not include it so that pg_restore can create it. if opts.Create { var err error - opts.ConnectionString, err = removeDatabaseFromConnectionString(opts.ConnectionString) + opts.ConnectionString, err = RemoveDatabaseFromConnectionString(opts.ConnectionString) if err != nil { return "", err } @@ -165,7 +165,7 @@ func tailOutput(out []byte, maxBytes int) string { return "..." + redacted[len(redacted)-maxBytes:] } -func removeDatabaseFromConnectionString(url string) (string, error) { +func RemoveDatabaseFromConnectionString(url string) (string, error) { dbName, err := extractDatabase(url) if err != nil { return "", err diff --git a/internal/postgres/pg_restore_test.go b/internal/postgres/pg_restore_test.go index 4064d80a..2d069fe0 100644 --- a/internal/postgres/pg_restore_test.go +++ b/internal/postgres/pg_restore_test.go @@ -13,6 +13,38 @@ import ( "github.com/stretchr/testify/require" ) +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 TestPGRestoreOptionsToPGOptions(t *testing.T) { t.Parallel() diff --git a/pkg/stream/config.go b/pkg/stream/config.go index 682525d5..7fb77aa5 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..3810d50b 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 OR rolsuper +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..65706618 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,106 @@ 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_SuperuserHasCreateDB(t *testing.T) { + t.Parallel() + + check := &TargetCreateDBPrivilegeCheck{ + Target: func(context.Context) (postgres.Querier, error) { + return &mocks.Querier{ + QueryRowFn: func(_ context.Context, dest []any, query string, _ ...any) error { + require.Contains(t, query, "rolcreatedb OR rolsuper") + require.Len(t, dest, 2) + role, ok := dest[0].(*string) + require.True(t, ok) + hasCreateDB, ok := dest[1].(*bool) + require.True(t, ok) + *role = "postgres" + *hasCreateDB = true + return nil + }, + }, nil + }, + } + + 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 +547,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 +580,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 +643,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) + } }) } } @@ -583,6 +717,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 2bd8deb5..789ce668 100644 --- a/pkg/stream/preflight/builder.go +++ b/pkg/stream/preflight/builder.go @@ -117,13 +117,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 +132,29 @@ 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 = postgres.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, joinCleanups(cleanups) } // BuildSchemaChecks returns the schema-preflight checks applicable to cfg,