diff --git a/steward/backup/restore.py b/steward/backup/restore.py index 26697b1..de4b0d0 100644 --- a/steward/backup/restore.py +++ b/steward/backup/restore.py @@ -80,32 +80,61 @@ def dangling_fixes(metadata: MetaData) -> list[str]: return out +# pg_restore 17 opens every script with this; a Postgres 16 server doesn't +# know the setting and refuses it, so the line is dropped on the way through. +_NEWER_THAN_SERVER = {b"SET transaction_timeout = 0;"} +# A COPY row is one line of the script; this bounds the longest one read. +_LINE_LIMIT = 256 * 1024 * 1024 + + async def _pg_restore(database_url: str, dump: Path, section: str) -> None: - args = ["pg_restore", "--no-owner", "--no-privileges", "--exit-on-error", - f"--section={section}", "--dbname", make_dbname(database_url), str(dump)] + """One section of the dump into the database: pg_restore writes it as a + script, psql runs it and stops at the first error.""" + env = pg_env(database_url) try: - proc = await asyncio.create_subprocess_exec( - *args, env=pg_env(database_url), + script = await asyncio.create_subprocess_exec( + "pg_restore", "--no-owner", "--no-privileges", f"--section={section}", "--file=-", str(dump), + env=env, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, limit=_LINE_LIMIT) + psql = await asyncio.create_subprocess_exec( + "psql", "-X", "-q", "-v", "ON_ERROR_STOP=1", + env=env, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE) except FileNotFoundError: - raise BackupError("pg_restore isn't installed here.") from None + raise BackupError("pg_restore or psql isn't installed here.") from None + + async def pump() -> None: + try: + while line := await script.stdout.readline(): + if line.strip() in _NEWER_THAN_SERVER: + continue + psql.stdin.write(line) + await psql.stdin.drain() + except (BrokenPipeError, ConnectionResetError): + # psql stopped on an error (its stderr says which); pg_restore + # would block on a pipe nobody reads. + script.kill() + finally: + psql.stdin.close() + + async def run() -> tuple[bytes, bytes]: + _, script_err, psql_err = await asyncio.gather( + pump(), script.stderr.read(), psql.stderr.read()) + await asyncio.gather(script.wait(), psql.wait()) + return script_err, psql_err + try: - _, err = await asyncio.wait_for(proc.communicate(), RESTORE_DEADLINE_SECONDS) + script_err, psql_err = await asyncio.wait_for(run(), RESTORE_DEADLINE_SECONDS) except TimeoutError: - proc.kill() - await proc.wait() - raise BackupError(f"pg_restore ({section}) didn't finish within " + for proc in (script, psql): + if proc.returncode is None: + proc.kill() + await proc.wait() + raise BackupError(f"Restoring {section} didn't finish within " f"{RESTORE_DEADLINE_SECONDS // 60} minutes.") from None - if proc.returncode != 0: - lines = [ln for ln in err.decode(errors="replace").splitlines() if ln.strip()] - raise BackupError(f"pg_restore ({section}): " + (lines[-1] if lines else f"exited {proc.returncode}.")) - - -def make_dbname(database_url: str) -> str: - """The database name for pg_restore's --dbname; host, user and password - come from the environment pg_env builds.""" - from sqlalchemy.engine import make_url - return make_url(database_url).database or "" + for proc, err, name in ((psql, psql_err, "psql"), (script, script_err, "pg_restore")): + if proc.returncode != 0: + lines = [ln for ln in err.decode(errors="replace").splitlines() if ln.strip()] + raise BackupError(f"Restoring {section}: " + (lines[-1] if lines else f"{name} exited {proc.returncode}.")) def _async_url(database_url: str) -> str: