AgentLab/backend/app/data/engine.py

"""数据源连接层。各家数据库的方言差异全部收敛在这个文件里。

上层（工具、探查、API）只认 DataSource 这个模型，不需要知道 Oracle 要 service_name、
MySQL 要 charset、asyncpg 不认 sslmode 这些事。哪天要加一种库，改这里一处。
"""
from __future__ import annotations

import asyncio
import json
import logging
import re
import time
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any, AsyncIterator, Awaitable, Callable

from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
from sqlalchemy import text

from app.core.crypto import decrypt
from app.data import guard
from app.data.guard import QueryLimits, SqlRejected, check, is_write

logger = logging.getLogger(__name__)

# 各方言的 SQLAlchemy 驱动。都选 asyncio 原生驱动，没有一个需要装数据库客户端：
# oracledb 走 thin 模式（纯 Python 协议实现），省掉了 Oracle Instant Client 这个
# 历来最劝退的一步。
_DRIVERS = {
    "mysql": "mysql+aiomysql",
    "mariadb": "mysql+aiomysql",
    "postgres": "postgresql+asyncpg",
    "postgresql": "postgresql+asyncpg",
    "oracle": "oracle+oracledb",
    "sqlite": "sqlite+aiosqlite",
}

_DEFAULT_PORTS = {
    "mysql": 3306, "mariadb": 3306,
    "postgres": 5432, "postgresql": 5432,
    "oracle": 1521,
}

SUPPORTED_KINDS = tuple(sorted(set(_DRIVERS)))

#: 数据源 options 里按源配置的查询时限（秒）。表单上叫「查询时限」
QUERY_TIMEOUT_OPTION = "query_timeout_s"
#: 按源配置时能填的范围。再长的查询不该让一个 agent 干等：该做成离线任务了
MIN_QUERY_TIMEOUT_S = 1
MAX_QUERY_TIMEOUT_S = 600

#: 数据源 options 里要遮罩的列（列名列表，或者逗号分隔的一段文字）。证据面板展示原始行时，
#: 这些列的值一律换成「已遮罩」。在有身份体系之前，遮罩只减少暴露，不是安全边界：
#: 完整快照仍能按工件 id 取到，SQL 里给列起个别名也能绕开
MASK_COLUMNS_OPTION = "mask_columns"

# options 里这些 key 是 AgentLab 自己的配置，不是驱动参数，拼 URL 时要摘掉。
# schema：探查哪个 schema（企业库里只读账号名下常常什么都没有，数据在别处）
# query_timeout_s：查询时限，由数据层按语句下发给数据库，见 _server_deadline
# mask_columns：证据面板遮罩的列，驱动不认识它，拼进连接串会被当成未知参数拒掉
_NON_DRIVER_OPTIONS = frozenset({"schema", QUERY_TIMEOUT_OPTION, MASK_COLUMNS_OPTION})

_MASK_SPLIT = re.compile(r"[,，、;；\n]+")


@dataclass
class QueryResult:
    columns: list[str]
    rows: list[list[Any]]
    row_count: int
    truncated: bool
    elapsed_ms: int
    sql: str
    #: 列名 → number / text / date / datetime / time / boolean。拿不准的列（全是空值、类型混着）不在里面
    column_types: dict[str, str] | None = None

    def to_payload(self) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "columns": self.columns,
            "rows": self.rows,
            "row_count": self.row_count,
            "truncated": self.truncated,
            "elapsed_ms": self.elapsed_ms,
            "sql": self.sql,
        }
        if self.column_types:
            payload["column_types"] = self.column_types
        return payload


