tools/copy_database.py (2723 bytes)
1 """Copy every row of the app database to another database. 2 3 The web server needs this once, when it turns on burst workers: the workers 4 share a Postgres database with it, and until then the data is in SQLite (see 5 infra/burst/README.md). 6 7 python -m tools.copy_database sqlite:////opt/subplz-web/data/subplz.db \\ 8 "postgresql+psycopg://user:password@host:port/db?sslmode=require" 9 10 The target gets the tables of the app models first. The copy refuses a target 11 that has rows already, so it never mixes two databases. 12 """ 13 14 from __future__ import annotations 15 16 import sys 17 from datetime import datetime, timezone 18 19 from sqlalchemy import DateTime, create_engine, func, insert, select 20 21 # The runbook runs this tool in a new interpreter. There, Base.metadata holds 22 # only the tables of the modules that Python imported. backend.main imports 23 # every model module, as the app does, so the copy gets every table. 24 import backend.main # noqa: F401 25 from backend.db import Base 26 27 BATCH = 1000 28 29 30 def _as_utc(table, row: dict) -> dict: 31 """SQLite keeps no time zone. The app writes UTC, so say so for Postgres.""" 32 for column in table.columns: 33 value = row.get(column.name) 34 if (isinstance(column.type, DateTime) and column.type.timezone 35 and isinstance(value, datetime) and value.tzinfo is None): 36 row[column.name] = value.replace(tzinfo=timezone.utc) 37 return row 38 39 40 def copy_database(source_url: str, target_url: str) -> dict[str, int]: 41 """Copy each table of the app from `source_url` to `target_url`. 42 43 Returns the number of rows copied for each table. 44 """ 45 source = create_engine(source_url) 46 target = create_engine(target_url) 47 Base.metadata.create_all(target) 48 counts: dict[str, int] = {} 49 with source.connect() as src, target.begin() as dst: 50 for table in Base.metadata.sorted_tables: 51 if dst.execute(select(func.count()).select_from(table)).scalar(): 52 raise RuntimeError(f"table {table.name} has rows in the target already") 53 # sorted_tables puts each table after the tables that it refers to. 54 for table in Base.metadata.sorted_tables: 55 rows = [_as_utc(table, dict(r._mapping)) for r in src.execute(select(table))] 56 for start in range(0, len(rows), BATCH): 57 dst.execute(insert(table), rows[start:start + BATCH]) 58 counts[table.name] = len(rows) 59 return counts 60 61 62 def main(argv: list[str]) -> int: 63 if len(argv) != 2: 64 print(__doc__, file=sys.stderr) 65 return 2 66 for table, count in copy_database(argv[0], argv[1]).items(): 67 print(f"{table}: {count}") 68 return 0 69 70 71 if __name__ == "__main__": 72 sys.exit(main(sys.argv[1:]))