"""Driver ABC + canonical row TypedDicts + shared helpers for datasource drivers.

This module is the contract every concrete datasource driver in 3tears
implements. Read it top-to-bottom before authoring a new driver --
``IMPLEMENTING_DRIVERS.md`` in this package walks the same surface from
the implementer's perspective.

The ABC is intentionally minimal:

- **No pool / no connection on the public surface.** BigQuery is stateless
  HTTPS; Snowflake has its own connection management; asyncpg uses a Pool
  internally but never exposes it. Keeping the surface backend-agnostic
  is how the contract stays honest across four very different backends.
- **Every method is ``async def``** even when the backend library is sync.
  Sync-bridged drivers (Redshift / Snowflake / BigQuery) route through
  :class:`threetears.datasources.drivers._sync_bridge.AsyncSyncBridge`.
- **``$1``-style placeholders** are the contract; concrete drivers
  translate to their dialect via
  :func:`threetears.datasources.drivers._util._translate_placeholders`.
- **Cancellation propagation** is mandatory and routes through
  :meth:`Driver._with_cancellation` so all four drivers share one
  implementation rather than three drifting copies.
- **Row shapes are pinned** via :class:`TableRow` and :class:`ColumnRow`.
  The ``is_nullable`` field is the raw warehouse string
  (``'YES'``/``'NO'``/``''``), NOT a boolean -- the Tier-2 column hash
  from datasource-task-02 depends on byte-equality with the warehouse-
  side MD5, which uses the raw value.

Observability contract (DS-09-11):

- ``datasource.driver.query.duration{driver_type, datasource_name}`` -- histogram
- ``datasource.driver.cancellation.fired{driver_type}`` -- counter
- ``datasource.driver.error{driver_type, error_kind}`` -- counter
- ``datasource.driver.executor.saturation{datasource_name}`` -- gauge
  (sync-bridged drivers only)
- ``datasource.driver.cache.{hit,miss}{datasource_name}`` -- counters
  (sync-bridged drivers only)

The :func:`_observed` decorator wraps method bodies to emit the always-on
duration + error metrics. Cache / saturation / cancellation-fired
metrics are driver-specific and emitted manually by the concrete driver.

``close()`` concurrency semantics (DS-09-12):

Single-shot. Concurrent calls and concurrent in-flight
``fetch``/``execute``/``fetch_iter`` while ``close()`` is running are
undefined behaviour. Drivers SHOULD set a ``_closed: bool`` flag and
reject subsequent operations with :class:`RuntimeError`. Drivers SHOULD
NOT call ``executor.shutdown(wait=True)`` from within an asyncio
coroutine (deadlocks the loop); use ``executor.shutdown(wait=False)``
and let the threads drain naturally.
"""

from __future__ import annotations

import asyncio
import functools
import inspect
import time
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator, Awaitable, Callable
from types import TracebackType
from typing import Any, TypeAlias, TypedDict, TypeVar

from threetears.observe import get_logger

__all__ = [
    "CallbackTransaction",
    "ColumnCoverage",
    "ColumnRow",
    "Driver",
    "TableRow",
    "Transaction",
    "TransactionContext",
]

log = get_logger(__name__)

F = TypeVar("F", bound=Callable[..., Awaitable[Any]])


# ---------------------------------------------------------------------------
# Pinned row shapes (DS-09-10)
# ---------------------------------------------------------------------------


class TableRow(TypedDict):
    """canonical row shape returned by :meth:`Driver.list_tables`.

    pinned via TypedDict so every backend driver agrees on the keys
    exposed to schema-introspection callers. matches the columns
    selected from ``information_schema.tables`` by the standard
    driver implementations.

    :key table_schema: schema name (matches the warehouse's
        ``information_schema.tables.table_schema``)
    :key table_name: table name within the schema
    """

    table_schema: str
    table_name: str


class ColumnRow(TypedDict):
    """canonical row shape returned by :meth:`Driver.list_columns`.

    ``is_nullable`` is the RAW warehouse value (``'YES'`` / ``'NO'`` /
    ``''``), NOT a boolean. the Tier-2 column hash in
    datasource-task-02 computes MD5 over a concatenation of the column
    metadata WITH the raw nullable string; converting to bool here
    would make the python-side hash diverge from the warehouse-side
    MD5, breaking the change-probe contract.

    :key table_schema: schema name (matches warehouse
        ``information_schema.columns.table_schema``)
    :key table_name: table name
    :key column_name: column name
    :key data_type: warehouse-reported data type string (varies by
        backend; e.g. asyncpg returns ``'integer'``, Redshift returns
        ``'INT4'``). drivers MUST surface the raw value
    :key is_nullable: raw warehouse nullable indicator -- ``'YES'``,
        ``'NO'``, or ``''``. NEVER a bool
    :key ordinal_position: 1-indexed column position from the
        information schema
    """

    table_schema: str
    table_name: str
    column_name: str
    data_type: str
    is_nullable: str
    ordinal_position: int