def build_url(source: Any, *, reveal: bool = False) -> str:
    """拼连接串。reveal=False 时把密码换成星号——日志和错误信息里用这个。

    密码永远不该出现在日志、事件流、或者返回给前端的任何地方。默认遮掉，
    真要连库时才显式 reveal=True。
    """
    kind = (source.kind or "").lower()
    driver = _DRIVERS.get(kind)
    if not driver:
        raise ValueError(f"不支持的数据库类型：{source.kind}（支持 {', '.join(SUPPORTED_KINDS)}）")

    options: dict[str, Any] = dict(source.options or {})
    # options 混着两类东西：驱动参数（charset、service_name…）和我们自己的配置
    # （schema 指探查哪个 schema）。后者不能拼进连接串，否则驱动会把它当成
    # 未知的连接参数直接拒掉。
    for key in _NON_DRIVER_OPTIONS:
        options.pop(key, None)

    if kind == "sqlite":
        # SQLite 没有主机和账号，database 就是文件路径
        return f"{driver}:///{source.database or ':memory:'}"

    password = decrypt(source.password) or "" if reveal else "***"
    from urllib.parse import quote_plus

    user = quote_plus(source.username or "")
    pw = quote_plus(password) if reveal else password
    auth = f"{user}:{pw}@" if user else ""
    port = source.port or _DEFAULT_PORTS.get(kind)
    host = f"{source.host or 'localhost'}:{port}" if port else (source.host or "localhost")

    if kind == "oracle":
        # Oracle 用 service_name 或 sid，不是"数据库名"。这是最常见的配错点，
        # 所以两种都认，并且在都没填时给出明确提示而不是让驱动报晦涩的 ORA 错误。
        service = options.pop("service_name", None) or source.database
        sid = options.pop("sid", None)
        if not service and not sid:
            raise ValueError("Oracle 需要在 options 里给 service_name 或 sid（也可以填在库名里）")
        tail = f"?service_name={service}" if service else f"?sid={sid}"
        query = "&".join(f"{k}={v}" for k, v in options.items())
        return f"{driver}://{auth}{host}/{tail}{('&' + query) if query else ''}"

    query = "&".join(f"{k}={v}" for k, v in options.items())
    return f"{driver}://{auth}{host}/{source.database or ''}{('?' + query) if query else ''}"


def engine_args(source: Any) -> tuple[str, dict[str, Any]]:
    """真正连库用的连接串和额外参数。只读源在这里拿到**连接层**的只读。

    守卫按关键字判定，而模型能写出什么 SQL 事前无法穷举——`WITH … DELETE`
    就曾从首关键字白名单下面钻过去，在 SQLite 上真把数据删了（pysqlite 只在
    INSERT/UPDATE/DELETE 开头的语句前隐式开事务，WITH 开头的直接自动提交）。
    所以只读不能只靠猜，连接本身就得写不进去：

    - SQLite：以 mode=ro 打开文件，写操作由 SQLite 自己拒绝。
    - PostgreSQL：会话的默认事务只读。能关掉它的 SET / set_config() 守卫都拦。
    - MySQL / MariaDB：会话事务只读。能关掉它的只有 SET，守卫拦。
    - Oracle：没有会话级只读开关，只剩守卫这一道——只读数据源请配只读账号。
    """
    url = build_url(source, reveal=True)
    if not source.readonly:
        return url, {}
    kind = (source.kind or "").lower()
    if kind == "sqlite":
        path = source.database or ""
        if path and path != ":memory:":
            from urllib.parse import quote
            url = f"{_DRIVERS['sqlite']}:///file:{quote(path)}?mode=ro&uri=true"
        return url, {}
    if kind in ("postgres", "postgresql"):
        return url, {"connect_args": {"server_settings": {"default_transaction_read_only": "on"}}}
    if kind in ("mysql", "mariadb"):
        return url, {"connect_args": {"init_command": "SET SESSION TRANSACTION READ ONLY"}}
    return url, {}


def _timeout_value(value: Any) -> float | None:
    """options 里的查询时限：表单存的是字符串，也认数字。填得不对返回 None。"""
    if value is None or isinstance(value, bool):
        return None
    try:
        seconds = float(str(value).strip())
    except ValueError:
        return None
    if not MIN_QUERY_TIMEOUT_S <= seconds <= MAX_QUERY_TIMEOUT_S:
        return None
    return int(seconds) if seconds.is_integer() else seconds


def query_timeout_problem(options: dict[str, Any] | None) -> str | None:
    """保存数据源时查一下查询时限填得对不对；没填、或者填对了返回 None。"""
    value = (options or {}).get(QUERY_TIMEOUT_OPTION)
    if value is None or str(value).strip() == "":
        return None
    if _timeout_value(value) is None:
        return (f"查询时限要填 {MIN_QUERY_TIMEOUT_S} 到 {MAX_QUERY_TIMEOUT_S} 之间的秒数，"
                f"比如 60；现在填的是「{value}」")
    return None


def masked_columns(options: dict[str, Any] | None) -> list[str]:
    """options.mask_columns → 列名列表（去空白、去重，保持顺序）。列表和「a, b」这样的文字都认。"""
    raw = (options or {}).get(MASK_COLUMNS_OPTION)
    items = _MASK_SPLIT.split(raw) if isinstance(raw, str) else raw if isinstance(raw, list) else []
    return list(dict.fromkeys(str(c).strip() for c in items if isinstance(c, (str, int)) and str(c).strip()))


