FabledSteward/steward/backup/pgdata.py

"""A backup's rows as COPY text: read out of the dump by pg_restore, without a
database, and parsed or handed on as they are. The check (check.py) parses
them to compare; the in-place restore (restore.py) hands them to COPY
unparsed, so what goes back in is byte for byte what the dump holds.
"""
from __future__ import annotations

import asyncio
import re
from pathlib import Path

from .dump import BackupError

SCRIPT_DEADLINE_SECONDS = 5 * 60
_COPY = re.compile(r'^COPY ([^ ]+) \((.*)\) FROM stdin;$')
_ESCAPES = {ord("b"): b"\b", ord("f"): b"\f", ord("n"): b"\n", ord("r"): b"\r",
            ord("t"): b"\t", ord("v"): b"\v"}
_HEX = "0123456789abcdefABCDEF"


def async_url(database_url: str) -> str:
    from sqlalchemy.engine import make_url
    url = make_url(database_url)
    return str(url.set(drivername="postgresql+asyncpg")) if url.drivername == "postgresql" else database_url


async def restore_script(dump: Path) -> bytes:
    """The dump's rows as pg_restore would load them: its --data-only script.
    Nothing connects; pg_restore only writes the script."""
    try:
        proc = await asyncio.create_subprocess_exec(
            "pg_restore", "--data-only", "--no-owner", "--no-privileges", "--file=-", str(dump),
            stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
    except FileNotFoundError:
        raise BackupError("pg_restore isn't installed here.") from None
    try:
        out, err = await asyncio.wait_for(proc.communicate(), SCRIPT_DEADLINE_SECONDS)
    except TimeoutError:
        proc.kill()
        await proc.wait()
        raise BackupError(f"Reading the backup didn't finish within "
                          f"{SCRIPT_DEADLINE_SECONDS // 60} minutes.") from None
    if proc.returncode != 0:
        lines = [ln for ln in err.decode(errors="replace").splitlines() if ln.strip()]
        raise BackupError("The backup can't be read: " + (lines[-1] if lines else
                                                           f"pg_restore exited {proc.returncode}."))
    return out


def _ident(name: str) -> str:
    name = name.strip()
    return name[1:-1].replace('""', '"') if name.startswith('"') and name.endswith('"') else name


def script_blocks(script: bytes) -> dict[str, tuple[list[str], bytes]]:
    """Each table's (columns, COPY text) from a --data-only script, the text
    as it stands: one line per row, ready for COPY … FROM in text format."""
    out: dict[str, tuple[list[str], bytes]] = {}
    lines = iter(script.split(b"\n"))
    for line in lines:
        m = _COPY.match(line.decode("utf-8", errors="replace"))
        if not m:
            continue
        table = _ident(m.group(1).split(".")[-1])
        columns = [_ident(c) for c in m.group(2).split(",")]
        body = []
        for row in lines:
            if row == b"\\.":
                break
            body.append(row + b"\n")
        out[table] = (columns, b"".join(body))
    return out


def unescape(field: bytes) -> str | None:
    """One field of COPY's text format: \\N is NULL, backslash escapes are
    undone (named, octal \\ooo and hex \\xhh)."""
    if field == b"\\N":
        return None
    if b"\\" not in field:
        return field.decode("utf-8", errors="replace")
    out = bytearray()
    i, n = 0, len(field)
    while i < n:
        c = field[i]
        if c != 0x5C or i + 1 == n:
            out.append(c)
            i += 1
            continue
        nxt = field[i + 1]
        if nxt in _ESCAPES:
            out += _ESCAPES[nxt]
            i += 2
        elif 0x30 <= nxt <= 0x37:
            j = i + 1
            while j < n and j < i + 4 and 0x30 <= field[j] <= 0x37:
                j += 1
            out.append(int(field[i + 1:j], 8) & 0xFF)
            i = j
        elif nxt == ord("x") and i + 2 < n and chr(field[i + 2]) in _HEX:
            j = i + 2
            while j < n and j < i + 4 and chr(field[j]) in _HEX:
                j += 1
            out.append(int(field[i + 2:j], 16))
            i = j
        else:
            out.append(nxt)
            i += 2
    return bytes(out).decode("utf-8", errors="replace")


def copy_rows(data: bytes) -> list[list[str | None]]:
    """COPY text rows (no header, no terminator) as lists of fields."""
    return [[unescape(f) for f in line.split(b"\t")]
            for line in data.split(b"\n") if line and line != b"\\."]


def script_tables(script: bytes) -> dict[str, tuple[list[str], list[list[str | None]]]]:
    """Each table's (columns, rows) from a --data-only script."""
    return {t: (cols, copy_rows(raw)) for t, (cols, raw) in script_blocks(script).items()}


def script_from(blocks: dict[str, tuple[list[str], bytes]]) -> bytes:
    """A --data-only script holding `blocks`: the shape pg_restore writes,
    so a test can stand in for a dump."""
    out = []
    for table, (columns, raw) in blocks.items():
        out.append(f"COPY public.{table} ({', '.join(columns)}) FROM stdin;\n".encode())
        out.append(raw)
        out.append(b"\\.\n\n")
    return b"".join(out)