Skip to content

Commit cb0efdb

Browse files
feat(pg-pg): Automated schema dump mode (#4283)
### Why - For Postgres to Postgres migration use-cases, users would prefer an apples-to-apples mirroring of their source setup on the target instance. - A well-recognized method to migrate schema only is pg_dump. ### What - This PR adds an activity in `SetupFlowWorkflow` which runs pg_dump on the source database and pipes its output to the target database via PSQL. - It adds a **new dynamic flag setting to gate the above**. Also the above is gated to PG type system mirrors. - When the above flag is set, destination validation and destination create normalized tables steps are skipped. - It sets **--no-owners and --no-privileges**, leaving it to the user to add the desired roles later. `pg_dumpall` requires exact major version matching between pg_dumpall version and target version, which we cannot guarantee. - E2E tests added
1 parent af8af34 commit cb0efdb

10 files changed

Lines changed: 1229 additions & 2 deletions

File tree

flow/activities/flowable.go

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2141,3 +2141,72 @@ func (a *FlowableActivity) MigratePostgresTableOIDs(
21412141

21422142
return nil
21432143
}
2144+
2145+
func (a *FlowableActivity) PeerDBPGAutomatedSchemaDump(ctx context.Context, env map[string]string) (bool, error) {
2146+
return internal.PeerDBPGAutomatedSchemaDump(ctx, env)
2147+
}
2148+
2149+
func (a *FlowableActivity) RunPgDumpSchema(
2150+
ctx context.Context,
2151+
input *protos.RunPgDumpSchemaInput,
2152+
) (bool, error) {
2153+
logger := internal.LoggerFromCtx(ctx)
2154+
ctx = context.WithValue(ctx, shared.FlowNameKey, input.FlowName)
2155+
2156+
srcPeer, err := connectors.LoadPeer(ctx, a.CatalogPool, input.SourceName)
2157+
if err != nil {
2158+
return false, a.Alerter.LogFlowError(ctx, input.FlowName, fmt.Errorf("failed to load source peer: %w", err))
2159+
}
2160+
2161+
dstPeer, err := connectors.LoadPeer(ctx, a.CatalogPool, input.DestinationName)
2162+
if err != nil {
2163+
return false, a.Alerter.LogFlowError(ctx, input.FlowName, fmt.Errorf("failed to load destination peer: %w", err))
2164+
}
2165+
2166+
srcPgConfig, ok := srcPeer.Config.(*protos.Peer_PostgresConfig)
2167+
if !ok {
2168+
return false, a.Alerter.LogFlowError(ctx, input.FlowName, fmt.Errorf("source peer %s is not a PostgreSQL peer", input.SourceName))
2169+
}
2170+
2171+
dstPgConfig, ok := dstPeer.Config.(*protos.Peer_PostgresConfig)
2172+
if !ok {
2173+
return false, a.Alerter.LogFlowError(ctx, input.FlowName,
2174+
fmt.Errorf("destination peer %s is not a PostgreSQL peer", input.DestinationName))
2175+
}
2176+
2177+
// skip schema migration for peers using SSH tunnels
2178+
if srcPgConfig.PostgresConfig.SshConfig != nil {
2179+
logger.Info("skipping pg_dump schema migration: source peer uses SSH tunnel")
2180+
return false, nil
2181+
}
2182+
if dstPgConfig.PostgresConfig.SshConfig != nil {
2183+
logger.Info("skipping pg_dump schema migration: destination peer uses SSH tunnel")
2184+
return false, nil
2185+
}
2186+
2187+
// skip schema migration for non-password auth (e.g. IAM)
2188+
if srcPgConfig.PostgresConfig.AuthType != protos.PostgresAuthType_POSTGRES_PASSWORD {
2189+
logger.Info("skipping pg_dump schema migration: source peer uses non-password auth")
2190+
return false, nil
2191+
}
2192+
if dstPgConfig.PostgresConfig.AuthType != protos.PostgresAuthType_POSTGRES_PASSWORD {
2193+
logger.Info("skipping pg_dump schema migration: destination peer uses non-password auth")
2194+
return false, nil
2195+
}
2196+
2197+
logger.Info("running pg_dump schema migration from source to destination",
2198+
slog.String("source", input.SourceName), slog.String("destination", input.DestinationName))
2199+
a.Alerter.LogFlowInfo(ctx, input.FlowName,
2200+
fmt.Sprintf("starting pg_dump schema migration from %s to %s", input.SourceName, input.DestinationName))
2201+
2202+
start := time.Now()
2203+
if err := connpostgres.RunPgDumpSchema(ctx, srcPgConfig.PostgresConfig, dstPgConfig.PostgresConfig); err != nil {
2204+
return false, a.Alerter.LogFlowError(ctx, input.FlowName, fmt.Errorf("pg_dump schema migration failed: %w", err))
2205+
}
2206+
2207+
elapsed := time.Since(start).Round(time.Millisecond)
2208+
logger.Info("pg_dump schema migration completed successfully", slog.Duration("elapsed", elapsed))
2209+
a.Alerter.LogFlowInfo(ctx, input.FlowName,
2210+
fmt.Sprintf("pg_dump schema migration completed successfully in %s", elapsed))
2211+
return true, nil
2212+
}
Lines changed: 302 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,302 @@
1+
package connpostgres
2+
3+
import (
4+
"bufio"
5+
"bytes"
6+
"context"
7+
"fmt"
8+
"io"
9+
"log/slog"
10+
"os"
11+
"os/exec"
12+
"regexp"
13+
"strconv"
14+
15+
"github.com/PeerDB-io/peerdb/flow/generated/protos"
16+
)
17+
18+
// pg_dump from newer Postgres versions emits statements that older
19+
// destinations don't recognize:
20+
// - SET transaction_timeout = 0; (PG17+ session GUC)
21+
// - \restrict / \unrestrict <token> (pg_dump 17.6+ psql meta-commands
22+
// that gate replay against an unrelated psql session; older psql treats
23+
// them as unknown backslash commands and aborts under ON_ERROR_STOP)
24+
//
25+
// These are session/replay housekeeping and safe to drop on the wire so we
26+
// keep ON_ERROR_STOP=1 for genuine DDL failures while remaining cross-version.
27+
var incompatibleLineRE = regexp.MustCompile(`^(SET\s+transaction_timeout\s*=|\\(?:un)?restrict(\s|$))`)
28+
29+
// RunPgDumpSchema streams a schema-only pg_dump from source directly into psql
30+
// on the destination, piping stdout into stdin without intermediate files.
31+
func RunPgDumpSchema(ctx context.Context, srcConfig *protos.PostgresConfig, dstConfig *protos.PostgresConfig) error {
32+
if err := pipeCommand(ctx, srcConfig, dstConfig, "pg_dump", buildPgDumpArgs(srcConfig)); err != nil {
33+
return fmt.Errorf("pg_dump schema migration failed: %w", err)
34+
}
35+
36+
return nil
37+
}
38+
39+
// pipeCommand runs srcBinary with the given args, piping its stdout into psql on the destination.
40+
func pipeCommand(
41+
ctx context.Context,
42+
srcConfig *protos.PostgresConfig,
43+
dstConfig *protos.PostgresConfig,
44+
srcBinary string,
45+
srcArgs []string,
46+
) error {
47+
psqlArgs := buildPsqlArgs(dstConfig)
48+
49+
srcCmd := exec.CommandContext(ctx, srcBinary, srcArgs...)
50+
psqlCmd := exec.CommandContext(ctx, "psql", psqlArgs...)
51+
52+
// set PGPASSWORD for each command via separate env slices
53+
srcCmd.Env = append(os.Environ(), "PGPASSWORD="+srcConfig.Password)
54+
psqlCmd.Env = append(os.Environ(), "PGPASSWORD="+dstConfig.Password)
55+
56+
// handle TLS env vars
57+
appendTLSEnv(ctx, srcCmd, srcConfig)
58+
appendTLSEnv(ctx, psqlCmd, dstConfig)
59+
60+
return runPipeline(ctx, srcCmd, psqlCmd, srcBinary, "psql", filterIncompatibleLines)
61+
}
62+
63+
// filterIncompatibleLines copies r->w line by line, dropping statements that
64+
// are valid in newer pg_dump output but rejected by older psql/destinations.
65+
func filterIncompatibleLines(ctx context.Context, r io.Reader, w io.Writer) error {
66+
br := bufio.NewReaderSize(r, 64*1024)
67+
for {
68+
line, err := br.ReadBytes('\n')
69+
if len(line) > 0 {
70+
if !incompatibleLineRE.Match(line) {
71+
if _, werr := w.Write(line); werr != nil {
72+
return werr
73+
}
74+
} else {
75+
slog.DebugContext(ctx, "dropping incompatible line from pg_dump stream",
76+
slog.String("line", string(bytes.TrimRight(line, "\n"))))
77+
}
78+
}
79+
if err != nil {
80+
if err == io.EOF {
81+
return nil
82+
}
83+
return err
84+
}
85+
}
86+
}
87+
88+
// runPipeline wires srcCmd's stdout into dstCmd's stdin (optionally through a
89+
// filter goroutine) and waits for both processes.
90+
//
91+
// Pipe topology:
92+
//
93+
// without filter: src.stdout -> srcW |--pipe--| srcR -> dst.stdin
94+
// with filter: src.stdout -> srcW |--pipe--| srcR -> filter -> dstW |--pipe--| dstR -> dst.stdin
95+
//
96+
// File descriptor ownership matters here -- if the parent keeps a write end
97+
// open after the child consumer dies, the producer can hang forever on a
98+
// blocked write. We close each fd as soon as the child or filter goroutine
99+
// owns it.
100+
func runPipeline(
101+
ctx context.Context,
102+
srcCmd, dstCmd *exec.Cmd,
103+
srcName, dstName string,
104+
filter func(context.Context, io.Reader, io.Writer) error,
105+
) error {
106+
srcR, srcW, err := os.Pipe()
107+
if err != nil {
108+
return fmt.Errorf("create src pipe: %w", err)
109+
}
110+
srcCmd.Stdout = srcW
111+
112+
var (
113+
dstR, dstW *os.File
114+
filterDone chan error
115+
)
116+
if filter == nil {
117+
dstCmd.Stdin = srcR
118+
} else {
119+
dstR, dstW, err = os.Pipe()
120+
if err != nil {
121+
srcR.Close()
122+
srcW.Close()
123+
return fmt.Errorf("create dst pipe: %w", err)
124+
}
125+
dstCmd.Stdin = dstR
126+
filterDone = make(chan error, 1)
127+
}
128+
129+
var srcStderr, dstStderr bytes.Buffer
130+
srcCmd.Stderr = &srcStderr
131+
dstCmd.Stderr = &dstStderr
132+
133+
// Start dst first so it's ready to read.
134+
if err := dstCmd.Start(); err != nil {
135+
srcR.Close()
136+
srcW.Close()
137+
if dstW != nil {
138+
dstR.Close()
139+
dstW.Close()
140+
}
141+
return fmt.Errorf("start %s: %w", dstName, err)
142+
}
143+
// dst owns its stdin fd in its child; close our copy.
144+
if filter == nil {
145+
srcR.Close()
146+
} else {
147+
dstR.Close()
148+
}
149+
150+
if err := srcCmd.Start(); err != nil {
151+
srcW.Close()
152+
if dstW != nil {
153+
// filter never started; close its writer so dst sees EOF.
154+
dstW.Close()
155+
// and the read side we still hold if filter==nil path wasn't taken.
156+
if filter != nil {
157+
srcR.Close()
158+
}
159+
}
160+
_ = dstCmd.Process.Kill()
161+
_ = dstCmd.Wait()
162+
return fmt.Errorf("start %s: %w", srcName, err)
163+
}
164+
// src owns its stdout fd in its child; close our copy.
165+
srcW.Close()
166+
167+
// Run the filter goroutine if configured. It bridges srcR -> dstW.
168+
if filter != nil {
169+
go func() {
170+
err := filter(ctx, srcR, dstW)
171+
// Always close both ends so the producer/consumer unblock.
172+
srcR.Close()
173+
dstW.Close()
174+
filterDone <- err
175+
}()
176+
}
177+
178+
srcDone := make(chan error, 1)
179+
dstDone := make(chan error, 1)
180+
go func() { srcDone <- srcCmd.Wait() }()
181+
go func() { dstDone <- dstCmd.Wait() }()
182+
183+
var (
184+
srcErr, dstErr error
185+
srcKilled, dstKilled bool
186+
)
187+
for range 2 {
188+
select {
189+
case err := <-srcDone:
190+
srcErr = err
191+
if err != nil && dstCmd.ProcessState == nil {
192+
_ = dstCmd.Process.Kill()
193+
dstKilled = true
194+
}
195+
case err := <-dstDone:
196+
dstErr = err
197+
if srcCmd.ProcessState == nil {
198+
// dst exited (success or failure) while src is still running;
199+
// kill src so it doesn't block on a pipe with no reader.
200+
_ = srcCmd.Process.Kill()
201+
srcKilled = true
202+
}
203+
}
204+
}
205+
206+
// Wait for the filter to finish so we surface any I/O error and so the
207+
// goroutine doesn't outlive this function.
208+
var filterErr error
209+
if filterDone != nil {
210+
filterErr = <-filterDone
211+
}
212+
213+
// Report the original cause, not the side we killed in response.
214+
if dstErr != nil && !dstKilled {
215+
return fmt.Errorf("%s failed: %w\nstderr:\n%s", dstName, dstErr, dstStderr.String())
216+
}
217+
if srcErr != nil && !srcKilled {
218+
return fmt.Errorf("%s failed: %w\nstderr:\n%s", srcName, srcErr, srcStderr.String())
219+
}
220+
if filterErr != nil {
221+
return fmt.Errorf("filter failed: %w", filterErr)
222+
}
223+
// Fallback: both sides killed (e.g. ctx cancel) — surface whichever error we have.
224+
if srcErr != nil {
225+
return fmt.Errorf("%s failed: %w\nstderr:\n%s", srcName, srcErr, srcStderr.String())
226+
}
227+
if dstErr != nil {
228+
return fmt.Errorf("%s failed: %w\nstderr:\n%s", dstName, dstErr, dstStderr.String())
229+
}
230+
return nil
231+
}
232+
233+
func buildPgDumpArgs(config *protos.PostgresConfig) []string {
234+
port := config.Port
235+
if port == 0 {
236+
port = 5432
237+
}
238+
239+
args := []string{
240+
"--schema-only",
241+
"--no-owner",
242+
"--no-privileges",
243+
"-h", config.Host,
244+
"-p", strconv.FormatUint(uint64(port), 10),
245+
"-d", config.Database,
246+
}
247+
if config.User != "" {
248+
args = append(args, "-U", config.User)
249+
}
250+
return args
251+
}
252+
253+
func buildPsqlArgs(config *protos.PostgresConfig) []string {
254+
port := config.Port
255+
if port == 0 {
256+
port = 5432
257+
}
258+
259+
args := []string{
260+
"-h", config.Host,
261+
"-p", strconv.FormatUint(uint64(port), 10),
262+
"-d", config.Database,
263+
// Wrap the entire dump in a single transaction so partial failures
264+
// roll back cleanly (makes the activity safely retryable) and avoid
265+
// per-statement autocommit overhead on high-latency links.
266+
"--single-transaction",
267+
// Without this, psql logs errors to stderr but exits 0, so a half-
268+
// applied schema would be reported as success. ON_ERROR_STOP=1 makes
269+
// psql exit non-zero on the first failed statement.
270+
"-v", "ON_ERROR_STOP=1",
271+
// Quiet informational chatter; errors still go to stderr.
272+
"--quiet",
273+
}
274+
if config.User != "" {
275+
args = append(args, "-U", config.User)
276+
}
277+
return args
278+
}
279+
280+
func appendTLSEnv(ctx context.Context, cmd *exec.Cmd, config *protos.PostgresConfig) {
281+
if config.RequireTls {
282+
cmd.Env = append(cmd.Env, "PGSSLMODE=require")
283+
284+
if config.RootCa != nil && *config.RootCa != "" {
285+
// write root CA to a temp file
286+
tmpFile, err := os.CreateTemp("", "peerdb-root-ca-*.pem")
287+
if err != nil {
288+
slog.WarnContext(ctx, "failed to create temp file for root CA, skipping sslrootcert", slog.Any("error", err))
289+
return
290+
}
291+
if _, err := tmpFile.WriteString(*config.RootCa); err != nil {
292+
slog.WarnContext(ctx, "failed to write root CA to temp file", slog.Any("error", err))
293+
tmpFile.Close()
294+
os.Remove(tmpFile.Name())
295+
return
296+
}
297+
tmpFile.Close()
298+
cmd.Env = append(cmd.Env, "PGSSLROOTCERT="+tmpFile.Name())
299+
// note: temp file is cleaned up when the process exits
300+
}
301+
}
302+
}

0 commit comments

Comments
 (0)