def mask_columns_problem(options: dict[str, Any] | None) -> str | None:
    """保存数据源时查一下遮罩列填得对不对；没填、或者填对了返回 None。"""
    raw = (options or {}).get(MASK_COLUMNS_OPTION)
    if raw is None or raw == "" or raw == []:
        return None
    if isinstance(raw, str) or (isinstance(raw, list) and all(isinstance(c, str) for c in raw)):
        return None
    return f"遮罩的列要填列名，用逗号分开，比如「phone, email」；现在填的是「{raw}」"


def query_timeout(source: Any) -> float:
    """这个数据源上一条查询最多跑多久（秒）。没配、或者库里存着的值不对，用缺省。

    缺省按模块属性现取：只有 guard.QueryLimits 这一处定义。
    """
    value = _timeout_value((source.options or {}).get(QUERY_TIMEOUT_OPTION))
    return value if value is not None else guard.QueryLimits().timeout_seconds


def _seconds(value: float) -> str:
    return str(int(value)) if float(value).is_integer() else f"{value:g}"


class EngineCache:
    """按数据源 id 缓存 engine。连接池的建立不便宜，不该每次查询都重来一遍。

    配置变了要显式 invalidate——改了密码却还在用旧连接，排查起来很费神；
    改了只读开关却还在用旧连接，只读就形同虚设。
    """

    def __init__(self) -> None:
        self._engines: dict[str, AsyncEngine] = {}
        self._lock = asyncio.Lock()

    async def get(self, source: Any) -> AsyncEngine:
        key = str(source.id)
        cached = self._engines.get(key)
        if cached is not None:
            return cached
        async with self._lock:
            if key in self._engines:
                return self._engines[key]
            url, extra = engine_args(source)
            engine = create_async_engine(
                url,
                pool_size=3,
                max_overflow=2,
                pool_pre_ping=True,   # 长时间空闲后连接会被数据库掐掉，先探活再用
                pool_recycle=1800,
                **extra,
            )
            self._engines[key] = engine
            return engine

    async def invalidate(self, source_id: str) -> None:
        engine = self._engines.pop(str(source_id), None)
        if engine is not None:
            await engine.dispose()

    async def close(self) -> None:
        for engine in list(self._engines.values()):
            await engine.dispose()
        self._engines.clear()


engines = EngineCache()


async def test_connection(source: Any) -> dict[str, Any]:
    """测连接。失败时把驱动的原始错误带出来——这类问题九成靠错误信息定位。"""
    started = time.perf_counter()
    try:
        engine = await engines.get(source)
        async with engine.connect() as conn:
            probe = "SELECT 1 FROM DUAL" if (source.kind or "").lower() == "oracle" else "SELECT 1"
            await conn.execute(text(probe))
        return {
            "ok": True,
            "elapsed_ms": int((time.perf_counter() - started) * 1000),
            "url": build_url(source),  # 遮掉密码的版本
        }
    except Exception as e:  # noqa: BLE001
        return {
            "ok": False,
            "error": f"{type(e).__name__}: {e}",
            "elapsed_ms": int((time.perf_counter() - started) * 1000),
            "url": build_url(source),
        }


#: 数据库按时限停下语句之后，数据层自己的兜底再多等多久。要比引擎的宽限
#: （toolcalls.GRACE_S）短：数据库报上来的「超过 Ns 被中断」才是给人看的那句，
#: 引擎那边只在数据层也卡住时才出面
BACKSTOP_S = 0.5

#: SQLite 每执行多少条虚拟机指令问一次进度回调。回调只比一次时间，几毫秒一问足够准
_SQLITE_STEPS = 10_000

#: 各家数据库按时限停下语句时的报错。认出来就说「超时」，而不是把驱动原文当成
#: SQL 写错了交给模型——它会去改一条本来没错的 SQL。
#:
#: 先认错误码：服务器的报错文字会跟着它的语言设置走，中文环境的 PostgreSQL 报的是
#: 「由于语句执行超时，正在取消查询命令」，按英文原文找一个也认不出来
_PG_CANCELED = "57014"                    # PostgreSQL query_canceled（statement_timeout 到点）
_MYSQL_TIMEOUT_CODES = frozenset({3024,   # MySQL：max_execution_time 到点
                                  1969})  # MariaDB：max_statement_time 到点
