272 lines
10 KiB
Python
Executable file
272 lines
10 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
"""Generate a PostgreSQL restore script from a Teleo SQLite backup.
|
|
|
|
This is a shadow-restore converter for restore drills while the canonical
|
|
runtime database is still SQLite. It preserves tables and rows, infers safe
|
|
Postgres column types from actual SQLite storage classes, and avoids blind
|
|
string rewrites of `.dump` output.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import gzip
|
|
import json
|
|
import shutil
|
|
import sqlite3
|
|
import tempfile
|
|
from contextlib import contextmanager
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
PG_NUMERIC_TYPES = {"bigint", "double precision"}
|
|
|
|
|
|
def pg_quote_identifier(value: str) -> str:
|
|
return '"' + value.replace('"', '""') + '"'
|
|
|
|
|
|
def sqlite_quote_identifier(value: str) -> str:
|
|
return '"' + value.replace('"', '""') + '"'
|
|
|
|
|
|
def copy_escape(value: Any, pg_type: str, stats: dict[str, int]) -> str:
|
|
if value is None:
|
|
return r"\N"
|
|
if isinstance(value, bytes):
|
|
text = value.hex()
|
|
stats["blob_values_hex_encoded"] += 1
|
|
elif pg_type in PG_NUMERIC_TYPES:
|
|
text = str(value)
|
|
else:
|
|
text = str(value)
|
|
|
|
if "\x00" in text:
|
|
stats["nul_chars_replaced"] += text.count("\x00")
|
|
text = text.replace("\x00", r"\u0000")
|
|
|
|
return (
|
|
text.replace("\\", "\\\\")
|
|
.replace("\t", r"\t")
|
|
.replace("\n", r"\n")
|
|
.replace("\r", r"\r")
|
|
)
|
|
|
|
|
|
def infer_pg_type(declared_type: str | None, storage_counts: dict[str, int]) -> tuple[str, str]:
|
|
storage_types = {name for name, count in storage_counts.items() if count and name != "null"}
|
|
declared = (declared_type or "").upper()
|
|
|
|
if not storage_types:
|
|
if "INT" in declared:
|
|
return "bigint", "all_null_declared_integer"
|
|
if any(token in declared for token in ("REAL", "FLOA", "DOUB")):
|
|
return "double precision", "all_null_declared_real"
|
|
return "text", "all_null_text_default"
|
|
|
|
if storage_types == {"integer"}:
|
|
return "bigint", "sqlite_integer_storage"
|
|
if storage_types <= {"integer", "real"}:
|
|
return "double precision", "sqlite_numeric_storage"
|
|
return "text", "sqlite_text_or_mixed_storage"
|
|
|
|
|
|
def read_storage_counts(conn: sqlite3.Connection, table: str, column: str) -> dict[str, int]:
|
|
sql = (
|
|
f"select typeof({sqlite_quote_identifier(column)}) as storage_type, count(*) "
|
|
f"from {sqlite_quote_identifier(table)} group by storage_type"
|
|
)
|
|
return {str(row[0]): int(row[1]) for row in conn.execute(sql)}
|
|
|
|
|
|
def user_tables(conn: sqlite3.Connection) -> list[dict[str, str]]:
|
|
rows = conn.execute(
|
|
"""
|
|
select name, sql
|
|
from sqlite_master
|
|
where type = 'table'
|
|
and name not like 'sqlite_%'
|
|
order by name
|
|
"""
|
|
).fetchall()
|
|
return [{"name": str(row[0]), "sqlite_sql": row[1] or ""} for row in rows]
|
|
|
|
|
|
@contextmanager
|
|
def sqlite_database_path(path: Path):
|
|
if path.suffix != ".gz":
|
|
yield path
|
|
return
|
|
|
|
tmpdir = Path(tempfile.mkdtemp(prefix="teleo-sqlite-restore-"))
|
|
try:
|
|
db_path = tmpdir / "pipeline.db"
|
|
with gzip.open(path, "rb") as source, db_path.open("wb") as target:
|
|
shutil.copyfileobj(source, target)
|
|
yield db_path
|
|
finally:
|
|
shutil.rmtree(tmpdir, ignore_errors=True)
|
|
|
|
|
|
def build_table_metadata(conn: sqlite3.Connection, table: str) -> dict[str, Any]:
|
|
columns: list[dict[str, Any]] = []
|
|
for cid, name, declared_type, notnull, default_value, primary_key in conn.execute(
|
|
f"pragma table_info({sqlite_quote_identifier(table)})"
|
|
):
|
|
storage_counts = read_storage_counts(conn, table, name)
|
|
pg_type, type_reason = infer_pg_type(declared_type, storage_counts)
|
|
columns.append(
|
|
{
|
|
"cid": cid,
|
|
"name": name,
|
|
"sqlite_declared_type": declared_type or "",
|
|
"sqlite_notnull": bool(notnull),
|
|
"sqlite_default": default_value,
|
|
"sqlite_primary_key_position": primary_key,
|
|
"postgres_type": pg_type,
|
|
"postgres_type_reason": type_reason,
|
|
"sqlite_storage_counts": storage_counts,
|
|
}
|
|
)
|
|
row_count = conn.execute(f"select count(*) from {sqlite_quote_identifier(table)}").fetchone()[0]
|
|
return {"name": table, "row_count": int(row_count), "columns": columns}
|
|
|
|
|
|
def write_target_counts_sql(path: Path, schema: str, tables: list[dict[str, Any]]) -> None:
|
|
selects = []
|
|
for table in tables:
|
|
table_name = table["name"]
|
|
table_literal = table_name.replace("'", "''")
|
|
table_ref = f"{pg_quote_identifier(schema)}.{pg_quote_identifier(table_name)}"
|
|
selects.append(f"select '{table_literal}' as table_name, count(*)::bigint as row_count from {table_ref}")
|
|
body = "\nunion all\n".join(selects) or "select '' as table_name, 0::bigint as row_count where false"
|
|
path.write_text(
|
|
"copy (\n"
|
|
f"{body}\n"
|
|
"order by table_name\n"
|
|
") to stdout with (format csv, header true);\n"
|
|
)
|
|
|
|
|
|
def convert_sqlite_to_postgres(
|
|
sqlite_path: Path,
|
|
output_sql: Path,
|
|
summary_json: Path,
|
|
schema: str,
|
|
target_counts_sql: Path | None = None,
|
|
) -> dict[str, Any]:
|
|
output_sql.parent.mkdir(parents=True, exist_ok=True)
|
|
summary_json.parent.mkdir(parents=True, exist_ok=True)
|
|
if target_counts_sql:
|
|
target_counts_sql.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
with sqlite_database_path(sqlite_path) as db_path:
|
|
conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
|
|
conn.row_factory = sqlite3.Row
|
|
integrity = conn.execute("pragma integrity_check").fetchone()[0]
|
|
tables = []
|
|
table_names = [table["name"] for table in user_tables(conn)]
|
|
for table_name in table_names:
|
|
tables.append(build_table_metadata(conn, table_name))
|
|
|
|
stats = {"blob_values_hex_encoded": 0, "nul_chars_replaced": 0}
|
|
with output_sql.open("w", encoding="utf-8", newline="\n") as out:
|
|
out.write("\\set ON_ERROR_STOP on\n")
|
|
out.write("begin;\n")
|
|
out.write(f"drop schema if exists {pg_quote_identifier(schema)} cascade;\n")
|
|
out.write(f"create schema {pg_quote_identifier(schema)};\n")
|
|
out.write(f"set search_path to {pg_quote_identifier(schema)};\n")
|
|
|
|
for table in tables:
|
|
table_ref = f"{pg_quote_identifier(schema)}.{pg_quote_identifier(table['name'])}"
|
|
column_defs = [
|
|
f"{pg_quote_identifier(column['name'])} {column['postgres_type']}"
|
|
for column in table["columns"]
|
|
]
|
|
out.write(f"create table {table_ref} (\n " + ",\n ".join(column_defs) + "\n);\n")
|
|
|
|
if not table["columns"]:
|
|
continue
|
|
|
|
column_list = ", ".join(pg_quote_identifier(column["name"]) for column in table["columns"])
|
|
out.write(f"copy {table_ref} ({column_list}) from stdin with (format text, delimiter E'\\t', null '\\N');\n")
|
|
|
|
select_columns = ", ".join(sqlite_quote_identifier(column["name"]) for column in table["columns"])
|
|
for row in conn.execute(f"select {select_columns} from {sqlite_quote_identifier(table['name'])}"):
|
|
fields = [
|
|
copy_escape(row[column["name"]], column["postgres_type"], stats)
|
|
for column in table["columns"]
|
|
]
|
|
out.write("\t".join(fields))
|
|
out.write("\n")
|
|
out.write("\\.\n")
|
|
|
|
out.write("commit;\n")
|
|
out.write("analyze;\n")
|
|
|
|
if target_counts_sql:
|
|
write_target_counts_sql(target_counts_sql, schema, tables)
|
|
|
|
summary = {
|
|
"artifact": "teleo_sqlite_to_postgres_dump",
|
|
"generated_at_utc": datetime.now(UTC).isoformat(),
|
|
"source_sqlite_path": str(sqlite_path),
|
|
"source_integrity_check": integrity,
|
|
"schema": schema,
|
|
"table_count": len(tables),
|
|
"source_total_rows": sum(table["row_count"] for table in tables),
|
|
"tables": tables,
|
|
"conversion_notes": {
|
|
"mode": "shadow_restore",
|
|
"constraints_and_indexes": "not_recreated",
|
|
"blob_storage": "hex_encoded_text_when_needed",
|
|
"postgres_types": "inferred_from_sqlite_storage_classes",
|
|
},
|
|
"conversion_stats": stats,
|
|
"output_sql": str(output_sql),
|
|
"target_counts_sql": str(target_counts_sql) if target_counts_sql else None,
|
|
}
|
|
summary_json.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n")
|
|
return summary
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--sqlite-backup", required=True, type=Path, help="SQLite .db or .db.gz backup")
|
|
parser.add_argument("--output-sql", required=True, type=Path, help="Postgres import SQL path")
|
|
parser.add_argument("--summary-json", required=True, type=Path, help="Non-secret conversion summary JSON")
|
|
parser.add_argument("--schema", default="teleo_restore", help="Postgres schema to recreate")
|
|
parser.add_argument("--target-counts-sql", type=Path, help="Optional psql query file for target row counts")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
summary = convert_sqlite_to_postgres(
|
|
sqlite_path=args.sqlite_backup,
|
|
output_sql=args.output_sql,
|
|
summary_json=args.summary_json,
|
|
schema=args.schema,
|
|
target_counts_sql=args.target_counts_sql,
|
|
)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"artifact": summary["artifact"],
|
|
"source_integrity_check": summary["source_integrity_check"],
|
|
"schema": summary["schema"],
|
|
"table_count": summary["table_count"],
|
|
"source_total_rows": summary["source_total_rows"],
|
|
"output_sql": summary["output_sql"],
|
|
"target_counts_sql": summary["target_counts_sql"],
|
|
},
|
|
indent=2,
|
|
sort_keys=True,
|
|
)
|
|
)
|
|
return 0 if summary["source_integrity_check"] == "ok" else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|