diff --git a/backend/packages/harness/deerflow/persistence/base.py b/backend/packages/harness/deerflow/persistence/base.py index fd99d5f74..1c88c0c73 100644 --- a/backend/packages/harness/deerflow/persistence/base.py +++ b/backend/packages/harness/deerflow/persistence/base.py @@ -9,10 +9,23 @@ LangGraph's checkpointer tables are NOT managed by this Base. from __future__ import annotations +from functools import cache + from sqlalchemy import inspect as sa_inspect from sqlalchemy.orm import DeclarativeBase +@cache +def _column_keys(cls: type) -> tuple[str, ...]: + """Mapped column keys for an ORM class, in mapper order. + + ``to_dict``/``__repr__`` run per row (e.g. once per event when serializing a + messages page), so the SQLAlchemy mapper reflection is cached per class — + the mapping is fixed at class-definition time, so this never goes stale. + """ + return tuple(c.key for c in sa_inspect(cls).mapper.column_attrs) + + class Base(DeclarativeBase): """Base class for all DeerFlow ORM models. @@ -24,7 +37,7 @@ class Base(DeclarativeBase): def to_dict(self, *, exclude: set[str] | None = None) -> dict: """Convert ORM instance to plain dict. - Uses SQLAlchemy's inspect() to iterate mapped column attributes. + Uses cached mapped-column keys (see :func:`_column_keys`). Args: exclude: Optional set of column keys to omit. @@ -32,9 +45,11 @@ class Base(DeclarativeBase): Returns: Dict of {column_key: value} for all mapped columns. """ - exclude = exclude or set() - return {c.key: getattr(self, c.key) for c in sa_inspect(type(self)).mapper.column_attrs if c.key not in exclude} + keys = _column_keys(type(self)) + if exclude: + return {k: getattr(self, k) for k in keys if k not in exclude} + return {k: getattr(self, k) for k in keys} def __repr__(self) -> str: - cols = ", ".join(f"{c.key}={getattr(self, c.key)!r}" for c in sa_inspect(type(self)).mapper.column_attrs) + cols = ", ".join(f"{k}={getattr(self, k)!r}" for k in _column_keys(type(self))) return f"{type(self).__name__}({cols})" diff --git a/backend/tests/test_base_to_dict.py b/backend/tests/test_base_to_dict.py new file mode 100644 index 000000000..a01ca8b5a --- /dev/null +++ b/backend/tests/test_base_to_dict.py @@ -0,0 +1,46 @@ +"""Base.to_dict()/__repr__ caching + behavior. + +These run once per row when serializing ORM rows (e.g. every event in a +messages page), so the mapped-column reflection is cached per class. Behavior +must stay identical. +""" + +from __future__ import annotations + +from sqlalchemy import Integer, String +from sqlalchemy.orm import Mapped, mapped_column + +from deerflow.persistence.base import Base, _column_keys + + +class _Widget(Base): + __tablename__ = "_widget_to_dict_test" + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + name: Mapped[str] = mapped_column(String(32)) + color: Mapped[str] = mapped_column(String(16)) + + +def test_to_dict_returns_all_columns(): + w = _Widget(id=1, name="gear", color="red") + assert w.to_dict() == {"id": 1, "name": "gear", "color": "red"} + + +def test_to_dict_exclude(): + w = _Widget(id=2, name="cog", color="blue") + assert w.to_dict(exclude={"color"}) == {"id": 2, "name": "cog"} + # Empty exclude behaves like no exclude. + assert w.to_dict(exclude=set()) == {"id": 2, "name": "cog", "color": "blue"} + + +def test_column_keys_are_cached_per_class(): + # Same tuple object returned across calls -> reflection ran once. + assert _column_keys(_Widget) is _column_keys(_Widget) + assert _column_keys(_Widget) == ("id", "name", "color") + + +def test_repr_lists_columns(): + w = _Widget(id=3, name="bolt", color="green") + r = repr(w) + assert r.startswith("_Widget(") + assert "id=3" in r and "name='bolt'" in r and "color='green'" in r