sqlmirror 0.2.0__py3-none-any.whl

sqlmirror/db_mssql.py

@@ -0,0 +1,530 @@
+"""SQL Server writer for the mirror (SQL Server 2016+, Azure SQL Database, Azure SQL Managed Instance).
+
+Connection string:
+
+  sqlserver://user:password@host:1433/database            local or on-premises
+  sqlserver://user:password@host:1433/database?encrypt=true  Azure SQL (TLS, verified with the system CAs)
+
+Other options (query string): instance=NAME (named instance through SQL Browser,
+no port), cafile=/path/ca.pem, trustservercertificate=true (skip the host-name
+check), schema=NAME (default dbo), login_timeout=SECONDS, timezone=IANA name
+(see below).
+
+The driver is python-tds (pure Python, no ODBC driver needed):
+  pip install "sqlmirror[sqlserver]"
+
+Differences from the Postgres writer, all invisible to report SQL:
+  - Upserts are MERGE statements; bulk loads are multi-row INSERTs.
+  - Every value is CAST to its column type inside the statement, so a batch
+    mixing NULLs, numbers and text never trips SQL Server's row-type rules.
+  - Data Center tables use datetime, as Data Center does on SQL Server. Times
+    are stored in UTC, or with ?timezone=America/New_York in that zone's local
+    time: Data Center stored the Jira server's local time, so reports that use
+    GETDATE() or group by day match the old numbers only with the old server's
+    zone. Dates before 1753, which datetime can't hold (typos such as year
+    0202), are stored as NULL.
+  - Mirror bookkeeping times use datetime2(6), so "newest snapshot wins" is exact.
+"""
+from __future__ import annotations
+
+import json
+import math
+from datetime import date, datetime, timedelta, timezone, tzinfo
+from decimal import Decimal
+from typing import Any, Callable, Iterable, Iterator
+from urllib.parse import parse_qs, unquote, urlsplit
+
+from .db import ApplyStats, Mirror
+from .mapper import SYNTHETIC_LINKTYPE_IDS, Mapper
+from .schema import SCHEMA_VERSION, TABLE_BY_NAME, TABLES, Table, column_type, ddl_statements, tsql_ident
+from .util import parse_ts, utcnow
+
+MAX_PARAMS = 2000  # SQL Server allows 2100 parameters per statement
+MAX_ROWS = 1000  # and 1000 rows per VALUES list
+ID_CHUNK = 1000  # ids inlined per IN (...) list
+DATETIME_MIN = datetime(1753, 1, 1)
+
+_TRUE = ("1", "true", "yes", "on", "strict", "mandatory")
+
+
+def parse_dsn(dsn: str) -> dict[str, Any]:
+    """sqlserver://user:pass@host:port/db?options -> keyword arguments for pytds.connect, plus 'schema'."""
+    u = urlsplit(dsn.strip())
+    q = {k.lower(): v[-1] for k, v in parse_qs(u.query, keep_blank_values=True).items()}
+    host = u.hostname or "localhost"
+    instance = q.get("instance") or q.get("instancename")
+    port = None if instance else (u.port or int(q.get("port") or 1433))
+    database = unquote(u.path.lstrip("/")) or q.get("database") or "master"
+    kwargs: dict[str, Any] = {
+        "server": f"{host}\\{instance}" if instance else host,
+        "port": port,
+        "database": database,
+        "user": unquote(u.username) if u.username else None,
+        "password": unquote(u.password) if u.password else None,
+        "login_timeout": float(q.get("login_timeout") or q.get("connect_timeout") or 15),
+        "appname": "sqlmirror",
+        "autocommit": False,
+    }
+    cafile = q.get("cafile")
+    if not cafile and q.get("encrypt", "").lower() in _TRUE:
+        import certifi
+
+        cafile = certifi.where()
+    if cafile:
+        kwargs["cafile"] = cafile
+        kwargs["validate_host"] = q.get("trustservercertificate", "").lower() not in _TRUE
+    kwargs["schema"] = q.get("schema") or None
+    kwargs["timezone"] = q.get("timezone") or q.get("tz") or None
+    return kwargs
+
+
+def _utc(v: datetime | None) -> datetime | None:
+    """A value read back from SQL Server (naive, UTC) as an aware datetime."""
+    if isinstance(v, datetime) and v.tzinfo is None:
+        return v.replace(tzinfo=timezone.utc)
+    return v
+
+
+def _to_naive_utc(v: Any, tz: tzinfo = timezone.utc) -> datetime | None:
+    """Aware or ISO value -> naive datetime in ``tz`` (UTC unless the mirror is set to a local zone)."""
+    if v is None:
+        return None
+    if not isinstance(v, datetime):
+        v = parse_ts(v) if not isinstance(v, date) else datetime(v.year, v.month, v.day, tzinfo=timezone.utc)
+        if v is None:
+            return None
+    if v.tzinfo is not None:
+        v = v.astimezone(tz).replace(tzinfo=None)
+    return v
+
+
+# Date-only Jira fields: stored as that calendar day's midnight, never shifted between zones.
+DATE_ONLY_COLUMNS = {("jiraissue", "duedate"), ("projectversion", "startdate"), ("projectversion", "releasedate")}
+
+
+def _utf16_len(v: str) -> int:
+    """Length as SQL Server counts nvarchar: UTF-16 code units (an emoji counts 2)."""
+    return len(v) + sum(1 for ch in v if ord(ch) > 0xFFFF)
+
+
+def _converter(sql_type: str, tz: tzinfo = timezone.utc, *, where: str = "", date_only: bool = False
+               ) -> Callable[[Any], Any]:
+    base = sql_type.split("(")[0].lower()
+    if base == "datetime":
+        def conv_dt(v: Any) -> Any:
+            if date_only or (isinstance(v, datetime) and v.tzinfo is not None and _is_utc_midnight(v)
+                             and tz is not timezone.utc and where.endswith(".datevalue")):
+                v = _to_naive_utc(v)  # a calendar date: keep the day as Jira shows it
+            else:
+                v = _to_naive_utc(v, tz)
+            return None if v is None or v < DATETIME_MIN else v
+        return conv_dt
+    if base == "datetime2":
+        return _to_naive_utc
+    if base == "bit":
+        return lambda v: None if v is None else bool(v)
+    if base == "float":
+        def conv_float(v: Any) -> Any:
+            if v is None or v == "":
+                return None
+            f = float(v)
+            return None if math.isnan(f) or math.isinf(f) else f  # SQL Server has no NaN / infinity
+        return conv_float
+    if base in ("numeric", "bigint", "int"):
+        def conv_num(v: Any) -> Any:
+            if v is None or v == "":
+                return None
+            if isinstance(v, bool):
+                return int(v)
+            if isinstance(v, float):
+                return None if math.isnan(v) or math.isinf(v) else int(v)
+            return v
+        return conv_num
+    # nvarchar(n): CAST would silently cut a longer value; refuse it, as Postgres does.
+    limit = sql_type[sql_type.find("(") + 1:-1] if "(" in sql_type else "max"
+    if limit.lower() == "max":
+        return lambda v: v
+    n = int(limit)
+
+    def conv_str(v: Any) -> Any:
+        if isinstance(v, str) and len(v) > n // 2 and _utf16_len(v) > n:
+            raise ValueError(f"Value too long for {where} ({sql_type}): {_utf16_len(v)} characters, "
+                             f"starting {v[:40]!r}")
+        return v
+    return conv_str
+
+
+def _is_utc_midnight(v: datetime) -> bool:
+    u = v.astimezone(timezone.utc)
+    return (u.hour, u.minute, u.second, u.microsecond) == (0, 0, 0, 0)
+
+
+class MssqlMirror(Mirror):
+    dialect = "tsql"
+    # datetime keeps 1/300 s: allow for rounding when comparing with Jira's milliseconds.
+    timestamp_resolution = timedelta(milliseconds=4)
+
+    def __init__(self, dsn: str, schema: str | None = None, *, connect: Callable[..., Any] | None = None):
+        kwargs = parse_dsn(dsn)
+        self.dsn = dsn
+        self.schema = schema or kwargs.pop("schema") or "dbo"
+        kwargs.pop("schema", None)
+        self.tz_name = kwargs.pop("timezone") or "UTC"
+        if self.tz_name.upper() == "UTC":
+            self.tz_name, self.tz = "UTC", timezone.utc
+        else:
+            from zoneinfo import ZoneInfo
+
+            self.tz = ZoneInfo(self.tz_name)
+        if connect is None:
+            try:
+                import pytds
+            except ImportError:  # pragma: no cover - depends on the install
+                raise SystemExit('SQL Server support needs the python-tds driver: pip install "sqlmirror[sqlserver]"')
+            connect = pytds.connect
+        self.conn = connect(**kwargs)
+        # Column types and value converters per table, from the shared table definitions.
+        self._cols: dict[str, list[tuple[str, str, Callable[[Any], Any]]]] = {
+            t.name: [(c, column_type(t, c, typ, "tsql"),
+                      _converter(column_type(t, c, typ, "tsql"), self.tz, where=f"{t.name}.{c}",
+                                 date_only=(t.name, c) in DATE_ONLY_COLUMNS))
+                     for c, typ in t.columns]
+            for t in TABLES
+        }
+
+    # ----------------------------------------------------------- naming
+    def t(self, name: str) -> str:
+        return f"{tsql_ident(self.schema)}.{tsql_ident(name)}"
+
+    def sql_name(self, name: str) -> str:
+        return tsql_ident(name)
+
+    def rollback(self) -> None:
+        try:
+            self.conn.rollback()
+        except Exception:  # noqa: BLE001 - the server may already have rolled back
+            pass
+
+    # ---------------------------------------------------------------- schema
+    def init_schema(self) -> None:
+        with self.conn.cursor() as cur:
+            for stmt in ddl_statements("tsql", self.schema):
+                cur.execute(stmt)
+        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(f"SELECT [value] FROM {self.t('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:
+            self._upsert(cur, TABLE_BY_NAME["sqlmirror_meta"], [{"key": key, "value": value}], overwrite=True)
+
+    def load_registry(self, mapper: Mapper) -> None:
+        with self.conn.cursor() as cur:
+            cur.execute(f"SELECT [id], [customfieldtypekey], [cfname] FROM {self.t('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)
+        out: set[str] = set()
+        with self.conn.cursor() as cur:
+            for i in range(0, len(keys), MAX_PARAMS):
+                chunk = keys[i:i + MAX_PARAMS]
+                cur.execute(f"SELECT [key] FROM {self.t('sqlmirror_applied_file')} WHERE [key] IN "
+                            f"({', '.join(['%s'] * len(chunk))})", tuple(chunk))
+                out.update(r[0] for r in cur.fetchall())
+        # The key column is case-insensitive on most servers; only exact matches count.
+        return out & set(keys)
+
+    def last_applied_file(self, prefix: str = "") -> str | None:
+        with self.conn.cursor() as cur:
+            if prefix:
+                cur.execute(f"SELECT MAX([key]) FROM {self.t('sqlmirror_applied_file')} WHERE LEFT([key], %s) = %s",
+                            (len(prefix), prefix))
+            else:
+                cur.execute(f"SELECT MAX([key]) FROM {self.t('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(
+                f"INSERT INTO {self.t('sqlmirror_applied_file')} ([key], [records], [applied_at]) "
+                f"SELECT %s, %s, %s WHERE NOT EXISTS (SELECT 1 FROM {self.t('sqlmirror_applied_file')} "
+                f"WHERE [key] = %s)",
+                (key, records, _to_naive_utc(utcnow()), key),
+            )
+
+    def issue_states(self, ids: list[int]) -> dict[int, tuple[datetime, bool]]:
+        out: dict[int, tuple[datetime, bool]] = {}
+        with self.conn.cursor() as cur:
+            for chunk in _id_chunks(ids):
+                cur.execute(f"SELECT [issue_id], [fetched_at], [deleted] FROM {self.t('sqlmirror_issue_state')} "
+                            f"WHERE [issue_id] IN ({chunk})")
+                for r in cur.fetchall():
+                    out[int(r[0])] = (_utc(r[1]), bool(r[2]) if r[2] is not None else None)
+        return out
+
+    def project_issue_ids(self, project_id: int | None) -> list[int]:
+        if project_id is None:
+            return []
+        with self.conn.cursor() as cur:
+            cur.execute(f"SELECT [id] FROM {self.t('jiraissue')} WHERE [project] = {int(project_id)}")
+            return [int(r[0]) for r in cur.fetchall()]
+
+    def _save_issue_states(self, cur: Any, rows: list[tuple]) -> None:
+        cols = ("issue_id", "fetched_at", "deleted", "source", "applied_at")
+        self._upsert(cur, TABLE_BY_NAME["sqlmirror_issue_state"], [dict(zip(cols, r)) for r in rows], overwrite=True)
+
+    # --------------------------------------------------------------- deletes
+    def _delete_issue_rows(self, cur: Any, ids: list[int], deleted_ids: list[int]) -> None:
+        synthetic = ", ".join(str(int(x)) for x in sorted(SYNTHETIC_LINKTYPE_IDS))
+        T = self.t
+        for chunk in _id_chunks(ids):
+            cur.execute(f"DELETE FROM {T('changeitem')} WHERE [groupid] IN "
+                        f"(SELECT [id] FROM {T('changegroup')} WHERE [issueid] IN ({chunk}))")
+            for table, col in (("changegroup", "issueid"), ("jiraaction", "issueid"), ("worklog", "issueid"),
+                               ("fileattachment", "issueid"), ("customfieldvalue", "issue"), ("label", "issue")):
+                cur.execute(f"DELETE FROM {T(table)} WHERE {tsql_ident(col)} IN ({chunk})")
+            cur.execute(f"DELETE FROM {T('nodeassociation')} WHERE [source_node_entity] = N'Issue' "
+                        f"AND [source_node_id] IN ({chunk})")
+            # Ordinary links belong to both issues; parent/child links belong to the child.
+            cur.execute(f"DELETE FROM {T('issuelink')} WHERE ([source] IN ({chunk}) OR [destination] IN ({chunk})) "
+                        f"AND [linktype] NOT IN ({synthetic})")
+            cur.execute(f"DELETE FROM {T('issuelink')} WHERE [destination] IN ({chunk}) AND [linktype] IN ({synthetic})")
+        for chunk in _id_chunks(deleted_ids):
+            cur.execute(f"DELETE FROM {T('issuelink')} WHERE [source] IN ({chunk})")
+        for chunk in _id_chunks(ids):
+            cur.execute(f"DELETE FROM {T('jiraissue')} WHERE [id] IN ({chunk})")
+
+    def _project_categories(self, cur: Any, 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
+        na = self.t("nodeassociation")
+        for chunk in _id_chunks(projects):
+            cur.execute(f"DELETE FROM {na} WHERE [source_node_entity] = N'Project' "
+                        f"AND [association_type] = N'ProjectCategory' AND [source_node_id] IN ({chunk})")
+        seen = set()
+        for a in assoc:
+            key = (int(a["project"]), int(a["category"]))
+            if key in seen:
+                continue
+            seen.add(key)
+            cur.execute(
+                f"INSERT INTO {na} ([source_node_id], [source_node_entity], [sink_node_id], [sink_node_entity], "
+                f"[association_type], [sequence]) SELECT {key[0]}, N'Project', {key[1]}, N'ProjectCategory', "
+                f"N'ProjectCategory', NULL WHERE NOT EXISTS (SELECT 1 FROM {na} WHERE [source_node_id] = {key[0]} "
+                f"AND [source_node_entity] = N'Project' AND [sink_node_id] = {key[1]} "
+                f"AND [sink_node_entity] = N'ProjectCategory' AND [association_type] = N'ProjectCategory')"
+            )
+
+    # ---------------------------------------------------------- low level
+    def _values(self, t: Table, rows: list[dict[str, Any]]) -> Iterator[tuple[list[str], list[Any]]]:
+        """Batches of "(CAST(.. AS type), ...)" row literals and their parameters.
+
+        Integers are inlined (they are ints, so nothing to escape); NULLs are
+        inlined; everything else is a parameter. Batches stay under SQL
+        Server's parameter and row limits.
+        """
+        cols = self._cols[t.name]
+        batch: list[str] = []
+        params: list[Any] = []
+        for r in rows:
+            parts: list[str] = []
+            row_params: list[Any] = []
+            for c, typ, conv in cols:
+                v = conv(r.get(c))
+                if v is None:
+                    parts.append(f"CAST(NULL AS {typ})")
+                elif isinstance(v, bool):
+                    parts.append(f"CAST({1 if v else 0} AS {typ})")
+                elif isinstance(v, int):
+                    parts.append(f"CAST({int(v)} AS {typ})")
+                else:
+                    parts.append(f"CAST(%s AS {typ})")
+                    row_params.append(v)
+            if batch and (len(batch) >= MAX_ROWS or len(params) + len(row_params) > MAX_PARAMS):
+                yield batch, params
+                batch, params = [], []
+            batch.append("(" + ", ".join(parts) + ")")
+            params.extend(row_params)
+        if batch:
+            yield batch, params
+
+    def _upsert(self, cur: Any, t: Table, rows: list[dict[str, Any]], overwrite: bool = False) -> int:
+        rows = self._dedupe(t, rows)
+        q = tsql_ident
+        cols = t.column_names
+        non_pk = [c for c in cols if c not in t.pk]
+        on = " AND ".join(f"tgt.{q(c)} = src.{q(c)}" for c in t.pk)
+        if overwrite:
+            sets = ", ".join(f"tgt.{q(c)} = src.{q(c)}" for c in non_pk)
+        else:
+            sets = ", ".join(f"tgt.{q(c)} = COALESCE(src.{q(c)}, tgt.{q(c)})" for c in non_pk)
+        col_list = ", ".join(q(c) for c in cols)
+        src_list = ", ".join(f"src.{q(c)}" for c in cols)
+        for values, params in self._values(t, rows):
+            cur.execute(
+                f"MERGE {self.t(t.name)} WITH (HOLDLOCK) AS tgt "
+                f"USING (VALUES {', '.join(values)}) AS src ({col_list}) ON {on} "
+                + (f"WHEN MATCHED THEN UPDATE SET {sets} " if non_pk else "")
+                + f"WHEN NOT MATCHED BY TARGET THEN INSERT ({col_list}) VALUES ({src_list});",
+                tuple(params),
+            )
+        return len(rows)
+
+    def _insert_rows(self, cur: Any, t: Table, rows: list[dict[str, Any]]) -> int:
+        rows = self._dedupe(t, rows)
+        col_list = ", ".join(tsql_ident(c) for c in t.column_names)
+        for values, params in self._values(t, rows):
+            cur.execute(f"INSERT INTO {self.t(t.name)} ({col_list}) VALUES {', '.join(values)};", tuple(params))
+        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:
+        self.rollback()  # in case the run failed mid-transaction
+        with self.conn.cursor() as cur:
+            self._insert_rows(cur, TABLE_BY_NAME["sqlmirror_sync_log"], [{
+                "id": int(started.timestamp() * 1_000_000), "started_at": started, "finished_at": utcnow(),
+                "kind": kind, "source": source, "files": files,
+                "issues": stats.snapshots if stats else 0, "tombstones": stats.tombstones if stats else 0,
+                "skipped": stats.skipped_stale if stats else 0, "error": error,
+            }])
+        self.conn.commit()
+
+    def recent_log(self, limit: int = 20) -> list[tuple]:
+        with self.conn.cursor() as cur:
+            cur.execute(f"SELECT TOP ({int(limit)}) [started_at], [finished_at], [kind], [source], [files], [issues], "
+                        f"[tombstones], [skipped], [error] FROM {self.t('sqlmirror_sync_log')} ORDER BY [started_at] DESC")
+            return [(_utc(r[0]), _utc(r[1]), *r[2:]) for r in cur.fetchall()]
+
+    # -------------------------------------------------------------- queries
+    def project_counts(self) -> dict[str, int]:
+        with self.conn.cursor() as cur:
+            cur.execute(f"SELECT p.[pkey], COUNT(i.[id]) FROM {self.t('project')} p "
+                        f"LEFT JOIN {self.t('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(f"SELECT i.[id], i.[updated] FROM {self.t('jiraissue')} i "
+                        f"JOIN {self.t('project')} p ON p.[id] = i.[project] WHERE p.[pkey] = %s", (project_key,))
+            return {int(r[0]): self.local_to_aware(r[1]) for r in cur.fetchall()}
+
+    def local_to_aware(self, v: Any) -> Any:
+        """A datetime read from a Data Center table (naive, in the mirror's zone) as an aware datetime."""
+        if isinstance(v, datetime) and v.tzinfo is None:
+            if self.tz is timezone.utc:
+                return v.replace(tzinfo=timezone.utc)
+            # In the autumn hour that happens twice, take the later reading, so reconcile
+            # never sees the mirror as older than Jira and re-fetches forever.
+            # Returned in UTC: Python never treats an ambiguous local time as equal to another zone's.
+            return max(v.replace(tzinfo=self.tz, fold=0).astimezone(timezone.utc),
+                       v.replace(tzinfo=self.tz, fold=1).astimezone(timezone.utc))
+        return v
+
+    def table_counts(self) -> dict[str, int]:
+        out = {}
+        with self.conn.cursor() as cur:
+            for name in TABLE_BY_NAME:
+                cur.execute(f"SELECT COUNT_BIG(*) FROM {self.t(name)}")
+                out[name] = int(cur.fetchone()[0])
+        return out
+
+    def freshness(self) -> tuple[datetime | None, datetime | None]:
+        with self.conn.cursor() as cur:
+            cur.execute(f"SELECT MAX([applied_at]), MAX([fetched_at]) FROM {self.t('sqlmirror_issue_state')}")
+            row = cur.fetchone()
+        return _utc(row[0]), _utc(row[1])
+
+    def read_table(self, t: Table) -> list[tuple]:
+        with self.conn.cursor() as cur:
+            cur.execute(f"SELECT {', '.join(tsql_ident(c) for c in t.column_names)} FROM {self.t(t.name)}")
+            return cur.fetchall()
+
+    def explain(self, sql_text: str) -> str:
+        """"ok" if SQL Server can compile the statement against the mirror, else the error."""
+        try:
+            with self.conn.cursor() as cur:
+                cur.execute("EXEC sp_describe_first_result_set @tsql = %s", (sql_text,))
+                cur.fetchall()
+            return "ok"
+        except Exception as exc:  # noqa: BLE001 - pytds.Error and friends
+            return str(exc).strip().splitlines()[0]
+        finally:
+            self.rollback()
+
+    def query(self, sql_text: str, params: tuple = (), timeout_s: int = 60) -> list[tuple]:
+        try:
+            with self.conn.cursor() as cur:
+                if params:
+                    cur.execute(sql_text, params)
+                else:
+                    cur.execute(sql_text)
+                return cur.fetchall() if cur.description else []
+        finally:
+            self.rollback()
+
+
+def _id_chunks(ids: Iterable[int]) -> Iterator[str]:
+    """Comma-separated id lists for IN (...). Ids are ints, so inlining them is safe."""
+    ids = [int(i) for i in ids]
+    for i in range(0, len(ids), ID_CHUNK):
+        yield ", ".join(str(x) for x in ids[i:i + ID_CHUNK])
+
+
+# ------------------------------------------------------------ server admin
+def server_connect(dsn: str, *, autocommit: bool = True) -> Any:
+    """A connection to the server itself (database from the DSN, usually master), for test setup."""
+    import pytds
+
+    kwargs = parse_dsn(dsn)
+    kwargs.pop("schema", None)
+    kwargs.pop("timezone", None)
+    kwargs["autocommit"] = autocommit
+    return pytds.connect(**kwargs)
+
+
+def dsn_with_database(dsn: str, database: str) -> str:
+    u = urlsplit(dsn)
+    return u._replace(path="/" + database).geturl()
+
+
+def recreate_database(server_dsn: str, name: str) -> str:
+    q = tsql_ident(name)
+    lit = name.replace("'", "''")
+    conn = server_connect(server_dsn)
+    try:
+        with conn.cursor() as cur:
+            cur.execute(f"IF DB_ID(N'{lit}') IS NOT NULL BEGIN ALTER DATABASE {q} SET SINGLE_USER WITH ROLLBACK IMMEDIATE; "
+                        f"DROP DATABASE {q}; END")
+            cur.execute(f"CREATE DATABASE {q}")
+    finally:
+        conn.close()
+    return dsn_with_database(server_dsn, name)
+
+
+def drop_database(server_dsn: str, name: str) -> None:
+    q = tsql_ident(name)
+    lit = name.replace("'", "''")
+    conn = server_connect(server_dsn)
+    try:
+        with conn.cursor() as cur:
+            cur.execute(f"IF DB_ID(N'{lit}') IS NOT NULL BEGIN ALTER DATABASE {q} SET SINGLE_USER WITH ROLLBACK IMMEDIATE; "
+                        f"DROP DATABASE {q}; END")
+    finally:
+        conn.close()
+
+
+__all__ = ["MssqlMirror", "parse_dsn", "recreate_database", "drop_database", "dsn_with_database", "Decimal"]

sqlmirror/dev/__init__.py

@@ -0,0 +1 @@
+"""Development tools: synthetic tenants, a mock Jira server, site loaders, golden queries."""