sqlmirror 0.2.0__py3-none-any.whl

sqlmirror/db.py

@@ -0,0 +1,469 @@
+"""Database writers for the mirror (Postgres here, SQL Server in db_mssql.py).
+
+Every write goes through ``apply_records`` (issue snapshots and tombstones) or
+``apply_metadata``. Both are idempotent: applying the same data twice leaves
+the database unchanged, and an older snapshot never overwrites a newer one.
+That logic lives in ``Mirror``; each database class supplies the SQL.
+
+``open_mirror(dsn)`` picks the class from the connection string:
+  postgresql://user:pass@host:5432/db           -> PgMirror
+  sqlserver://user:pass@host:1433/db (mssql://) -> MssqlMirror
+"""
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from typing import Any, Iterable
+
+import psycopg
+from psycopg import sql
+
+from .mapper import SYNTHETIC_LINKTYPE_IDS, Mapper
+from .schema import ISSUE_TABLES, SCHEMA_VERSION, TABLE_BY_NAME, UPSERT_TABLES, Table, ddl
+from .util import parse_ts, to_int, utcnow
+
+
+@dataclass
+class ApplyStats:
+    snapshots: int = 0
+    tombstones: int = 0
+    skipped_stale: int = 0
+    rows: dict[str, int] = field(default_factory=lambda: defaultdict(int))
+
+    def add(self, other: "ApplyStats") -> None:
+        self.snapshots += other.snapshots
+        self.tombstones += other.tombstones
+        self.skipped_stale += other.skipped_stale
+        for k, v in other.rows.items():
+            self.rows[k] += v
+
+    def summary(self) -> str:
+        return (f"{self.snapshots} issues written, {self.tombstones} deleted, "
+                f"{self.skipped_stale} skipped (older than mirror)")
+
+
+def record_time(rec: dict[str, Any]) -> datetime | None:
+    return parse_ts(rec.get("fetchedAt") or rec.get("deletedAt"))
+
+
+def is_sqlserver_dsn(dsn: str) -> bool:
+    return dsn.strip().lower().split("://", 1)[0] in ("sqlserver", "mssql", "mssql+pytds", "mssql+pyodbc", "mssql+pymssql")
+
+
+def open_mirror(dsn: str, schema: str | None = None) -> "Mirror":
+    """Connect to the mirror database named by the connection string."""
+    if is_sqlserver_dsn(dsn):
+        from .db_mssql import MssqlMirror
+
+        return MssqlMirror(dsn, schema=schema)
+    return PgMirror(dsn, schema=schema)
+
+
+class Mirror:
+    """Shared apply logic. Subclasses provide the SQL for one database."""
+
+    dialect = ""  # "postgres" or "tsql"
+    # Smallest difference between two stored timestamps the database can show
+    # (SQL Server's datetime keeps 1/300 s), so reconcile doesn't chase rounding.
+    timestamp_resolution = timedelta(0)
+    conn: Any
+    dsn: str
+
+    def close(self) -> None:
+        self.conn.close()
+
+    def commit(self) -> None:
+        self.conn.commit()
+
+    def rollback(self) -> None:
+        self.conn.rollback()
+
+    # -------------------------------------------------------------- metadata
+    def apply_metadata(self, mapper: Mapper, meta: dict[str, Any]) -> dict[str, int]:
+        rows = mapper.map_metadata(meta)
+        counts: dict[str, int] = {}
+        with self.conn.cursor() as cur:
+            for t in UPSERT_TABLES:
+                if rows.get(t.name):
+                    counts[t.name] = self._upsert(cur, t, rows[t.name])
+            self._project_categories(cur, rows.get("project_category_assoc", []),
+                                     [r["id"] for r in rows.get("project", [])])
+        self.set_meta("metadata_fetched_at", meta.get("fetchedAt") or "")
+        self.conn.commit()
+        return counts
+
+    # ---------------------------------------------------------------- issues
+    def apply_records(self, mapper: Mapper, records: list[dict[str, Any]], source: str) -> ApplyStats:
+        """Apply issue snapshots and tombstones in one transaction."""
+        stats = ApplyStats()
+        # A deleted project becomes a tombstone for every issue the mirror has in it.
+        expanded: list[dict[str, Any]] = []
+        for rec in records:
+            if rec.get("type") == "project_deleted":
+                ids = self.project_issue_ids(to_int(rec.get("projectId")))
+                expanded.extend({"v": 1, "type": "deleted", "id": str(i), "deletedAt": rec.get("deletedAt"),
+                                 "reason": "project_deleted"} for i in ids)
+            else:
+                expanded.append(rec)
+        records = expanded
+        latest: dict[int, dict[str, Any]] = {}
+        for rec in records:
+            iid = to_int(rec.get("id") or (rec.get("issue") or {}).get("id"))
+            t = record_time(rec)
+            if iid is None or t is None:
+                continue
+            cur_best = latest.get(iid)
+            if cur_best is None or record_time(cur_best) < t:
+                latest[iid] = rec
+
+        states = self.issue_states(list(latest))
+        keep: dict[int, dict[str, Any]] = {}
+        for iid, rec in latest.items():
+            prev = states.get(iid)
+            if prev is not None and prev[0] is not None and prev[0] >= record_time(rec):
+                stats.skipped_stale += 1
+                continue
+            keep[iid] = rec
+        if not keep:
+            return stats
+
+        snapshots = {i: r for i, r in keep.items() if r.get("type") == "issue"}
+        tombstones = {i: r for i, r in keep.items() if r.get("type") == "deleted"}
+
+        merged: dict[str, list[dict[str, Any]]] = defaultdict(list)
+        for rec in snapshots.values():
+            for table, rows in mapper.map_issue(rec).items():
+                merged[table].extend(rows)
+
+        ids = list(keep)
+        with self.conn.cursor() as cur:
+            self._delete_issue_rows(cur, ids, list(tombstones))
+            for t in UPSERT_TABLES:
+                if merged.get(t.name):
+                    stats.rows[t.name] += self._upsert(cur, t, merged[t.name])
+            self._project_categories(cur, merged.get("project_category_assoc", []), [])
+            for t in ISSUE_TABLES:
+                rows = merged.get(t.name)
+                if not rows:
+                    continue
+                if t.name == "issuelink":
+                    stats.rows[t.name] += self._upsert(cur, t, rows, overwrite=True)
+                else:
+                    stats.rows[t.name] += self._insert_rows(cur, t, rows)
+            now = utcnow()
+            self._save_issue_states(
+                cur, [(iid, record_time(rec), rec.get("type") == "deleted", source, now) for iid, rec in keep.items()])
+        self.conn.commit()
+        stats.snapshots = len(snapshots)
+        stats.tombstones = len(tombstones)
+        return stats
+
+    @staticmethod
+    def _dedupe(t: Table, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Merge rows that share a primary key: later non-null values win."""
+        out: dict[tuple, dict[str, Any]] = {}
+        for r in rows:
+            key = tuple(r.get(c) for c in t.pk)
+            if key in out:
+                for k, v in r.items():
+                    if v is not None:
+                        out[key][k] = v
+            else:
+                out[key] = dict(r)
+        return list(out.values())
+
+    # ------------------------------------------------- per-database SQL hooks
+    def init_schema(self) -> None: raise NotImplementedError
+    def get_meta(self, key: str) -> str | None: raise NotImplementedError
+    def set_meta(self, key: str, value: Any) -> None: raise NotImplementedError
+    def load_registry(self, mapper: Mapper) -> None: raise NotImplementedError
+    def applied_files(self, keys: Iterable[str]) -> set[str]: raise NotImplementedError
+    def last_applied_file(self, prefix: str = "") -> str | None: raise NotImplementedError
+    def mark_file_applied(self, key: str, records: int) -> None: raise NotImplementedError
+    def issue_states(self, ids: list[int]) -> dict[int, tuple[datetime, bool]]: raise NotImplementedError
+    def project_issue_ids(self, project_id: int | None) -> list[int]: raise NotImplementedError
+    def _delete_issue_rows(self, cur: Any, ids: list[int], deleted_ids: list[int]) -> None: raise NotImplementedError
+    def _project_categories(self, cur: Any, assoc: list[dict[str, Any]], all_projects: list[int]) -> None:
+        raise NotImplementedError
+    def _upsert(self, cur: Any, t: Table, rows: list[dict[str, Any]], overwrite: bool = False) -> int:
+        raise NotImplementedError
+    def _insert_rows(self, cur: Any, t: Table, rows: list[dict[str, Any]]) -> int: raise NotImplementedError
+    def _save_issue_states(self, cur: Any, rows: list[tuple]) -> None: raise NotImplementedError
+    def log_run(self, kind: str, source: str, started: datetime, *, files: int = 0, stats: "ApplyStats | None" = None,
+                error: str | None = None) -> None: raise NotImplementedError
+    def recent_log(self, limit: int = 20) -> list[tuple]: raise NotImplementedError
+    def project_counts(self) -> dict[str, int]: raise NotImplementedError
+    def issue_updated(self, project_key: str) -> dict[int, datetime]: raise NotImplementedError
+    def table_counts(self) -> dict[str, int]: raise NotImplementedError
+    def freshness(self) -> tuple[datetime | None, datetime | None]: raise NotImplementedError
+    def read_table(self, t: Table) -> list[tuple]: raise NotImplementedError
+    def explain(self, sql_text: str) -> str: raise NotImplementedError
+    def query(self, sql_text: str, params: tuple = (), timeout_s: int = 60) -> list[tuple]: raise NotImplementedError
+    def sql_name(self, name: str) -> str:
+        """A table or column name quoted for this database when it needs quoting (Active Objects tables)."""
+        raise NotImplementedError
+
+
+class PgMirror(Mirror):
+    dialect = "postgres"
+
+    def __init__(self, dsn: str, schema: str | None = None):
+        self.dsn = dsn
+        self.schema = schema
+        self.conn = psycopg.connect(dsn, autocommit=False)
+        if schema:
+            with self.conn.cursor() as cur:
+                cur.execute(sql.SQL("CREATE SCHEMA IF NOT EXISTS {}").format(sql.Identifier(schema)))
+                cur.execute(sql.SQL("SET search_path TO {}").format(sql.Identifier(schema)))
+            self.conn.commit()
+
+    # ---------------------------------------------------------------- schema
+    def init_schema(self) -> None:
+        with self.conn.cursor() as cur:
+            cur.execute(ddl("postgres"))
+        self.conn.commit()
+        self.set_meta("schema_version", SCHEMA_VERSION)
+        self.conn.commit()
+
+    def get_meta(self, key: str) -> str | None:
+        with self.conn.cursor() as cur:
+            cur.execute("SELECT value FROM sqlmirror_meta WHERE key = %s", (key,))
+            row = cur.fetchone()
+        return row[0] if row else None
+
+    def set_meta(self, key: str, value: Any) -> None:
+        if not isinstance(value, str):
+            value = json.dumps(value, default=str)
+        with self.conn.cursor() as cur:
+            cur.execute(
+                "INSERT INTO sqlmirror_meta (key, value) VALUES (%s, %s) "
+                "ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value",
+                (key, value),
+            )
+
+    def load_registry(self, mapper: Mapper) -> None:
+        """Teach the mapper about custom fields already stored in the mirror."""
+        with self.conn.cursor() as cur:
+            cur.execute("SELECT id, customfieldtypekey, cfname FROM customfield")
+            for num, typekey, name in cur.fetchall():
+                mapper.register_custom_field(int(num), typekey, name)
+
+    # ------------------------------------------------------------ bookkeeping
+    def applied_files(self, keys: Iterable[str]) -> set[str]:
+        keys = list(keys)
+        if not keys:
+            return set()
+        with self.conn.cursor() as cur:
+            cur.execute("SELECT key FROM sqlmirror_applied_file WHERE key = ANY(%s)", (keys,))
+            return {r[0] for r in cur.fetchall()}
+
+    def last_applied_file(self, prefix: str = "") -> str | None:
+        with self.conn.cursor() as cur:
+            if prefix:
+                cur.execute("SELECT max(key) FROM sqlmirror_applied_file WHERE key LIKE %s", (prefix + "%",))
+            else:
+                cur.execute("SELECT max(key) FROM sqlmirror_applied_file")
+            row = cur.fetchone()
+        return row[0] if row else None
+
+    def mark_file_applied(self, key: str, records: int) -> None:
+        with self.conn.cursor() as cur:
+            cur.execute(
+                "INSERT INTO sqlmirror_applied_file (key, records, applied_at) VALUES (%s, %s, %s) "
+                "ON CONFLICT (key) DO NOTHING",
+                (key, records, utcnow()),
+            )
+
+    def issue_states(self, ids: list[int]) -> dict[int, tuple[datetime, bool]]:
+        if not ids:
+            return {}
+        with self.conn.cursor() as cur:
+            cur.execute(
+                "SELECT issue_id, fetched_at, deleted FROM sqlmirror_issue_state WHERE issue_id = ANY(%s)",
+                (ids,),
+            )
+            return {int(r[0]): (r[1], r[2]) for r in cur.fetchall()}
+
+    def project_issue_ids(self, project_id: int | None) -> list[int]:
+        with self.conn.cursor() as cur:
+            cur.execute("SELECT id FROM jiraissue WHERE project = %s", (project_id,))
+            return [int(r[0]) for r in cur.fetchall()]
+
+    def _save_issue_states(self, cur: psycopg.Cursor, rows: list[tuple]) -> None:
+        cur.executemany(
+            "INSERT INTO sqlmirror_issue_state (issue_id, fetched_at, deleted, source, applied_at) "
+            "VALUES (%s, %s, %s, %s, %s) ON CONFLICT (issue_id) DO UPDATE SET "
+            "fetched_at = EXCLUDED.fetched_at, deleted = EXCLUDED.deleted, "
+            "source = EXCLUDED.source, applied_at = EXCLUDED.applied_at",
+            rows,
+        )
+
+    def _delete_issue_rows(self, cur: psycopg.Cursor, ids: list[int], deleted_ids: list[int]) -> None:
+        cur.execute(
+            "DELETE FROM changeitem WHERE groupid IN (SELECT id FROM changegroup WHERE issueid = ANY(%s))", (ids,)
+        )
+        for table, col in (("changegroup", "issueid"), ("jiraaction", "issueid"), ("worklog", "issueid"),
+                           ("fileattachment", "issueid"), ("customfieldvalue", "issue"), ("label", "issue")):
+            cur.execute(sql.SQL("DELETE FROM {} WHERE {} = ANY(%s)").format(sql.Identifier(table),
+                                                                              sql.Identifier(col)), (ids,))
+        cur.execute(
+            "DELETE FROM nodeassociation WHERE source_node_entity = 'Issue' AND source_node_id = ANY(%s)", (ids,)
+        )
+        synthetic = list(SYNTHETIC_LINKTYPE_IDS)
+        # Ordinary links belong to both issues; parent/child links belong to the child.
+        cur.execute(
+            "DELETE FROM issuelink WHERE (source = ANY(%s) OR destination = ANY(%s)) AND NOT (linktype = ANY(%s))",
+            (ids, ids, synthetic),
+        )
+        cur.execute("DELETE FROM issuelink WHERE destination = ANY(%s) AND linktype = ANY(%s)", (ids, synthetic))
+        if deleted_ids:
+            cur.execute("DELETE FROM issuelink WHERE source = ANY(%s)", (deleted_ids,))
+        cur.execute("DELETE FROM jiraissue WHERE id = ANY(%s)", (ids,))
+
+    def _project_categories(self, cur: psycopg.Cursor, assoc: list[dict[str, Any]], all_projects: list[int]) -> None:
+        projects = sorted({a["project"] for a in assoc} | set(all_projects))
+        if not projects:
+            return
+        cur.execute(
+            "DELETE FROM nodeassociation WHERE source_node_entity = 'Project' AND association_type = 'ProjectCategory' "
+            "AND source_node_id = ANY(%s)",
+            (projects,),
+        )
+        seen = set()
+        for a in assoc:
+            key = (a["project"], a["category"])
+            if key in seen:
+                continue
+            seen.add(key)
+            cur.execute(
+                "INSERT INTO nodeassociation (source_node_id, source_node_entity, sink_node_id, sink_node_entity, "
+                "association_type, sequence) VALUES (%s, 'Project', %s, 'ProjectCategory', 'ProjectCategory', NULL) "
+                "ON CONFLICT DO NOTHING",
+                key,
+            )
+
+    # ---------------------------------------------------------- low level
+    def _upsert(self, cur: psycopg.Cursor, t: Table, rows: list[dict[str, Any]], overwrite: bool = False) -> int:
+        rows = self._dedupe(t, rows)
+        cols = t.column_names
+        non_pk = [c for c in cols if c not in t.pk]
+        col_sql = sql.SQL(", ").join(sql.Identifier(c) if t.quoted else sql.SQL(c) for c in cols)
+        ph = sql.SQL(", ").join(sql.Placeholder() for _ in cols)
+        pk_sql = sql.SQL(", ").join(sql.Identifier(c) if t.quoted else sql.SQL(c) for c in t.pk)
+        tname = sql.Identifier(t.name) if t.quoted else sql.SQL(t.name)
+
+        def ident(c: str) -> sql.Composable:
+            return sql.Identifier(c) if t.quoted else sql.SQL(c)
+
+        if overwrite:
+            sets = [sql.SQL("{c} = EXCLUDED.{c}").format(c=ident(c)) for c in non_pk]
+        else:
+            sets = [sql.SQL("{c} = COALESCE(EXCLUDED.{c}, {t}.{c})").format(c=ident(c), t=tname) for c in non_pk]
+        stmt = sql.SQL("INSERT INTO {t} ({cols}) VALUES ({ph}) ON CONFLICT ({pk}) DO UPDATE SET {sets}").format(
+            t=tname, cols=col_sql, ph=ph, pk=pk_sql, sets=sql.SQL(", ").join(sets)
+        )
+        cur.executemany(stmt, [tuple(r.get(c) for c in cols) for r in rows])
+        return len(rows)
+
+    def _insert_rows(self, cur: psycopg.Cursor, t: Table, rows: list[dict[str, Any]]) -> int:
+        return self._copy(cur, t, rows)
+
+    def _copy(self, cur: psycopg.Cursor, t: Table, rows: list[dict[str, Any]]) -> int:
+        rows = self._dedupe(t, rows)
+        cols = t.column_names
+        stmt = sql.SQL("COPY {t} ({cols}) FROM STDIN").format(
+            t=sql.Identifier(t.name) if t.quoted else sql.SQL(t.name),
+            cols=sql.SQL(", ").join(sql.Identifier(c) if t.quoted else sql.SQL(c) for c in cols),
+        )
+        with cur.copy(stmt) as copy:
+            for r in rows:
+                copy.write_row(tuple(r.get(c) for c in cols))
+        return len(rows)
+
+    # ------------------------------------------------------------- sync log
+    def log_run(self, kind: str, source: str, started: datetime, *, files: int = 0, stats: "ApplyStats | None" = None,
+                error: str | None = None) -> None:
+        """Record one apply / backfill / reconcile run in sqlmirror_sync_log."""
+        self.conn.rollback()  # in case the run failed mid-transaction
+        with self.conn.cursor() as cur:
+            cur.execute(
+                "INSERT INTO sqlmirror_sync_log (id, started_at, finished_at, kind, source, files, issues, tombstones, "
+                "skipped, error) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)",
+                (int(started.timestamp() * 1_000_000), started, utcnow(), kind, source, files,
+                 stats.snapshots if stats else 0, stats.tombstones if stats else 0,
+                 stats.skipped_stale if stats else 0, error),
+            )
+        self.conn.commit()
+
+    def recent_log(self, limit: int = 20) -> list[tuple]:
+        with self.conn.cursor() as cur:
+            cur.execute("SELECT started_at, finished_at, kind, source, files, issues, tombstones, skipped, error "
+                        "FROM sqlmirror_sync_log ORDER BY started_at DESC LIMIT %s", (limit,))
+            return cur.fetchall()
+
+    # -------------------------------------------------------------- queries
+    def project_counts(self) -> dict[str, int]:
+        with self.conn.cursor() as cur:
+            cur.execute(
+                "SELECT p.pkey, count(i.id) FROM project p LEFT JOIN jiraissue i ON i.project = p.id "
+                "GROUP BY p.pkey ORDER BY p.pkey"
+            )
+            return {r[0]: int(r[1]) for r in cur.fetchall()}
+
+    def issue_updated(self, project_key: str) -> dict[int, datetime]:
+        with self.conn.cursor() as cur:
+            cur.execute(
+                "SELECT i.id, i.updated FROM jiraissue i JOIN project p ON p.id = i.project WHERE p.pkey = %s",
+                (project_key,),
+            )
+            return {int(r[0]): r[1] for r in cur.fetchall()}
+
+    def table_counts(self) -> dict[str, int]:
+        out = {}
+        with self.conn.cursor() as cur:
+            for name, t in TABLE_BY_NAME.items():
+                cur.execute(sql.SQL("SELECT count(*) FROM {}").format(
+                    sql.Identifier(t.name) if t.quoted else sql.SQL(t.name)))
+                out[name] = int(cur.fetchone()[0])
+        return out
+
+    def freshness(self) -> tuple[datetime | None, datetime | None]:
+        """(last write, newest data fetched from Jira)."""
+        with self.conn.cursor() as cur:
+            cur.execute("SELECT max(applied_at), max(fetched_at) FROM sqlmirror_issue_state")
+            row = cur.fetchone()
+        return row[0], row[1]
+
+    def sql_name(self, name: str) -> str:
+        t = TABLE_BY_NAME.get(name)
+        return f'"{name}"' if (t is not None and t.quoted) or name != name.lower() else name
+
+    def read_table(self, t: Table) -> list[tuple]:
+        ident = (lambda c: sql.Identifier(c)) if t.quoted else (lambda c: sql.SQL(c))
+        q = sql.SQL("SELECT {} FROM {}").format(sql.SQL(", ").join(ident(c) for c in t.column_names), ident(t.name))
+        with self.conn.cursor() as cur:
+            cur.execute(q)
+            return cur.fetchall()
+
+    def explain(self, sql_text: str) -> str:
+        """"ok" if the statement compiles against the mirror, else the error."""
+        try:
+            with self.conn.cursor() as cur:
+                cur.execute("EXPLAIN " + sql_text)
+            self.conn.rollback()
+            return "ok"
+        except psycopg.Error as exc:
+            self.conn.rollback()
+            return str(exc).strip().splitlines()[0]
+
+    def query(self, sql_text: str, params: tuple = (), timeout_s: int = 60) -> list[tuple]:
+        try:
+            with self.conn.cursor() as cur:
+                cur.execute(f"SET LOCAL statement_timeout = '{int(timeout_s)}s'")
+                cur.execute(sql_text, params or None)
+                rows = cur.fetchall() if cur.description else []
+            return rows
+        finally:
+            self.conn.rollback()