class ColumnCoverage(TypedDict):
    """per-column value-coverage counts from :meth:`Driver.column_value_coverage`.

    the raw facts a warehouse scan can decide about one numeric column: how many
    rows the table holds, and in how many of them the column is a non-null,
    non-zero value. ``total_rows > 0 and nonzero_count == 0`` is the unloaded /
    dead-column signal -- a ``0`` that means "not loaded for this table", not a
    measured zero. the driver returns counts only; the loaded-vs-unloaded VERDICT
    is the caller's, never the driver's.

    :key total_rows: row count of the scanned table
    :key nonzero_count: number of rows where the column is NOT NULL and ``<> 0``
    """

    total_rows: int
    nonzero_count: int


def _quote_ident(name: str) -> str:
    """double-quote a SQL identifier, escaping embedded quotes by doubling.

    identifiers (schema / table / column names) cannot be passed as bind
    parameters, so the coverage probe interpolates them -- quoting defends the
    interpolation even though the names come from the warehouse catalog.

    :param name: raw identifier
    :ptype name: str
    :return: a safely double-quoted identifier
    :rtype: str
    """
    return '"' + name.replace('"', '""') + '"'


def _build_coverage_sql(schema: str, table: str, columns: list[str]) -> str:
    """build the single-scan coverage aggregate for ``columns`` in ``schema.table``.

    one pass over the table: ``COUNT(*)`` plus, per column, a
    ``COUNT(NULLIF(col, 0))`` (which counts only non-null, non-zero rows --
    ``NULLIF`` maps both ``0`` and ``NULL`` to ``NULL`` and ``COUNT`` skips
    ``NULL``). columns alias to ``nz_<index>`` so the caller maps results back by
    position without re-quoting. portable standard SQL (Redshift / Postgres /
    Yugabyte / Snowflake); BigQuery's backtick identifiers would need an override.

    :param schema: schema name
    :ptype schema: str
    :param table: table name
    :ptype table: str
    :param columns: numeric columns to probe (non-empty)
    :ptype columns: list[str]
    :return: the coverage SELECT statement
    :rtype: str
    """
    qualified = f"{_quote_ident(schema)}.{_quote_ident(table)}"
    selects = ["COUNT(*) AS total_rows"]
    selects.extend(f"COUNT(NULLIF({_quote_ident(col)}, 0)) AS nz_{i}" for i, col in enumerate(columns))
    return f"SELECT {', '.join(selects)} FROM {qualified}"


def _build_coverage_by_dimension_sql(
    schema: str,
    table: str,
    dimension_column: str,
    columns: list[str],
) -> str:
    """build the single grouped-scan coverage aggregate for ``columns`` per dimension.

    one pass over the table grouped by ``dimension_column``: the dimension value,
    ``COUNT(*)``, and per column a ``COUNT(NULLIF(col, 0))`` (non-null, non-zero
    rows). the dimension aliases to ``dim_value`` and columns to ``nz_<index>`` so
    the caller maps results back by position without re-quoting. portable standard
    SQL (Redshift / Postgres / Yugabyte / Snowflake); BigQuery's backtick
    identifiers would need an override, same as :func:`_build_coverage_sql`.

    :param schema: schema name
    :ptype schema: str
    :param table: table name
    :ptype table: str
    :param dimension_column: column to group coverage by (non-empty)
    :ptype dimension_column: str
    :param columns: numeric columns to probe (non-empty)
    :ptype columns: list[str]
    :return: the grouped coverage SELECT statement
    :rtype: str
    """
    qualified = f"{_quote_ident(schema)}.{_quote_ident(table)}"
    dimension = _quote_ident(dimension_column)
    selects = [f"{dimension} AS dim_value", "COUNT(*) AS total_rows"]
    selects.extend(f"COUNT(NULLIF({_quote_ident(col)}, 0)) AS nz_{i}" for i, col in enumerate(columns))
    return f"SELECT {', '.join(selects)} FROM {qualified} GROUP BY {dimension}"


# ---------------------------------------------------------------------------
# @_observed decorator (DS-09-11)
# ---------------------------------------------------------------------------


# OTel metrics availability is cached after first probe -- zero-cost path
# when ``opentelemetry`` is not installed, matching the ``@traced``
# decorator's discipline in :mod:`threetears.observe.tracing`. drivers
# don't pay an import cost or a per-call attribute lookup beyond a single
# bool check.
_otel_metrics_available: bool | None = None


def _check_otel_metrics() -> bool:
    """probe for ``opentelemetry.metrics``; cache the result.

    :return: True iff ``opentelemetry.metrics`` is importable
    :rtype: bool
    """
    global _otel_metrics_available  # noqa: PLW0603
    if _otel_metrics_available is None:
        try:
            import opentelemetry.metrics  # noqa: F401

            _otel_metrics_available = True
        except ImportError:
            _otel_metrics_available = False
    return _otel_metrics_available


# instrument cache so we don't recreate Histogram / Counter objects on
# every call. keyed by ``(driver_type, metric_name)``. populated lazily
# the first time a metric fires for a given driver type.
_instrument_cache: dict[tuple[str, str], Any] = {}