_TIMEOUT_MARKERS = (
    "statement timeout",                   # PostgreSQL（英文环境）
    "maximum statement execution time",    # MySQL 3024
    "max_statement_time",                  # MariaDB 1969
    "dpy-4024", "dpi-1067", "call timeout",  # python-oracledb：call timeout of N ms exceeded（驱动自己的话，不随服务器语言变）
)


def _causes(e: BaseException) -> list[BaseException]:
    """异常连同它包着的原始驱动异常。SQLAlchemy 包一层（.orig），适配器再包一层（__cause__），
    有的路径（服务端游标取行）又原样抛出驱动异常，所以几条链都走一遍。"""
    out: list[BaseException] = []
    todo: list[Any] = [e]
    while todo and len(out) < 8:
        err = todo.pop(0)
        if isinstance(err, BaseException) and err not in out:
            out.append(err)
            todo += [getattr(err, "orig", None), err.__cause__, err.__context__]
    return out


class _Deadline:
    """一条语句的时限。数据库那边按它停下；停下时报的错据此认成超时。

    从取连接之前起算，和后端兜底同一个起点：取连接花掉的也算在时限里，数据库只拿到
    剩下的那段（remaining_ms）。否则取连接一慢，兜底就抢在数据库前面取消语句。
    """

    def __init__(self, seconds: float) -> None:
        self.seconds = seconds
        self.ends = time.monotonic() + seconds
        #: SQLite 的进度回调真的掐过
        self.fired = False
        #: 语句已经结束。回调这时要是还挂在连接上（摘除没成功），也不能再掐别人的语句
        self.over = False

    def remaining_ms(self) -> int:
        return max(1, int((self.ends - time.monotonic()) * 1000))

    def remaining_seconds(self) -> float:
        return max(0.001, round(self.ends - time.monotonic(), 3))

    def expired(self) -> bool:
        return time.monotonic() >= self.ends

    def sqlite_progress(self) -> int:
        # 在 SQLite 执行语句的那个线程里被调用；返回非零，SQLite 就中止当前语句
        if not self.over and time.monotonic() >= self.ends:
            self.fired = True
            return 1
        return 0

    def stopped(self, e: BaseException) -> bool:
        """这个报错是不是数据库按时限停下语句造成的。"""
        if self.fired:
            return True
        for err in _causes(e):
            if _PG_CANCELED in (getattr(err, "sqlstate", None), getattr(err, "pgcode", None)):
                return True
            args = getattr(err, "args", None) or ()
            if args and isinstance(args[0], int) and args[0] in _MYSQL_TIMEOUT_CODES:
                return True
            low = str(err).lower()
            if any(marker in low for marker in _TIMEOUT_MARKERS):
                return True
            # 各家的措辞不一（interrupted / canceled，或者干脆只在类名里）；
            # 过了时限才报出来的，就是时限掐的
            if self.expired() and any(w in f"{type(err).__name__} {low}".lower()
                                      for w in ("interrupt", "cancel")):
                return True
        return False


async def _driver(conn: AsyncConnection) -> Any:
    return (await conn.get_raw_connection()).driver_connection


