AgentLab/backend/app/db/postgres.py

"""应用库是 PostgreSQL 时，让它在 SQLite 上写成的那些假设下照样成立。

只装在应用库那个引擎的方言实例上（db/base.make_engine），业务数据源的 PostgreSQL 引擎不受影响。

**建表的列类型。** 应用是在 SQLite 上写成的：SQLite 不查 VARCHAR 的长度，INTEGER 是 64 位。PostgreSQL 上
VARCHAR(n) 超长直接报错、INTEGER 只有 32 位。模型里的 String(n) 是写给人看的「大致多长」，从来没有被库约束过：
逐列核对只核得到测试走到的值，漏一处就是线上的写入失败，从 SQLite 导过来的存量数据里也可能已经有超长的。所以
应用库上 String(n) 建成不限长的 VARCHAR（PostgreSQL 里 varchar、varchar(n)、text 存储和性能一样，限长只多一次
检查）。CAST(x AS VARCHAR(n)) 这种查询里的写法不动——那是截断的意思。Integer 列在 PostgreSQL 上是 BIGINT：
那是在模型这一层换的（db/base._portable_integer），绑定参数的类型转换跟着列走，大数不在半路溢出。

**文本怎么比大小、怎么排序。** SQLite 缺省按字节比（BINARY）；PostgreSQL 按建库时的区域设置（常见的 en_US.UTF-8 不分
大小写地先比字母、忽略标点），同一列排出来的先后、`key < '…'` 这种比较的结果都会变。应用库上的文本列建表时一律带
COLLATE "C"（按字节比），和 SQLite 一致，不看建库时选了什么区域；LIKE 'xx%' 这种前缀匹配也因此用得上普通索引。
代价是列上的 ILIKE、lower() 也按 C 的规则，只认 ASCII 的大小写（É/é、全角Ａ/ａ 不认）：要不分大小写地找，用
db/portable.contains_ci，模式带 COLLATE "default"，按建库时的区域设置折大小写。

**ORDER BY 里 NULL 的位置。** SQLite 升序 NULL 在前、降序 NULL 在后，PostgreSQL 正好相反。应用里两百多处排序，
可空的列（开始时间、到期时间、发布时间……）排出来的先后要和 SQLite 一致：可空列上没写 nulls_first / nulls_last
的 ASC 补 NULLS FIRST、DESC 补 NULLS LAST，ORDER BY 里直接写列名（隐含 ASC）的同样补上。不可空的列不动——
补了会让索引用不上，又没有 NULL 可排。窗口函数 OVER (ORDER BY …) 里同样补。

只认直接写的列（Column，含它的 asc() / desc()）。这几种不补，两种库排出来的先后可能不同，要紧的地方自己写明
nulls_first / nulls_last：label 出来的聚合（max(...) 之类）、coalesce / case 这类表达式、外连接进来的表上的列（模型里
不可空，外连接没连上时却是 NULL）、原生 text() 里的 ORDER BY（db/base 的去重语句里写明了）。
"""

from __future__ import annotations

from typing import Any

from sqlalchemy import Column
from sqlalchemy.dialects.postgresql.asyncpg import PGCompiler_asyncpg
from sqlalchemy.dialects.postgresql.base import PGTypeCompiler
from sqlalchemy.sql import operators
from sqlalchemy.sql.elements import ClauseList, UnaryExpression

#: 外层已经写了 nulls_first / nulls_last：里面那一层 ASC / DESC 不再补
_EXPLICIT = "_agentlab_explicit_nulls"


def _nullable(element: Any) -> bool:
    return isinstance(element, Column) and bool(element.nullable) and not element.primary_key


#: 建表时文本列的排序规则：按字节比，同 SQLite 的 BINARY（见模块说明）
_COLLATE = ' COLLATE "C"'


class AppTypeCompiler(PGTypeCompiler):
    def visit_string(self, type_: Any, **kw: Any) -> str:
        if isinstance(kw.get("type_expression"), Column):
            return ("VARCHAR" if type_.length else super().visit_string(type_, **kw)) + _COLLATE
        return super().visit_string(type_, **kw)

    def visit_text(self, type_: Any, **kw: Any) -> str:
        text = super().visit_text(type_, **kw)
        return text + _COLLATE if isinstance(kw.get("type_expression"), Column) else text


def _ordered(clauses: Any) -> list[Any]:
    """ORDER BY 里直接写的可空列补成 ASC（visit_unary 再给它补 NULLS FIRST）。"""
    return [c.asc() if _nullable(c) else c for c in clauses]


class AppCompiler(PGCompiler_asyncpg):
    def order_by_clause(self, select: Any, **kw: Any) -> str:
        order_by = self._generate_delimited_list(_ordered(select._order_by_clauses), ", ", **kw)
        return " ORDER BY " + order_by if order_by else ""

    def visit_over(self, over: Any, **kw: Any) -> str:
        clauses = list(over.order_by) if over.order_by is not None else []
        if any(_nullable(c) for c in clauses):
            # 在副本上换：语句对象可能被缓存、被别处共用，不能改原件
            over = over._clone()
            over.order_by = ClauseList(*_ordered(clauses))
        return super().visit_over(over, **kw)

    def visit_unary(self, unary: UnaryExpression, **kw: Any) -> str:
        modifier = unary.modifier
        if modifier in (operators.nulls_first_op, operators.nulls_last_op):
            return super().visit_unary(unary, **{**kw, _EXPLICIT: True})
        explicit = kw.pop(_EXPLICIT, False)
        text = super().visit_unary(unary, **kw)
        if not explicit and modifier in (operators.asc_op, operators.desc_op) and _nullable(unary.element):
            text += " NULLS FIRST" if modifier is operators.asc_op else " NULLS LAST"
        return text


def _has_nul(value: Any) -> bool:
    return isinstance(value, str) and "\x00" in value


def _clean_row(params: Any) -> Any:
    if isinstance(params, dict):
        return {k: v.replace("\x00", "") if _has_nul(v) else v for k, v in params.items()} \
            if any(_has_nul(v) for v in params.values()) else params
    if isinstance(params, (list, tuple)) and any(_has_nul(v) for v in params):
        return type(params)(v.replace("\x00", "") if _has_nul(v) else v for v in params)
    return params


def strip_nul(conn: Any, cursor: Any, statement: str, parameters: Any, context: Any,
              executemany: bool) -> tuple[str, Any]:
    """写进 PostgreSQL 之前去掉文本参数里的 NUL 字符（before_cursor_execute，retval=True）。

    PostgreSQL 的 text / varchar 不收 \\x00（报 invalid byte sequence），SQLite 照存：一段带 NUL 的模型输出、一个从二进制
    文件里粘出来的问题，在 SQLite 上写得进去，在 PostgreSQL 上整条写入失败。所以和 JSON 列（db/base.pg_json_dumps）一样
    去掉。只看参数里的字符串，bytes（LargeBinary）不动；没有 NUL 的参数原样交回，不多一次拷贝。"""
    if executemany:
        if any(_clean_row(row) is not row for row in parameters):
            return statement, [_clean_row(row) for row in parameters]
        return statement, parameters
    return statement, _clean_row(parameters)


def install(dialect: Any) -> None:
    """装到应用库引擎的方言实例上。"""
    dialect.statement_compiler = AppCompiler
    dialect.type_compiler_instance = dialect.type_compiler = AppTypeCompiler(dialect)