def _get_query_duration_histogram(driver_type: str) -> Any:
    """fetch or create the per-driver-type query.duration histogram.

    :param driver_type: driver type label (e.g. ``"asyncpg"``, ``"redshift"``)
    :ptype driver_type: str
    :return: OTel Histogram instrument (or None if OTel not available)
    :rtype: Any
    """
    result: Any = None
    if _check_otel_metrics():
        key = (driver_type, "datasource.driver.query.duration")
        instrument = _instrument_cache.get(key)
        if instrument is None:
            from opentelemetry import metrics

            meter = metrics.get_meter("threetears.datasources.drivers")
            instrument = meter.create_histogram(
                name="datasource.driver.query.duration",
                description="datasource driver query duration in seconds",
                unit="s",
            )
            _instrument_cache[key] = instrument
        result = instrument
    return result


def _get_error_counter(driver_type: str) -> Any:
    """fetch or create the per-driver-type error counter.

    :param driver_type: driver type label
    :ptype driver_type: str
    :return: OTel Counter instrument (or None if OTel not available)
    :rtype: Any
    """
    result: Any = None
    if _check_otel_metrics():
        key = (driver_type, "datasource.driver.error")
        instrument = _instrument_cache.get(key)
        if instrument is None:
            from opentelemetry import metrics

            meter = metrics.get_meter("threetears.datasources.drivers")
            instrument = meter.create_counter(
                name="datasource.driver.error",
                description="datasource driver error count by error kind",
            )
            _instrument_cache[key] = instrument
        result = instrument
    return result


def _observed(driver_type: str) -> Callable[[F], F]:
    """decorator factory: wrap an async driver method with standard metric emission.

    emits :data:`datasource.driver.query.duration` (histogram) on every
    completion (success and exception) and
    :data:`datasource.driver.error` (counter) tagged with the exception
    class name on any exception. cancellation is NOT counted as an
    error -- :class:`asyncio.CancelledError` re-raises without bumping
    the error counter (the driver's manual ``cancellation.fired``
    counter covers that). this keeps the always-on metrics zero-effort
    for concrete drivers; backend-specific gauges / counters are
    emitted manually inside the driver.

    when ``opentelemetry.metrics`` isn't installed the decorator is a
    pure passthrough (single bool check per call), matching
    :func:`threetears.observe.traced`.

    :param driver_type: stable label for the driver type (e.g.
        ``"asyncpg"`` / ``"redshift"`` / ``"snowflake"`` / ``"bigquery"``).
        used as the OTel attribute ``driver_type``
    :ptype driver_type: str
    :return: decorator suitable for ``async def`` driver methods
    :rtype: Callable[[F], F]
    """

    def decorator(fn: F) -> F:
        if not inspect.iscoroutinefunction(fn):
            raise TypeError(f"@_observed only wraps async functions; {fn.__qualname__} is not async")

        @functools.wraps(fn)
        async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
            if not _check_otel_metrics():
                return await fn(self, *args, **kwargs)
            histogram = _get_query_duration_histogram(driver_type)
            error_counter = _get_error_counter(driver_type)
            # datasource_name attribute is best-effort: drivers expose
            # ``_datasource_name`` if they want to participate, else
            # the metric carries only the driver_type label.
            datasource_name = getattr(self, "_datasource_name", "unknown")
            attrs = {"driver_type": driver_type, "datasource_name": datasource_name}
            start = time.monotonic()
            try:
                result = await fn(self, *args, **kwargs)
            except asyncio.CancelledError:
                # propagation only -- the per-driver
                # ``cancellation.fired`` counter is bumped explicitly
                # by concrete drivers from _with_cancellation hooks
                raise
            except Exception as exc:
                error_counter.add(
                    1,
                    attributes={
                        "driver_type": driver_type,
                        "error_kind": type(exc).__name__,
                    },
                )
                raise
            finally:
                histogram.record(time.monotonic() - start, attributes=attrs)
            return result

        return async_wrapper  # type: ignore[return-value]

    return decorator


# ---------------------------------------------------------------------------
# Transaction surface (DSD-01-01 / DSD-01-02)
# ---------------------------------------------------------------------------


#: signature of the driver-supplied fetch hook a :class:`CallbackTransaction`
#: forwards to. params arrive as a tuple rather than varargs so the hook can
#: be a plain bound method with a fixed arity.
TransactionFetchHook: TypeAlias = Callable[[str, tuple[Any, ...], "int | None"], Awaitable[list[dict[str, Any]]]]

#: signature of the driver-supplied execute hook.
TransactionExecuteHook: TypeAlias = Callable[[str, tuple[Any, ...], "int | None"], Awaitable[None]]

#: signature of the driver-supplied finish hook. the bool is True for
#: commit, False for rollback. the hook owns ending the backend
#: transaction AND returning the pinned connection to its pool / cache,
#: because those two must happen exactly once and in that order.
TransactionFinishHook: TypeAlias = Callable[[bool], Awaitable[None]]


