From f042452fe64a017b471ec9bd662815391ebd75c5 Mon Sep 17 00:00:00 2001 From: Donny Date: Thu, 5 Nov 2026 02:30:11 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E8=87=AA=E7=A0=94=20agtype=20=E8=A7=A3?= =?UTF-8?q?=E6=9E=90=E4=B8=8E=20psycopg3=20loader?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - graph/agtype.py：解析 ::vertex / ::edge / ::path 后缀，返回 Vertex/Edge/Path - register_agtype() 直接查 ag_catalog.agtype 的 OID 注册 loader（TypeInfo.fetch 取不到） - 18 条单元测试，样本全部从真实 AGE 库抓取；7 条畸形输入用例 - tests/test_agtype_live.py：连真库的集成测试，默认跳过（ADAPTER_TEST_DSN） - AGENTS.md 技术栈表：图驱动改为自研解析，注明不用 apache-age-python 的原因 --- AGENTS.md | 2 +- src/adapter_python_pg/__init__.py | 27 +-- src/adapter_python_pg/graph/agtype.py | 251 ++++++++++++++++++++++++++ tests/test_agtype.py | 104 +++++++++++ tests/test_agtype_live.py | 73 ++++++++ 5 files changed, 443 insertions(+), 14 deletions(-) create mode 100644 src/adapter_python_pg/graph/agtype.py create mode 100644 tests/test_agtype.py create mode 100644 tests/test_agtype_live.py diff --git a/AGENTS.md b/AGENTS.md index 34753bf..29c4e2b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -35,7 +35,7 @@ app.include_router(build_router(adapter), prefix="/ontology") | 包管理 | **uv**（`uv.lock` 入库） | 不用手工 venv + pip | | 构建 | hatchling | `src/` 布局，带 `py.typed` | | 驱动 | **psycopg 3**（async + `psycopg_pool`） | 不用 asyncpg、不用任何 ORM | -| 图 | **`age`**（AGE 官方 psycopg3 驱动） | `AgeLoader` 处理 agtype，不手写解析 | +| 图 | **自研 agtype 解析**（`graph/agtype.py`），注册 psycopg3 loader | 不用 PyPI 的 `apache-age-python`（0.0.7 依赖 psycopg2+antlr4）；AGE master 的 psycopg3 驱动未发布 | | 对象 | **pydantic v2** + `create_model()` | 运行期按本体动态建类 | | 本体校验 | pydantic 模型 + 导出 JSON Schema | 本体先自查，schema 供其他语言消费 | | 表达式 | **SQL 下推** + **CEL**（`cel-python`） | 不自研表达式解析器，不 `eval` | diff --git a/src/adapter_python_pg/__init__.py b/src/adapter_python_pg/__init__.py index 2ecb78c..47e4327 100644 --- a/src/adapter_python_pg/__init__.py +++ b/src/adapter_python_pg/__init__.py @@ -5,24 +5,25 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +import importlib +from typing import Any __version__ = "0.1.0" __all__ = ["Adapter", "__version__", "build_router"] -if TYPE_CHECKING: # pragma: no cover - 仅类型检查期可见 - from adapter_python_pg.api.router import build_router - from adapter_python_pg.core import Adapter +_LAZY: dict[str, str] = { + "Adapter": "adapter_python_pg.core", + "build_router": "adapter_python_pg.api.router", +} def __getattr__(name: str) -> Any: - """延迟导入：`import adapter_python_pg` 不应拉起 psycopg / fastapi 之外的任何子系统。""" - if name == "Adapter": - from adapter_python_pg.core import Adapter + """延迟导入：`import adapter_python_pg` 不应拉起任何额外的子系统。 - return Adapter - if name == "build_router": - from adapter_python_pg.api.router import build_router - - return build_router - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + `Adapter` / `build_router` 在实现到位后由 core / api 子包提供；在那之前 + 属性访问会抛 AttributeError，这是刻意的（避免半成品被当成可用）。 + """ + module = _LAZY.get(name) + if module is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(importlib.import_module(module), name) diff --git a/src/adapter_python_pg/graph/agtype.py b/src/adapter_python_pg/graph/agtype.py new file mode 100644 index 0000000..374219c --- /dev/null +++ b/src/adapter_python_pg/graph/agtype.py @@ -0,0 +1,251 @@ +"""agtype 解析：把 Apache AGE 的图值文本解析成 Python 对象。 + +AGE 的类型输出函数把 agtype 序列化成 **JSON + 类型后缀**： + + 顶点 {"id": 844424930131971, "label": "Person", "properties": {"name": "Alice"}}::vertex + 边 {"id": 1125899906842626, "label": "KNOWS", "end_id": ..., "start_id": ..., + "properties": {"since": 2021}}::edge + 路径 [{...}::vertex, {...}::edge, {...}::vertex]::path + 标量 ["Alice", 42, 3.5, true, null, {"k": "v", "n": 1}] + +所以不能直接 `json.loads`：数组里可能夹着带后缀的元素。这里用一个小扫描器， +对象/标量交给 `json.raw_decode`，数组自己走一遍，后缀单独消费。 + +为什么不用 PyPI 上的 `apache-age-python`：那个包停在 0.0.7（2023-09），依赖 psycopg2 + +antlr4，与本项目的 psycopg3 决策冲突；AGE master 里基于 psycopg3 的驱动未发布。 +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Final + +from adapter_python_pg.errors import Configuration, QueryError + +__all__ = [ + "AGTypeError", + "Edge", + "Path", + "Vertex", + "parse_agtype", + "register_agtype", +] + +_VERTEX: Final = "vertex" +_EDGE: Final = "edge" +_PATH: Final = "path" + +_DECODER: Final = json.JSONDecoder() + + +class AGTypeError(QueryError): + """agtype 文本不符合预期结构。""" + + +@dataclass(frozen=True, slots=True) +class Vertex: + """图节点。`id` 是 AGE 内部的 bigint，不等于业务主键。""" + + id: int + label: str + properties: dict[str, Any] + + def get(self, name: str, default: Any = None) -> Any: + return self.properties.get(name, default) + + +@dataclass(frozen=True, slots=True) +class Edge: + """图边。`start_id` → `end_id`，对应 from → to。""" + + id: int + label: str + start_id: int + end_id: int + properties: dict[str, Any] + + def get(self, name: str, default: Any = None) -> Any: + return self.properties.get(name, default) + + +@dataclass(frozen=True, slots=True) +class Path: + """路径：顶点与边交替出现，顺序由 AGE 保证。""" + + elements: tuple[Vertex | Edge, ...] + + @property + def vertices(self) -> tuple[Vertex, ...]: + return tuple(item for item in self.elements if isinstance(item, Vertex)) + + @property + def edges(self) -> tuple[Edge, ...]: + return tuple(item for item in self.elements if isinstance(item, Edge)) + + @property + def start(self) -> Vertex: + return self.vertices[0] + + @property + def end(self) -> Vertex: + return self.vertices[-1] + + +class _Scanner: + """按需推进的扫描器：对象与标量复用 json，数组手工处理注解。""" + + def __init__(self, text: str) -> None: + self.text = text + self.pos = 0 + + def parse(self) -> Any: + value = self._value() + self._skip_ws() + if self.pos != len(self.text): + raise AGTypeError(f"agtype 尾部有多余内容：{self.text[self.pos : self.pos + 32]!r}") + return value + + # -- 内部 -------------------------------------------------------------- + + def _skip_ws(self) -> None: + while self.pos < len(self.text) and self.text[self.pos].isspace(): + self.pos += 1 + + def _value(self) -> Any: + self._skip_ws() + if self.pos >= len(self.text): + raise AGTypeError("agtype 意外结束") + + char = self.text[self.pos] + if char == "[": + return self._array() + if char == "{": + obj, self.pos = _DECODER.raw_decode(self.text, self.pos) + return self._annotated(obj) + value, self.pos = _DECODER.raw_decode(self.text, self.pos) + return self._annotated(value) + + def _array(self) -> Any: + self.pos += 1 # 跳过 '[' + items: list[Any] = [] + self._skip_ws() + if self.pos < len(self.text) and self.text[self.pos] == "]": + self.pos += 1 + return self._annotated(items) + + while True: + items.append(self._value()) + self._skip_ws() + if self.pos >= len(self.text): + raise AGTypeError("数组未闭合") + char = self.text[self.pos] + self.pos += 1 + if char == "]": + break + if char != ",": + raise AGTypeError(f"数组分隔符非法：{char!r}") + return self._annotated(items) + + def _annotated(self, value: Any) -> Any: + """消费形如 `::vertex` 的后缀注解，并按注解包装值。""" + self._skip_ws() + if not self.text.startswith("::", self.pos): + return value + self.pos += 2 + start = self.pos + while self.pos < len(self.text) and ( + self.text[self.pos].isalnum() or self.text[self.pos] == "_" + ): + self.pos += 1 + tag = self.text[start : self.pos] + return _wrap(tag, value) + + +def _wrap(tag: str, value: Any) -> Any: + if tag == _VERTEX: + return _to_vertex(value) + if tag == _EDGE: + return _to_edge(value) + if tag == _PATH: + if not isinstance(value, list): + raise AGTypeError("::path 只应作用于数组") + return Path(elements=tuple(value)) + # 其它注解（如 ::numeric）对调用方无意义，去掉即可 + return value + + +def _require(obj: Any, key: str, kind: str) -> Any: + if not isinstance(obj, dict): + raise AGTypeError(f"{kind} 不是 JSON 对象：{type(obj).__name__}") + if key not in obj: + raise AGTypeError(f"{kind} 缺少字段 {key!r}：{sorted(obj)}") + return obj[key] + + +def _properties(obj: dict[str, Any], kind: str) -> dict[str, Any]: + props = _require(obj, "properties", kind) + if props is None: + return {} + if not isinstance(props, dict): + raise AGTypeError(f"{kind}.properties 不是对象：{type(props).__name__}") + return props + + +def _to_vertex(value: Any) -> Vertex: + return Vertex( + id=int(_require(value, "id", _VERTEX)), + label=str(_require(value, "label", _VERTEX)), + properties=_properties(value, _VERTEX), + ) + + +def _to_edge(value: Any) -> Edge: + return Edge( + id=int(_require(value, "id", _EDGE)), + label=str(_require(value, "label", _EDGE)), + start_id=int(_require(value, "start_id", _EDGE)), + end_id=int(_require(value, "end_id", _EDGE)), + properties=_properties(value, _EDGE), + ) + + +def parse_agtype(text: str | bytes | bytearray | memoryview) -> Any: + """把 agtype 文本解析成 Vertex / Edge / Path / list / 标量。""" + if isinstance(text, memoryview): + text = text.tobytes() + if isinstance(text, (bytes, bytearray)): + text = text.decode("utf-8") + try: + return _Scanner(text).parse() + except AGTypeError: + raise + except (ValueError, TypeError) as exc: # json 层的错误统一成 AGTypeError + raise AGTypeError(f"agtype 解析失败：{exc}") from exc + + +def register_agtype(conn: Any) -> None: + """在当前连接上注册 agtype 的 psycopg3 loader。 + + 直接用 `pg_type` 查 OID：agtype 位于 `ag_catalog` schema，而新连接的 + `search_path` 里没有它，`TypeInfo.fetch(conn, "agtype")` 会返回 None。 + 缺扩展、连错库时抛 `Configuration`；调用方应在连接池的 `configure` 回调里调用。 + """ + from psycopg.adapt import Loader + + row = conn.execute( + """ + SELECT t.oid + FROM pg_type AS t + JOIN pg_namespace AS n ON n.oid = t.typnamespace + WHERE t.typname = 'agtype' AND n.nspname = 'ag_catalog' + """ + ).fetchone() + if row is None: + raise Configuration("当前库没有 ag_catalog.agtype：AGE 扩展未安装或未加载") + + class _Loader(Loader): + def load(self, data: bytes | bytearray | memoryview) -> Any: + return parse_agtype(data) + + conn.adapters.register_loader(row[0], _Loader) diff --git a/tests/test_agtype.py b/tests/test_agtype.py new file mode 100644 index 0000000..9a4ef73 --- /dev/null +++ b/tests/test_agtype.py @@ -0,0 +1,104 @@ +"""agtype 解析测试。 + +下面的原始文本全部是从真实 AGE 库（pg-multimodel / demo_graph）里 `select v` 抓下来的， +不是手写的近似值——格式一旦对不上，AGE 升级就会在这里先报错。 +""" + +from __future__ import annotations + +import pytest + +from adapter_python_pg.graph.agtype import ( + AGTypeError, + Edge, + Path, + Vertex, + parse_agtype, +) + +VERTEX_RAW = '{"id": 844424930131971, "label": "Person", "properties": {"name": "Alice"}}::vertex' +EDGE_RAW = ( + '{"id": 1125899906842626, "label": "KNOWS", "end_id": 844424930131972, ' + '"start_id": 844424930131971, "properties": {"since": 2021}}::edge' +) +PATH_RAW = ( + '[{"id": 844424930131971, "label": "Person", "properties": {"name": "Alice"}}::vertex, ' + '{"id": 1125899906842626, "label": "KNOWS", "end_id": 844424930131972, ' + '"start_id": 844424930131971, "properties": {"since": 2021}}::edge, ' + '{"id": 844424930131972, "label": "Person", "properties": {"name": "Bob"}}::vertex]::path' +) +SCALARS_RAW = '["Alice", 42, 3.5, true, null, {"k": "v", "n": 1}]' + + +def test_parse_vertex() -> None: + value = parse_agtype(VERTEX_RAW) + assert isinstance(value, Vertex) + assert value.id == 844424930131971 + assert value.label == "Person" + assert value.properties == {"name": "Alice"} + assert value.get("name") == "Alice" + assert value.get("missing", "fallback") == "fallback" + + +def test_parse_edge() -> None: + value = parse_agtype(EDGE_RAW) + assert isinstance(value, Edge) + assert value.id == 1125899906842626 + assert value.label == "KNOWS" + assert value.start_id == 844424930131971 + assert value.end_id == 844424930131972 + assert value.properties == {"since": 2021} + assert value.get("since") == 2021 + + +def test_parse_path() -> None: + value = parse_agtype(PATH_RAW) + assert isinstance(value, Path) + assert len(value.elements) == 3 + assert len(value.vertices) == 2 + assert len(value.edges) == 1 + assert value.start.properties["name"] == "Alice" + assert value.end.properties["name"] == "Bob" + assert value.edges[0].label == "KNOWS" + + +def test_parse_scalars_and_nested() -> None: + value = parse_agtype(SCALARS_RAW) + assert value == ["Alice", 42, 3.5, True, None, {"k": "v", "n": 1}] + + +def test_parse_accepts_bytes_and_memoryview() -> None: + assert isinstance(parse_agtype(VERTEX_RAW.encode()), Vertex) + assert isinstance(parse_agtype(memoryview(VERTEX_RAW.encode())), Vertex) + + +def test_empty_list() -> None: + assert parse_agtype("[]") == [] + + +def test_empty_properties_defaults_to_dict() -> None: + value = parse_agtype('{"id": 1, "label": "X", "properties": {}}::vertex') + assert isinstance(value, Vertex) + assert value.properties == {} + + +def test_unknown_annotation_is_stripped() -> None: + # 形如 ::numeric 的注解对调用方没有意义，直接去掉 + assert parse_agtype("42::numeric") == 42 + + +@pytest.mark.parametrize( + ("raw", "reason"), + [ + ('{"id": 1, "label": "X"}::vertex', "缺少 properties"), + ('{"label": "X", "properties": {}}::vertex', "缺少 id"), + ('{"id": 1, "label": "X", "properties": {}}::edge', "边缺少 start_id"), + ('["a", "b"::vertex', "数组未闭合"), + ('["a" "b"]', "数组分隔符非法"), + ("[] trailing", "尾部有多余内容"), + ('{"id": 1, "label": "X", "properties": {}}::path', "path 应为数组"), + ], +) +def test_malformed_input_raises(raw: str, reason: str) -> None: + with pytest.raises(AGTypeError): + parse_agtype(raw) diff --git a/tests/test_agtype_live.py b/tests/test_agtype_live.py new file mode 100644 index 0000000..e0fa716 --- /dev/null +++ b/tests/test_agtype_live.py @@ -0,0 +1,73 @@ +"""agtype loader 的集成测试：连真实 AGE 库，验证 psycopg3 侧的类型注册。 + +默认跳过（需要环境变量 `ADAPTER_TEST_DSN`）： + ADAPTER_TEST_DSN="host=127.0.0.1 port=5432 dbname=multimodel user=postgres password=..." \ + uv run pytest -m integration +""" + +from __future__ import annotations + +import os + +import pytest + +from adapter_python_pg.graph.agtype import Edge, Path, Vertex, register_agtype + +DSN = os.environ.get("ADAPTER_TEST_DSN") + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif(not DSN, reason="未设置 ADAPTER_TEST_DSN，跳过集成测试"), +] + +CYPHER_VERTEX = """ +SELECT v FROM cypher('demo_graph', $$ MATCH (p:Person {name:'Alice'}) RETURN p $$) AS (v agtype) +""" +CYPHER_EDGE = """ +SELECT e FROM cypher('demo_graph', $$ + MATCH (:Person {name:'Alice'})-[r:KNOWS]->(:Person) RETURN r +$$) AS (e agtype) +""" +CYPHER_PATH = """ +SELECT p FROM cypher('demo_graph', $$ + MATCH p=(:Person {name:'Alice'})-[:KNOWS]->(:Person) RETURN p LIMIT 1 +$$) AS (p agtype) +""" + + +@pytest.fixture() +def conn() -> object: + psycopg = pytest.importorskip("psycopg") + + with psycopg.connect(DSN) as connection: + register_agtype(connection) + with connection.cursor() as cur: + cur.execute('SET search_path = ag_catalog, "$user", public') + yield connection + + +def _one(conn: object, sql: str) -> object: + with conn.cursor() as cur: # type: ignore[attr-defined] + cur.execute(sql) + return cur.fetchone()[0] + + +def test_vertex_roundtrip(conn: object) -> None: + value = _one(conn, CYPHER_VERTEX) + assert isinstance(value, Vertex) + assert value.label == "Person" + assert value.properties["name"] == "Alice" + + +def test_edge_roundtrip(conn: object) -> None: + value = _one(conn, CYPHER_EDGE) + assert isinstance(value, Edge) + assert value.label == "KNOWS" + assert value.start_id != value.end_id + + +def test_path_roundtrip(conn: object) -> None: + value = _one(conn, CYPHER_PATH) + assert isinstance(value, Path) + assert len(value.elements) == 3 + assert value.start.properties["name"] == "Alice"