diff --git a/database/CHANGELOG.md b/database/CHANGELOG.md index 4f78737..f77c3a5 100644 --- a/database/CHANGELOG.md +++ b/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 diff --git a/database/macrostrat/database/transfer/__init__.py b/database/macrostrat/database/transfer/__init__.py index 57437a6..b278a1b 100644 --- a/database/macrostrat/database/transfer/__init__.py +++ b/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", +] diff --git a/database/macrostrat/database/transfer/cli.py b/database/macrostrat/database/transfer/cli.py new file mode 100644 index 0000000..eb385b2 --- /dev/null +++ b/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() diff --git a/database/macrostrat/database/transfer/dump_database.py b/database/macrostrat/database/transfer/dump_database.py index 0e28fdf..c3aaa33 100644 --- a/database/macrostrat/database/transfer/dump_database.py +++ b/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()) diff --git a/database/macrostrat/database/transfer/restore_database.py b/database/macrostrat/database/transfer/restore_database.py index a8327d4..0a16b5c 100644 --- a/database/macrostrat/database/transfer/restore_database.py +++ b/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: diff --git a/database/macrostrat/database/transfer/stream_utils.py b/database/macrostrat/database/transfer/stream_utils.py index a34f502..c4f3e2a 100644 --- a/database/macrostrat/database/transfer/stream_utils.py +++ b/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)""" diff --git a/database/pyproject.toml b/database/pyproject.toml index cd6e794..c50b7b8 100644 --- a/database/pyproject.toml +++ b/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"] diff --git a/database/tests/test_transfer.py b/database/tests/test_transfer.py new file mode 100644 index 0000000..98844b6 --- /dev/null +++ b/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"