class Transaction(ABC):
    """a multi-statement unit of work pinned to ONE backend session.

    the shape mirrors ``asyncpg``'s own transaction API so a reader who
    knows one knows the other. every statement issued through this
    handle runs on the same backend session for the transaction's
    lifetime -- which is the whole point: promotion's ``CREATE TABLE
    AS`` / ``DROP`` / ``ALTER`` / ``GRANT`` sequence is only atomic if
    the statements share a session.

    obtain one via :meth:`Driver.transaction` (the context manager,
    which commits on clean exit and rolls back on exception) or
    :meth:`Driver.begin` (explicit, when the commit point is decided
    somewhere the ``async with`` block cannot reach). a handle from
    :meth:`Driver.begin` MUST be committed or rolled back -- an
    abandoned one holds its pooled connection until the driver closes.

    holding a transaction open reduces the driver's effective pool
    capacity for its duration. that coupling is why the executor's
    admission control caps build concurrency at enqueue rather than
    letting arbitrarily many transactions contend for the cache; the
    two knobs must not be tuned independently.
    """

    @abstractmethod
    async def fetch(self, sql: str, *params: Any, timeout_seconds: int | None = None) -> list[dict[str, Any]]:
        """run a SELECT on this transaction's pinned session.

        :param sql: SQL text with ``$1``-style placeholders
        :ptype sql: str
        :param params: positional placeholder values
        :ptype params: Any
        :param timeout_seconds: per-statement timeout override in
            seconds. applied transaction-locally, so it bounds THIS
            statement and unwinds with the transaction. None leaves the
            connection's configured ceiling in force
        :ptype timeout_seconds: int | None
        :return: list of column-name -> value dicts in row order
        :rtype: list[dict[str, Any]]
        :raises RuntimeError: if the transaction already finished
        """

    @abstractmethod
    async def execute(self, sql: str, *params: Any, timeout_seconds: int | None = None) -> None:
        """run a DML / DDL statement on this transaction's pinned session.

        :param sql: SQL text with ``$1``-style placeholders
        :ptype sql: str
        :param params: positional placeholder values
        :ptype params: Any
        :param timeout_seconds: per-statement timeout override in
            seconds; transaction-local. None leaves the connection's
            configured ceiling in force
        :ptype timeout_seconds: int | None
        :return: nothing
        :rtype: None
        :raises RuntimeError: if the transaction already finished
        """

    @abstractmethod
    async def commit(self) -> None:
        """commit the transaction and release its pinned session.

        :return: nothing
        :rtype: None
        :raises RuntimeError: if the transaction already finished
        """

    @abstractmethod
    async def rollback(self) -> None:
        """roll the transaction back and release its pinned session.

        :return: nothing
        :rtype: None
        :raises RuntimeError: if the transaction already finished
        """


class CallbackTransaction(Transaction):
    """concrete :class:`Transaction` driven by three driver-supplied hooks.

    both concrete drivers construct one of these rather than each
    shipping its own :class:`Transaction` subclass. the finished-guard,
    the "released exactly once" invariant, and the argument shape then
    live in ONE place, which is what makes the two drivers provably the
    same surface instead of two implementations that agree today.

    the hooks are bound methods of the driver, so this class never
    reaches into driver internals.

    :param on_fetch: driver hook running a SELECT on the pinned session
    :ptype on_fetch: TransactionFetchHook
    :param on_execute: driver hook running a DML / DDL statement on the
        pinned session
    :ptype on_execute: TransactionExecuteHook
    :param on_finish: driver hook ending the backend transaction and
        releasing the pinned connection. receives True for commit and
        False for rollback
    :ptype on_finish: TransactionFinishHook
    """

    def __init__(
        self,
        *,
        on_fetch: TransactionFetchHook,
        on_execute: TransactionExecuteHook,
        on_finish: TransactionFinishHook,
    ) -> None:
        """capture the three driver hooks; mark the transaction live.

        :param on_fetch: driver hook running a SELECT
        :ptype on_fetch: TransactionFetchHook
        :param on_execute: driver hook running a DML / DDL statement
        :ptype on_execute: TransactionExecuteHook
        :param on_finish: driver hook ending + releasing the session
        :ptype on_finish: TransactionFinishHook
        :return: nothing
        :rtype: None
        """
        self._on_fetch = on_fetch
        self._on_execute = on_execute
        self._on_finish = on_finish
        self._finished = False

    @property
    def finished(self) -> bool:
        """whether the transaction has already committed or rolled back.

        :return: True once a finish hook has been dispatched
        :rtype: bool
        """
        return self._finished

    async def fetch(self, sql: str, *params: Any, timeout_seconds: int | None = None) -> list[dict[str, Any]]:
        """run a SELECT on the pinned session.

        :param sql: SQL text with ``$1``-style placeholders
        :ptype sql: str
        :param params: positional placeholder values
        :ptype params: Any
        :param timeout_seconds: transaction-local timeout override
        :ptype timeout_seconds: int | None
        :return: list of column-name -> value dicts in row order
        :rtype: list[dict[str, Any]]
        :raises RuntimeError: if the transaction already finished
        """
        self._reject_when_finished()
        return await self._on_fetch(sql, params, timeout_seconds)

    async def execute(self, sql: str, *params: Any, timeout_seconds: int | None = None) -> None:
        """run a DML / DDL statement on the pinned session.

        :param sql: SQL text with ``$1``-style placeholders
        :ptype sql: str
        :param params: positional placeholder values
        :ptype params: Any
        :param timeout_seconds: transaction-local timeout override
        :ptype timeout_seconds: int | None
        :return: nothing
        :rtype: None
        :raises RuntimeError: if the transaction already finished
        """
        self._reject_when_finished()
        await self._on_execute(sql, params, timeout_seconds)

    async def commit(self) -> None:
        """commit and release the pinned session.

        the finished flag is set BEFORE the hook runs: a hook that
        raises still consumed the connection release, and letting a
        caller retry the commit would double-release it.

        :return: nothing
        :rtype: None
        :raises RuntimeError: if the transaction already finished
        """
        self._reject_when_finished()
        self._finished = True
        await self._on_finish(True)

    async def rollback(self) -> None:
        """roll back and release the pinned session.

        :return: nothing
        :rtype: None
        :raises RuntimeError: if the transaction already finished
        """
        self._reject_when_finished()
        self._finished = True
        await self._on_finish(False)

    def _reject_when_finished(self) -> None:
        """raise when the transaction has already been committed or rolled back.

        :return: nothing
        :rtype: None
        :raises RuntimeError: if the transaction already finished
        """
        if self._finished:
            raise RuntimeError("transaction already committed or rolled back")