@asynccontextmanager
async def _server_deadline(conn: AsyncConnection, kind: str, deadline: _Deadline) -> AsyncIterator[None]:
    """让数据库自己在时限处停下这条语句。

    只靠后端这边 wait_for 放弃等待是不够的：数据库那边的语句照样在跑，占着连接池
    里的一个连接直到跑完（Oracle 实测多占了 60 秒）；aiosqlite 更糟，语句在线程里
    跑，取消之后关连接要排在它后面，一条不收敛的查询能让调用永远回不来。

    会话级的设置用完要恢复：连接会回到池子里，探查结构这类长操作也用它。
    下发失败（老版本的 MySQL 没有这个变量）不挡查询，只剩后端这边的兜底。
    """
    if kind in ("postgres", "postgresql"):
        # SET LOCAL 只管这一个事务：连接还回池子时回滚，设置随之作废
        await conn.execute(text(f"SET LOCAL statement_timeout = {deadline.remaining_ms()}"))
        yield
        return

    if kind in ("mysql", "mariadb"):
        # MySQL 的 max_execution_time 以毫秒计、只管 SELECT；MariaDB 没有它，
        # 用 max_statement_time（秒）。按连上的服务器认，不按表单里选的类型
        if getattr(conn.dialect, "is_mariadb", False):
            var, value = "max_statement_time", _seconds(deadline.remaining_seconds())
        else:
            var, value = "max_execution_time", str(deadline.remaining_ms())
        try:
            await conn.execute(text(f"SET SESSION {var} = {value}"))
        except Exception as e:  # noqa: BLE001
            logger.warning("MySQL 语句时限没下发成功（%s），只剩后端兜底", e)
            yield
            return
        cancelled = False
        try:
            yield
        except asyncio.CancelledError:
            # 连接还卡在那条语句上，这时再发一条只会排在它后面
            cancelled = True
            raise
        finally:
            if not cancelled:
                try:
                    await conn.execute(text(f"SET SESSION {var} = DEFAULT"))
                except Exception:  # noqa: BLE001 - 连接已经坏了的话，连接池会把它作废
                    pass
        return

    if kind == "oracle":
        driver = await _driver(conn)
        before = getattr(driver, "call_timeout", 0)
        try:
            driver.call_timeout = deadline.remaining_ms()
        except Exception as e:  # noqa: BLE001
            logger.warning("Oracle call_timeout 没设上（%s），只剩后端兜底", e)
            yield
            return
        try:
            yield
        finally:
            try:
                driver.call_timeout = before
            except Exception:  # noqa: BLE001 - 超时后连接可能已经被驱动关掉
                pass
        return

    if kind == "sqlite":
        driver = await _driver(conn)
        await driver.set_progress_handler(deadline.sqlite_progress, _SQLITE_STEPS)
        try:
            yield
        finally:
            deadline.over = True
            try:
                await driver.set_progress_handler(None, 0)
            except Exception:  # noqa: BLE001 - over 已经让它不再掐人
                pass
        return

    yield


async def _bounded(
    opener: Callable[[], Any], kind: str, limits: QueryLimits, over: str,
    work: Callable[[AsyncConnection], Awaitable[Any]],
) -> Any:
    """在 opener() 开的连接上跑 work，时限交给数据库执行，后端这边再留一道兜底。

    数据库按时限停下语句时报的错，翻成 over 这句话：「超过 Ns 被中断」。
    两道时限从同一刻（取连接之前）算起，数据库那道先到，兜底晚 BACKSTOP_S。
    """
    deadline: _Deadline | None = None

    async def _exec() -> Any:
        nonlocal deadline
        deadline = _Deadline(limits.timeout_seconds)
        async with opener() as conn:
            # 光等连接就把时限用完了：语句不必再发，发出去也只剩 1ms
            if deadline.expired():
                raise SqlRejected(over)
            async with _server_deadline(conn, kind, deadline):
                return await work(conn)

    try:
        return await asyncio.wait_for(_exec(), timeout=limits.timeout_seconds + BACKSTOP_S)
    except asyncio.TimeoutError as e:
        raise SqlRejected(over) from e
    except Exception as e:
        if deadline is not None and deadline.stopped(e):
            raise SqlRejected(over) from e
        raise


async def run_query(
    source: Any, sql: str, *, limits: QueryLimits | None = None
) -> QueryResult:
    """执行一条查询。守卫先过，再谈执行。

    注意执行的是 guard.check 的返回值而不是原始输入——守卫会把尾分号之类
    规范掉，执行原文等于绕过了守卫。

    时限没指定时按数据源自己的配置（query_timeout），由数据库执行（_server_deadline）。
    """
    limits = limits or QueryLimits(timeout_seconds=query_timeout(source))
    kind = (source.kind or "").lower()
    # 按数据源的方言扫：MySQL 的反斜杠转义、PG 的 $$ 字符串，守卫要和数据库读得一样
    statement = check(sql, readonly=bool(source.readonly), source_name=source.name, dialect=kind)

    engine = await engines.get(source)
    started = time.perf_counter()

    if not source.readonly and is_write(statement, dialect=kind):
        return await _run_write(engine, kind, statement, limits, started)

    async def _read(conn: AsyncConnection) -> tuple[list[str], list[list[Any]], bool, dict[str, str]]:
        cursor = await conn.stream(text(statement))
        columns = list(cursor.keys())
        rows: list[list[Any]] = []
        seen = _TypeTally(len(columns))
        truncated = False
        size = 0
        async for row in cursor:
            # 逐行累加，边收边判上限：一次性 fetchall 一个亿级表就晚了
            seen.add(row)
            values = [_jsonable(v) for v in row]
            rows.append(values)
            size += len(json.dumps(values, ensure_ascii=False, default=str))
            if len(rows) >= limits.max_rows or size >= limits.max_bytes:
                truncated = True
                break
        return columns, rows, truncated, seen.types(columns)

    columns, rows, truncated, types = await _bounded(
        engine.connect, kind, limits,
        f"查询超过 {_seconds(limits.timeout_seconds)}s 被中断。加上 WHERE 条件或 LIMIT 缩小范围。",
        _read,
    )
    return QueryResult(
        columns=columns,
        rows=rows,
        row_count=len(rows),
        truncated=truncated,
        elapsed_ms=int((time.perf_counter() - started) * 1000),
        sql=statement,
        column_types=types,
    )


