Support streaming PostgreSQL dump and restore

database/CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## Unreleased
+
+- Add `database-transfer`, a basic dump/restore CLI that supports `-` for
+  standard output and input, respectively.
+
 ## [4.6.1] - 2026-08-28
 
 - Fix `template_database(close_source_connections=True)` leaving the source

database/macrostrat/database/transfer/__init__.py

@@ -6,3 +6,11 @@
 from .dump_database import pg_dump, pg_dump_to_file
 from .move_tables import move_tables
 from .restore_database import pg_restore, pg_restore_from_file
+
+__all__ = [
+    "move_tables",
+    "pg_dump",
+    "pg_dump_to_file",
+    "pg_restore",
+    "pg_restore_from_file",
+]

database/macrostrat/database/transfer/cli.py

@@ -0,0 +1,70 @@
+"""A basic command-line interface for PostgreSQL database transfer.
+
+For example, stream a schema between databases:
+
+    database-transfer --database "$SOURCE_DATABASE" dump -n temp - \
+      | database-transfer --database "$TARGET_DATABASE" restore -
+"""
+
+import asyncio
+from pathlib import Path
+
+import click
+from sqlalchemy import create_engine
+
+from .dump_database import pg_dump_to_file
+from .restore_database import pg_restore_from_file
+
+
+@click.group()
+@click.option(
+    "--database",
+    envvar="DATABASE_URL",
+    required=True,
+    help="PostgreSQL connection URL.",
+)
+@click.pass_context
+def cli(ctx, database):
+    """Dump and restore PostgreSQL databases."""
+    ctx.ensure_object(dict)
+    ctx.obj["database"] = database
+
+
+def _engine(ctx):
+    return create_engine(ctx.obj["database"])
+
+
+@cli.command()
+@click.option("-n", "--schema", multiple=True, help="Schema to dump.")
+@click.argument("destination")
+@click.pass_context
+def dump(ctx, schema, destination):
+    """Dump to DESTINATION, or '-' for standard output."""
+    args = [arg for name in schema for arg in ("--schema", name)]
+    engine = _engine(ctx)
+    try:
+        asyncio.run(
+            pg_dump_to_file(
+                engine, None if destination == "-" else Path(destination), args=args
+            )
+        )
+    finally:
+        engine.dispose()
+
+
+@cli.command()
+@click.argument("source")
+@click.pass_context
+def restore(ctx, source):
+    """Restore from SOURCE, or '-' for standard input."""
+    engine = _engine(ctx)
+    try:
+        asyncio.run(
+            pg_restore_from_file(None if source == "-" else Path(source), engine)
+        )
+    finally:
+        engine.dispose()
+
+
+if __name__ == "__main__":
+    cli()

database/macrostrat/database/transfer/dump_database.py

@@ -1,5 +1,4 @@
 import asyncio
-import sys
 from pathlib import Path
 from typing import Optional
 
@@ -8,7 +7,7 @@
 
 from macrostrat.utils import get_logger
 
-from .stream_utils import print_stdout, print_stream_progress
+from .stream_utils import print_stderr, print_stdout, print_stream_progress
 from .utils import _create_command
 
 log = get_logger(__name__)