class TransactionContext:
    """async context manager returned by :meth:`Driver.transaction`.

    commits on a clean exit, rolls back on ANY exception including
    :class:`asyncio.CancelledError`. concrete on this layer rather than
    per driver, so the two drivers cannot disagree about what a bare
    ``async with driver.transaction()`` means.

    a rollback that itself fails is logged, never substituted for the
    caller's exception -- the caller acts on the failure they caused,
    and a failed rollback is observability.

    :param driver: driver whose :meth:`Driver.begin` opens the session
    :ptype driver: Driver
    """

    def __init__(self, driver: Driver) -> None:
        """capture the driver; open nothing yet.

        :param driver: driver whose :meth:`Driver.begin` opens the session
        :ptype driver: Driver
        :return: nothing
        :rtype: None
        """
        self._driver = driver
        self._transaction: Transaction | None = None

    async def __aenter__(self) -> Transaction:
        """open the transaction and hand back its handle.

        :return: the live transaction handle
        :rtype: Transaction
        """
        transaction = await self._driver.begin()
        self._transaction = transaction
        return transaction

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        traceback: TracebackType | None,
    ) -> bool:
        """commit on a clean exit; roll back on any exception.

        :param exc_type: exception class raised in the body, or None
        :ptype exc_type: type[BaseException] | None
        :param exc: exception instance raised in the body, or None
        :ptype exc: BaseException | None
        :param traceback: traceback of the raised exception, or None
        :ptype traceback: TracebackType | None
        :return: False -- the caller's exception is never suppressed
        :rtype: bool
        """
        transaction = self._transaction
        if transaction is not None:
            if exc_type is None:
                await transaction.commit()
            else:
                try:
                    await transaction.rollback()
                except Exception as rollback_exc:  # noqa: BLE001 -- the body's exception is the caller's answer
                    # the transaction may still be open server-side; the
                    # connection is evicted by the driver's finish hook
                    # either way, so this is observability, not control flow.
                    log.warning(
                        "transaction rollback failed; original error propagates",
                        extra={
                            "extra_data": {
                                "error": str(rollback_exc),
                                "error_type": type(rollback_exc).__name__,
                            }
                        },
                    )
        return False


# ---------------------------------------------------------------------------
# Driver ABC (DS-09-01..06)
# ---------------------------------------------------------------------------