class _TypeTally:
    """边读边记每列见过的值类型。只看驱动交回的原始值：_jsonable 之后 Decimal 就成了字符串，
    小数位为 0 的 DECIMAL（"45678"）和文本列的 "2026" 再也分不开。"""

    def __init__(self, width: int) -> None:
        self._kinds: list[set[str]] = [set() for _ in range(width)]

    def add(self, row: Any) -> None:
        for kinds, value in zip(self._kinds, row):
            if value is not None:
                kinds.add(_value_kind(value))

    def types(self, columns: list[str]) -> dict[str, str]:
        return {str(name): kind for name, kinds in zip(columns, self._kinds)
                if len(kinds) == 1 and (kind := next(iter(kinds))) != "other"}


def _value_kind(value: Any) -> str:
    from datetime import date, datetime, time as dtime
    from decimal import Decimal

    if isinstance(value, bool):
        return "boolean"
    if isinstance(value, (int, float, Decimal)):
        return "number"
    if isinstance(value, datetime):
        return "datetime"
    if isinstance(value, date):
        return "date"
    if isinstance(value, dtime):
        return "time"
    return "text" if isinstance(value, str) else "other"


def column_types(columns: list[str], rows: list[Any]) -> dict[str, str]:
    """一批原始行里每列的类型：number / text / date / datetime / time / boolean。

    全是空值、类型混着（同一列里有数也有文本）的列不记——拿不准就不说，下游照旧按值猜。
    """
    tally = _TypeTally(len(columns))
    for row in rows:
        tally.add(row)
    return tally.types(columns)


async def _run_write(
    engine: AsyncEngine, kind: str, statement: str, limits: QueryLimits, started: float
) -> QueryResult:
    """可写源上的写操作：真的提交。

    以前写和查询走同一条 stream 路径：拿不到结果集就抛 ResourceClosedError，
    连接关闭时再回滚——可写源上的写从来没落过库，调用方只看到一个报错。
    能走到这里的写都已经过了审批（registry.call_is_dangerous → 各节点的审批关卡），
    只读源上的写在 check 里就被拒了。
    """

    async def _write(conn: AsyncConnection) -> tuple[list[str], list[list[Any]]]:
        result = await conn.execute(text(statement))
        if result.returns_rows:          # RETURNING、PG 的数据修改 CTE
            return (list(result.keys()),
                    [[_jsonable(v) for v in row] for row in result.fetchmany(limits.max_rows)])
        return ["affected_rows"], [[result.rowcount]]

    # engine.begin()：正常退出提交，出错（包括数据库按时限停下）回滚
    columns, rows = await _bounded(
        engine.begin, kind, limits,
        f"写操作超过 {_seconds(limits.timeout_seconds)}s 被中断，已回滚。", _write,
    )
    return QueryResult(
        columns=columns,
        rows=rows,
        row_count=len(rows),
        truncated=len(rows) >= limits.max_rows,
        elapsed_ms=int((time.perf_counter() - started) * 1000),
        sql=statement,
    )


def _jsonable(value: Any) -> Any:
    """把驱动返回的类型转成能进 JSON 的形态。

    Decimal、date、datetime、bytes 这些直接 json.dumps 会炸，而它们在真实业务
    数据里遍地都是——金额是 Decimal、时间是 datetime，不处理等于不可用。
    """
    from datetime import date, datetime, time as dtime
    from decimal import Decimal

    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, Decimal):
        # 转 float 会丢精度，而金额恰恰不能丢。字符串保真，下游要算再自己转。
        return str(value)
    if isinstance(value, (datetime, date, dtime)):
        return value.isoformat()
    if isinstance(value, (bytes, bytearray, memoryview)):
        raw = bytes(value)
        try:
            return raw.decode()
        except UnicodeDecodeError:
            return f"<binary {len(raw)} bytes>"
    return str(value)