@@ -48,11 +47,11 @@ async def pg_dump(
 
 
 async def pg_dump_to_file(engine: Engine, dumpfile: Path | None, **kwargs):
-    proc = await pg_dump(engine, **kwargs)
-    if dumpfile is None or dumpfile == sys.stdout:
-        # If we have no dumpfile, just print to stdout
-        await _monitor_stdout(proc)
+    if dumpfile is None:
+        proc = await pg_dump(engine, stdout=None, **kwargs)
+        await _monitor_stderr(proc)
         return
+    proc = await pg_dump(engine, **kwargs)
     # Open dump file as an async stream
     async with aiofiles.open(dumpfile, mode="wb") as dest:
         await asyncio.gather(
@@ -61,8 +60,5 @@ async def pg_dump_to_file(engine: Engine, dumpfile: Path | None, **kwargs):
         )
 
 
-async def _monitor_stdout(proc):
-    await asyncio.gather(
-        asyncio.create_task(print_stdout(proc.stdout)),
-        asyncio.create_task(print_stream_progress(proc.stderr, None)),
-    )
+async def _monitor_stderr(proc):
+    await asyncio.gather(print_stderr(proc.stderr), proc.wait())

database/macrostrat/database/transfer/restore_database.py

@@ -8,7 +8,7 @@
 
 from macrostrat.utils import get_logger
 
-from .stream_utils import print_stdout, print_stream_progress
+from .stream_utils import print_stderr, print_stdout, print_stream_progress
 from .utils import _create_command, _create_database_if_not_exists
 
 console = Console()
@@ -23,6 +23,7 @@ async def pg_restore(
     command_prefix: Optional[list] = None,
     args: list = [],
     postgres_container: str = "postgres:15",
+    stdin=asyncio.subprocess.PIPE,
 ):
     # Pipe file to pg_restore, mimicking
 
@@ -46,13 +47,18 @@ async def pg_restore(
 
     return await asyncio.create_subprocess_exec(
         *_cmd,
-        stdin=asyncio.subprocess.PIPE,
+        stdin=stdin,
         stderr=asyncio.subprocess.PIPE,
         limit=1024 * 1024 * 1,  # 1 MB windows
     )
 
 
-async def pg_restore_from_file(dumpfile: Path, engine: Engine, **kwargs):
+async def pg_restore_from_file(dumpfile: Path | None, engine: Engine, **kwargs):
+    if dumpfile is None:
+        proc = await pg_restore(engine, stdin=None, **kwargs)
+        await asyncio.gather(print_stderr(proc.stderr), proc.wait())
+        return
+
     proc = await pg_restore(engine, **kwargs)
     # Open dump file as an async stream
     async with aiofiles.open(dumpfile, mode="rb") as source:

database/macrostrat/database/transfer/stream_utils.py

@@ -75,6 +75,13 @@ async def print_stdout(stream: asyncio.StreamReader):
         console.print(line.decode("utf-8"), style="dim")
 
 
+async def print_stderr(stream: asyncio.StreamReader):
+    async for line in stream:
+        log.info(line)
+        sys.stderr.buffer.write(line)
+        sys.stderr.buffer.flush()
+
+
 class DecodingStreamReader(asyncio.StreamReader):
     """A StreamReader that decompresses gzip files (if compressed)"""
 

database/pyproject.toml

@@ -25,6 +25,9 @@ dependencies = [
     "psycopg2>=2.9.11,<3",
 ]
 
+[project.scripts]
+database-transfer = "macrostrat.database.transfer.cli:cli"
+
 [dependency-groups]
 dev = ["macrostrat.utils"]
 

database/tests/test_transfer.py

@@ -0,0 +1,110 @@
+import asyncio
+from types import SimpleNamespace
+
+from click.testing import CliRunner
+
+from macrostrat.database.transfer import dump_database, restore_database, stream_utils
+from macrostrat.database.transfer.cli import cli
+
+
+def test_dump_to_stdout_inherits_standard_output(monkeypatch):
+    calls = {}
+
+    async def fake_dump(engine, **kwargs):
+        calls.update(kwargs)
+        return SimpleNamespace(stderr=object(), wait=fake_wait)
+
+    async def fake_wait():
+        calls["waited"] = True
+
+    async def fake_print_stderr(stream):
+        calls["stderr"] = stream
+
+    monkeypatch.setattr(dump_database, "pg_dump", fake_dump)
+    monkeypatch.setattr(dump_database, "print_stderr", fake_print_stderr)
+
+    asyncio.run(dump_database.pg_dump_to_file(object(), None))
+
+    assert calls["stdout"] is None
+    assert calls["stderr"] is not None
+    assert calls["waited"] is True
+
+
+def test_restore_from_stdin_inherits_standard_input(monkeypatch):
+    calls = {}
+
+    async def fake_restore(engine, **kwargs):
+        calls.update(kwargs)
+        return SimpleNamespace(stderr=object(), wait=fake_wait)
+
+    async def fake_wait():
+        calls["waited"] = True
+
+    async def fake_print_stderr(stream):
+        calls["stderr"] = stream
+
+    monkeypatch.setattr(restore_database, "pg_restore", fake_restore)
+    monkeypatch.setattr(restore_database, "print_stderr", fake_print_stderr)
+
+    asyncio.run(restore_database.pg_restore_from_file(None, object()))
+
+    assert calls["stdin"] is None
+    assert calls["stderr"] is not None
+    assert calls["waited"] is True
+
+
+def test_cli_streams_schema_dump_to_standard_output(monkeypatch):
+    calls = {}
+
+    async def fake_dump(engine, destination, **kwargs):
+        calls["url"] = str(engine.url)
+        calls["destination"] = destination
+        calls["args"] = kwargs["args"]
+
+    monkeypatch.setattr(
+        "macrostrat.database.transfer.cli.pg_dump_to_file", fake_dump
+    )
+
+    result = CliRunner().invoke(
+        cli, ["--database", "postgresql://localhost/source", "dump", "-n", "temp", "-"]
+    )
+
+    assert result.exit_code == 0, result.output
+    assert calls == {
+        "url": "postgresql://localhost/source",
+        "destination": None,
+        "args": ["--schema", "temp"],
+    }
+
+
+def test_cli_restores_from_standard_input(monkeypatch):
+    calls = {}
+
+    async def fake_restore(source, engine):
+        calls["source"] = source
+        calls["url"] = str(engine.url)
+
+    monkeypatch.setattr(
+        "macrostrat.database.transfer.cli.pg_restore_from_file", fake_restore
+    )
+
+    result = CliRunner().invoke(
+        cli, ["--database", "postgresql://localhost/target", "restore", "-"]
+    )
+
+    assert result.exit_code == 0, result.output
+    assert calls == {"source": None, "url": "postgresql://localhost/target"}
+
+
+def test_print_stderr_does_not_write_to_standard_output(capsysbinary):
+    async def write_stderr():
+        stream = asyncio.StreamReader()
+        stream.feed_data(b"pg_dump: warning\n")
+        stream.feed_eof()
+        await stream_utils.print_stderr(stream)
+
+    asyncio.run(write_stderr())
+
+    captured = capsysbinary.readouterr()
+    assert captured.out == b""
+    assert captured.err == b"pg_dump: warning\n"