class Driver(ABC):
    """abstract base for all datasource drivers.

    every concrete backend driver in
    :mod:`threetears.datasources.drivers` subclasses this ABC and
    implements the documented method surface. instances are created
    via :func:`threetears.datasources.drivers.create_driver` rather
    than direct instantiation -- the factory enforces lazy import of
    the backend library and provides the AGENT_INTERNAL pool-borrow
    plumbing.

    placeholder convention (DS-09-04):
        callers always pass ``$1``-style positional placeholders. each
        concrete driver translates internally via
        :func:`threetears.datasources.drivers._util._translate_placeholders`
        to its backend's expected dialect (``%s`` for pyformat /
        Redshift, ``:1`` for numeric, ``@p1`` for BigQuery named-at).

    cancellation contract (DS-09-05, amended by dsd-task-02):
        ``fetch`` / ``execute`` / ``fetch_iter`` MUST cooperate with
        :mod:`asyncio` cancellation: the awaiting coroutine sees
        :class:`asyncio.CancelledError`. concrete drivers route every
        backend call through :meth:`_with_cancellation` so the
        propagation logic lives in one place.

        what a driver MUST NOT do is claim a backend cancellation it
        cannot perform. where the backend library exposes a real
        primitive the driver uses it (Redshift terminates the
        server-side backend); where it does not, the driver says so
        rather than calling a method that does not exist. **no driver
        guarantees an in-flight warehouse statement stops.** callers
        that need bounded cancellation make it cooperative and bound it
        by statement completion.

    transaction contract (DSD-01-01 / DSD-01-02):
        :meth:`begin` opens a multi-statement unit of work pinned to
        ONE backend session; :meth:`transaction` wraps it in the shared
        async context manager that commits on clean exit and rolls back
        on exception. a promote is only atomic if its statements share
        a session, which is why the ABC carries this rather than each
        driver improvising.

    per-statement timeout contract (DSD-01-03 / DSD-01-04):
        :meth:`fetch` and :meth:`execute` take ``timeout_seconds``. the
        datasource's configured timeout is a connection-level CEILING;
        the override bounds one statement below it. drivers MUST apply
        the override transaction-locally (or reset it on release, in a
        ``finally``) -- a session-scoped ``SET`` on a pooled connection
        leaks the previous borrower's bound to the next one, and the
        resulting failure is nondeterministic.

    read-path contract (DSD-01-05):
        a completed statement MUST NOT return its connection to the
        pool holding an open transaction. an idle-in-transaction
        connection keeps its snapshot and locks, which is what makes a
        ``DROP`` or an ownership change block behind a *finished*
        query. this is a prerequisite for reaping and stats sharing one
        warehouse, not a tidy-up.

    row-shape contract (DS-09-10):
        :meth:`list_tables` returns :class:`TableRow` dicts and
        :meth:`list_columns` returns :class:`ColumnRow` dicts. the
        ``is_nullable`` field is the RAW warehouse string -- the
        Tier-2 hash depends on byte-equality with the warehouse-side
        MD5 which uses the raw value.

    observability contract (DS-09-11):
        concrete drivers SHOULD decorate query-emitting methods with
        :func:`_observed` to get the standard duration + error metrics
        for free. additional cache / saturation / cancellation-fired
        metrics are emitted manually.

    close concurrency semantics (DS-09-12):
        ``close()`` is single-shot. concurrent calls and concurrent
        in-flight queries while ``close()`` is running are undefined.
        drivers SHOULD set a ``_closed: bool`` flag and reject
        subsequent operations with :class:`RuntimeError`. drivers
        backed by a :class:`concurrent.futures.ThreadPoolExecutor`
        MUST call ``executor.shutdown(wait=False)``;
        ``wait=True`` deadlocks the asyncio event loop.
    """

    @abstractmethod
    async def fetch(self, sql: str, *params: Any, timeout_seconds: int | None = None) -> list[dict[str, Any]]:
        """run an arbitrary SELECT statement; materialize all rows in memory.

        for result sets large enough to risk OOM, prefer
        :meth:`fetch_iter`. ``fetch`` is right for small / bounded
        queries (introspection, single-row lookups, small joins).

        :param sql: SQL text with ``$1``-style placeholders
        :ptype sql: str
        :param params: positional placeholder values
        :ptype params: Any
        :param timeout_seconds: per-statement timeout override in
            seconds. the datasource's configured timeout is a
            connection-level CEILING, not the per-statement value: one
            number cannot govern an hour-long build statement, a
            reaper ``DROP``, and a bounded interactive aggregate alike.
            the override is applied transaction-locally so it cannot
            leak to the next borrower of a pooled connection. None
            leaves the ceiling in force
        :ptype timeout_seconds: int | None
        :return: list of column-name -> value dicts in row order
        :rtype: list[dict[str, Any]]
        :raises asyncio.CancelledError: propagated after best-effort
            backend cancellation
        :raises RuntimeError: if the driver was previously closed
        :raises ValueError: if ``timeout_seconds`` is not a positive int
        """

    @abstractmethod
    async def execute(self, sql: str, *params: Any, timeout_seconds: int | None = None) -> None:
        """run an arbitrary DML / DDL statement; discard any returned rows.

        :param sql: SQL text with ``$1``-style placeholders
        :ptype sql: str
        :param params: positional placeholder values
        :ptype params: Any
        :param timeout_seconds: per-statement timeout override in
            seconds; see :meth:`fetch`. None leaves the connection's
            configured ceiling in force
        :ptype timeout_seconds: int | None
        :return: nothing
        :rtype: None
        :raises asyncio.CancelledError: propagated after best-effort
            backend cancellation
        :raises RuntimeError: if the driver was previously closed
        :raises ValueError: if ``timeout_seconds`` is not a positive int
        """

    @abstractmethod
    async def begin(self) -> Transaction:
        """open a transaction pinned to ONE backend session.

        the caller MUST finish the returned handle with
        :meth:`Transaction.commit` or :meth:`Transaction.rollback` --
        an abandoned handle holds its pooled connection until the
        driver closes. prefer :meth:`transaction`, which guarantees the
        finish, and reach for ``begin`` only when the commit point is
        decided somewhere an ``async with`` block cannot reach.

        the pinned connection is out of the pool / cache for the
        transaction's whole life, so a transaction reduces effective
        pool capacity for its duration.

        :return: a live transaction handle
        :rtype: Transaction
        :raises RuntimeError: if the driver was previously closed
        """

    def transaction(self) -> TransactionContext:
        """return an async context manager wrapping :meth:`begin`.

        commits on a clean exit, rolls back on any exception::

            async with driver.transaction() as tx:
                await tx.execute("CREATE TABLE ... AS SELECT ...")
                await tx.execute("GRANT SELECT ON ... TO ...")

        concrete on the ABC so every driver presents exactly the same
        commit / rollback semantics. do NOT hold a transaction open
        across a tool-call boundary: transactions belong to the
        executor and the promoter, both of which run off the tool path.

        :return: the shared transaction context manager
        :rtype: TransactionContext
        """
        return TransactionContext(self)

    async def fetch_iter(self, sql: str, *params: Any) -> AsyncIterator[dict[str, Any]]:
        """stream rows for large result sets.

        default implementation calls :meth:`fetch` and yields each row
        -- correct but materializes the full result first, defeating
        the streaming purpose. drivers backed by native server-side
        cursors (``asyncpg``, ``redshift_connector``) override to
        stream incrementally. the default exists so simple stub
        drivers (Snowflake, BigQuery) have a working ``fetch_iter``
        from day one.

        :param sql: SQL text with ``$1``-style placeholders
        :ptype sql: str
        :param params: positional placeholder values
        :ptype params: Any
        :return: async iterator over column-name -> value dicts
        :rtype: AsyncIterator[dict[str, Any]]
        :raises asyncio.CancelledError: propagated after best-effort
            backend cancellation
        :raises RuntimeError: if the driver was previously closed
        """
        rows = await self.fetch(sql, *params)
        for row in rows:
            yield row

    @abstractmethod
    async def list_tables(self, schemas: list[str]) -> list[TableRow]:
        """list tables visible to the connection within the given schemas.

        :param schemas: schema-name allow-list. empty list means "no
            tables" (callers gate against the agent.yaml schema
            whitelist before calling)
        :ptype schemas: list[str]
        :return: list of :class:`TableRow` dicts
        :rtype: list[TableRow]
        :raises RuntimeError: if the driver was previously closed
        """

    @abstractmethod
    async def list_columns(self, schemas: list[str]) -> list[ColumnRow]:
        """list columns for every table in the given schemas.

        :param schemas: schema-name allow-list
        :ptype schemas: list[str]
        :return: list of :class:`ColumnRow` dicts with the raw
            ``is_nullable`` string preserved
        :rtype: list[ColumnRow]
        :raises RuntimeError: if the driver was previously closed
        """

    @abstractmethod
    async def table_hashes(self, schemas: list[str]) -> dict[tuple[str, str], str]:
        """per-table MD5 over the column shape; Tier-2 change-probe.

        MUST byte-equal the python-side ``_compute_column_hash`` from
        datasource-task-02 over identical input. the warehouse-side
        MD5 hashes the raw ``is_nullable`` string (``'YES'`` / ``'NO'``
        / ``''``) -- using a boolean here makes the python-side hash
        diverge and breaks the probe.

        :param schemas: schema-name allow-list
        :ptype schemas: list[str]
        :return: mapping of ``(schema, table)`` -> column-shape hex digest
        :rtype: dict[tuple[str, str], str]
        :raises RuntimeError: if the driver was previously closed
        """

    @abstractmethod
    async def test_connection(self) -> None:
        """cheapest possible round-trip; verifies credentials + reachability.

        intended for health-check use; safe to call repeatedly. concrete
        drivers implement as ``SELECT 1`` (or equivalent). NOT the same
        as opening the connection -- that happens in ``__init__`` or
        lazily on first ``fetch``.

        :return: nothing; raises on failure
        :rtype: None
        :raises Exception: backend-specific error if the round-trip fails
        """

    @abstractmethod
    async def close(self) -> None:
        """release driver resources (connections, executors, clients).

        single-shot per the close-concurrency contract in the class
        docstring. drivers SHOULD set ``self._closed = True`` and
        reject subsequent operations with :class:`RuntimeError`.

        sync-bridged drivers MUST call ``executor.shutdown(wait=False)``
        -- ``wait=True`` deadlocks the asyncio event loop because the
        worker threads may be awaiting a coroutine that can't run
        while the loop is blocked.

        :return: nothing
        :rtype: None
        """

    # -------------------------------------------------------------------
    # Concrete: value-coverage probe (datasource honesty)
    # -------------------------------------------------------------------

    async def column_value_coverage(
        self,
        schema: str,
        table: str,
        columns: list[str],
    ) -> dict[str, ColumnCoverage]:
        """scan ``schema.table`` once and report value-coverage per numeric column.

        the decidable half of the unloaded-zero problem: a numeric column that
        is non-null and non-zero in ZERO of the table's rows is unloaded (the
        Alaska ``agg_dummy_eip`` incident), and a ``0`` read from it is missing
        data, not a measured zero. this method returns the raw counts
        (:class:`ColumnCoverage`); the caller decides the verdict. CONCRETE on the
        ABC (portable SQL routed through :meth:`fetch`) so every backend inherits
        it; it does NOT carry its own ``@_observed`` because the inner
        :meth:`fetch` it delegates to is already instrumented.

        :param schema: schema name of the table to scan
        :ptype schema: str
        :param table: table name to scan
        :ptype table: str
        :param columns: numeric column names to probe; an empty list short-circuits
            with no query
        :ptype columns: list[str]
        :return: mapping of column name -> :class:`ColumnCoverage` counts; empty
            when ``columns`` is empty
        :rtype: dict[str, ColumnCoverage]
        :raises asyncio.CancelledError: propagated from :meth:`fetch`
        :raises RuntimeError: if the driver was previously closed
        """
        result: dict[str, ColumnCoverage] = {}
        if columns:
            sql = _build_coverage_sql(schema, table, columns)
            rows = await self.fetch(sql)
            row = rows[0] if rows else {}
            total_rows = int(row.get("total_rows") or 0)
            for index, column in enumerate(columns):
                nonzero = row.get(f"nz_{index}")
                result[column] = ColumnCoverage(
                    total_rows=total_rows,
                    nonzero_count=int(nonzero or 0),
                )
        return result

    async def column_value_coverage_by_dimension(
        self,
        schema: str,
        table: str,
        dimension_column: str,
        columns: list[str],
    ) -> dict[str, dict[str, ColumnCoverage]]:
        """scan ``schema.table`` once grouped by ``dimension_column`` and report
        value-coverage per numeric column per dimension value.

        the grouped sibling of :meth:`column_value_coverage`: the same decidable
        question, but a single ``GROUP BY`` pass so the caller can decide a column
        is unloaded for SOME dimension values while loaded for others -- the
        partial-coverage case the whole-table probe cannot see (a column loaded for
        49 states but all-zero for every row of one). returns the raw counts per
        ``(column, dimension_value)``; the loaded-vs-unloaded VERDICT is the
        caller's. CONCRETE on the ABC (portable SQL routed through :meth:`fetch`).

        every requested column maps to a (possibly empty) per-dimension dict so
        callers can index ``result[column]`` safely. a NULL dimension value (rows
        with no value for the dimension) keys on the empty string. an empty
        ``columns`` OR empty ``dimension_column`` short-circuits with no query.

        :param schema: schema name of the table to scan
        :ptype schema: str
        :param table: table name to scan
        :ptype table: str
        :param dimension_column: column to group coverage by (e.g. ``state_code``);
            empty short-circuits with no query
        :ptype dimension_column: str
        :param columns: numeric column names to probe; empty short-circuits
        :ptype columns: list[str]
        :return: mapping of column name -> {dimension_value -> :class:`ColumnCoverage`};
            empty when ``columns`` or ``dimension_column`` is empty
        :rtype: dict[str, dict[str, ColumnCoverage]]
        :raises asyncio.CancelledError: propagated from :meth:`fetch`
        :raises RuntimeError: if the driver was previously closed
        """
        result: dict[str, dict[str, ColumnCoverage]] = {}
        if columns and dimension_column:
            result = {column: {} for column in columns}
            sql = _build_coverage_by_dimension_sql(schema, table, dimension_column, columns)
            rows = await self.fetch(sql)
            for row in rows:
                raw_dimension = row.get("dim_value")
                dimension_value = "" if raw_dimension is None else str(raw_dimension)
                total_rows = int(row.get("total_rows") or 0)
                for index, column in enumerate(columns):
                    nonzero = row.get(f"nz_{index}")
                    result[column][dimension_value] = ColumnCoverage(
                        total_rows=total_rows,
                        nonzero_count=int(nonzero or 0),
                    )
        return result

    # -------------------------------------------------------------------
    # Shared helpers (DS-09-05)
    # -------------------------------------------------------------------

    async def _with_cancellation(
        self,
        coro_fn: Callable[[], Awaitable[Any]],
        *,
        cancel_callback: Callable[[], Any],
    ) -> Any:
        """run ``coro_fn`` and propagate cancellation to the backend.

        the canonical pattern every concrete driver routes through.
        without a shared helper four drivers would each implement the
        try/except/cancel-callback dance and three would get it
        slightly wrong (forget to suppress callback errors, forget to
        await an async callback, swallow the original
        :class:`CancelledError`). the shared helper makes the contract
        load-bearing in exactly one place.

        propagation rules:

        - on success: return the coroutine's result.
        - on :class:`asyncio.CancelledError`: invoke ``cancel_callback``,
          log any exception it raises at warning without letting it
          displace the cancellation, then re-raise. if the callback
          returns a coroutine, it is awaited.

        :param coro_fn: zero-arg callable returning the awaitable to
            run. taking a callable rather than the awaitable itself
            keeps construction inside the try block, which lets us
            short-circuit on cancellation that fires before the
            backend call starts
        :ptype coro_fn: Callable[[], Awaitable[Any]]
        :param cancel_callback: backend-specific hook invoked on
            cancellation. what it may honestly do is bounded by what the
            backend library actually exposes: the Redshift driver
            terminates the server-side backend and closes the socket,
            while the asyncpg driver only RECORDS the cancellation --
            :class:`asyncpg.Connection` has no cancel verb and asyncpg's
            protocol already requests the backend cancel itself (see
            dsd-task-02 in the asyncpg driver's module docstring). a
            callback MUST NOT imply a cancellation the backend cannot
            perform. may be sync or async; sync callables whose return
            is a coroutine are awaited
        :ptype cancel_callback: Callable[[], Any]
        :return: the wrapped coroutine's result on success
        :rtype: Any
        :raises asyncio.CancelledError: re-raised after best-effort
            cancellation. the original exception is preserved
        """
        try:
            result = await coro_fn()
        except asyncio.CancelledError:
            try:
                cancel_result = cancel_callback()
                if asyncio.iscoroutine(cancel_result):
                    await cancel_result
            except Exception as exc:  # noqa: BLE001 -- the CancelledError below is the caller's answer
                # The cancellation still propagates, but the backend statement it was meant to stop
                # may not have stopped: a failed cancel leaves work running server-side while this
                # side reports cancelled.
                log.warning(
                    "cancel callback failed; backend work may still be running",
                    extra={"extra_data": {"error": str(exc), "error_type": type(exc).__name__}},
                )
            raise
        return result
