commit f61e9113ae1f49bfba30fa6804e19410c4adbde5 equwal <truex@equwal.com> 2026-09-22 16:45:22 -0700 Add a tool that copies the app database to another database Burst workers share one Postgres database with the web server, and the web server keeps its data in SQLite today. tools/copy_database.py copies every row of every app table from one database URL to another: SQLite to Postgres when burst mode goes on, and back when it goes off. It makes the tables from the app models, marks the naive SQLite times as UTC, and refuses a target that has rows already. The test copies random databases and compares every row.
tests/test_copy_database.py | 83 +++++++++++++++++++++++++++++++++++++++++++++ tools/copy_database.py | 68 +++++++++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+)
diff --git a/tests/test_copy_database.py b/tests/test_copy_database.py new file mode 100644 index 0000000..bb70a67 --- /dev/null +++ b/tests/test_copy_database.py @@ -0,0 +1,83 @@ +"""tools/copy_database.py: a copy has every row of every table, unchanged.""" + +from __future__ import annotations + +from datetime import date, datetime, timedelta, timezone + +import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st +from sqlalchemy import create_engine, select +from sqlalchemy.orm import Session + +from backend.db import Account, Artifact, Base, Job, JobStatus, Purchase +from tools.copy_database import copy_database + +T0 = datetime(2026, 9, 22, tzinfo=timezone.utc) +text = st.text(st.characters(codec="utf-8", exclude_categories=("Cs", "Cc")), max_size=20) + + +@st.composite +def databases(draw): + accounts = [] + for a in range(draw(st.integers(1, 4))): + account = Account( + id=f"acct_{a}", device_token=f"dev_{a}", + email=draw(st.none() | st.just(f"reader{a}@example.com")), + purchased_credits=draw(st.integers(0, 500)), + daily_credit_on=draw(st.none() | st.dates(date(2026, 1, 1), date(2027, 1, 1))), + created_at=T0 + timedelta(seconds=draw(st.integers(0, 10**6))), + ) + accounts.append(account) + for j in range(draw(st.integers(0, 3))): + job_id = f"job_{a}_{j}" + accounts.append(Job( + id=job_id, account_id=account.id, status=draw(st.sampled_from(list(JobStatus))), + language="ja", splitter="pysbd", model="tiny", + audio_filename=draw(text), text_filename=draw(text), + audio_duration_seconds=draw(st.none() | st.floats(0, 1e5)), + progress=draw(st.floats(0, 1)), error=draw(st.none() | text), + )) + accounts.append(Artifact(id=f"art_{a}_{j}", job_id=job_id, kind="srt", + filename=draw(text), storage_key=f"{job_id}/out.srt")) + if draw(st.booleans()): + accounts.append(Purchase(id=f"buy_{a}", account_id=account.id, plan_id="pack10", + credits=10, amount_cents=499, stripe_session_id=f"cs_{a}")) + return accounts + + +def rows(url: str) -> dict[str, list[dict]]: + engine = create_engine(url) + with engine.connect() as c: + return {t.name: sorted((dict(r._mapping) for r in c.execute(select(t))), key=repr) + for t in Base.metadata.sorted_tables} + + +@settings(max_examples=40, deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture]) +@given(objects=databases()) +def test_the_copy_has_every_row_unchanged(tmp_path_factory, objects): + tmp = tmp_path_factory.mktemp("copy") + source, target = f"sqlite:///{tmp / 'source.db'}", f"sqlite:///{tmp / 'target.db'}" + engine = create_engine(source) + Base.metadata.create_all(engine) + with Session(engine) as s: + s.add_all(objects) + s.commit() + + counts = copy_database(source, target) + + assert rows(target) == rows(source) + assert counts["accounts"] == len([o for o in objects if isinstance(o, Account)]) + + +def test_the_copy_refuses_a_target_with_rows(tmp_path): + source, target = f"sqlite:///{tmp_path / 's.db'}", f"sqlite:///{tmp_path / 't.db'}" + for url in (source, target): + engine = create_engine(url) + Base.metadata.create_all(engine) + with Session(engine) as s: + s.add(Account(id="acct_x", device_token="dev_x")) + s.commit() + with pytest.raises(RuntimeError, match="accounts"): + copy_database(source, target) diff --git a/tools/copy_database.py b/tools/copy_database.py new file mode 100644 index 0000000..b59be97 --- /dev/null +++ b/tools/copy_database.py @@ -0,0 +1,68 @@ +"""Copy every row of the app database to another database. + +The web server needs this once, when it turns on burst workers: the workers +share a Postgres database with it, and until then the data is in SQLite (see +infra/burst/README.md). + + python -m tools.copy_database sqlite:////opt/subplz-web/data/subplz.db \\ + "postgresql+psycopg://user:password@host:port/db?sslmode=require" + +The target gets the tables of the app models first. The copy refuses a target +that has rows already, so it never mixes two databases. +""" + +from __future__ import annotations + +import sys +from datetime import datetime, timezone + +from sqlalchemy import DateTime, create_engine, func, insert, select + +from backend.db import Base + +BATCH = 1000 + + +def _as_utc(table, row: dict) -> dict: + """SQLite keeps no time zone. The app writes UTC, so say so for Postgres.""" + for column in table.columns: + value = row.get(column.name) + if (isinstance(column.type, DateTime) and column.type.timezone + and isinstance(value, datetime) and value.tzinfo is None): + row[column.name] = value.replace(tzinfo=timezone.utc) + return row + + +def copy_database(source_url: str, target_url: str) -> dict[str, int]: + """Copy each table of the app from `source_url` to `target_url`. + + Returns the number of rows copied for each table. + """ + source = create_engine(source_url) + target = create_engine(target_url) + Base.metadata.create_all(target) + counts: dict[str, int] = {} + with source.connect() as src, target.begin() as dst: + for table in Base.metadata.sorted_tables: + if dst.execute(select(func.count()).select_from(table)).scalar(): + raise RuntimeError(f"table {table.name} has rows in the target already") + # sorted_tables puts each table after the tables that it refers to. + for table in Base.metadata.sorted_tables: + rows = [_as_utc(table, dict(r._mapping)) for r in src.execute(select(table))] + for start in range(0, len(rows), BATCH): + dst.execute(insert(table), rows[start:start + BATCH]) + counts[table.name] = len(rows) + return counts + + +def main(argv: list[str]) -> int: + if len(argv) != 2: + print(__doc__, file=sys.stderr) + return 2 + for table, count in copy_database(argv[0], argv[1]).items(): + print(f"{table}: {count}") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:]))