mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-24 15:07:49 +00:00
perf(persistence): cache Base.to_dict column reflection per class (#3654)
Base.to_dict() and __repr__() ran sqlalchemy.inspect(type(self)).mapper reflection on every call, but to_dict() is invoked once per row when serializing ORM results (e.g. every event in a messages/events page). The mapped columns are fixed at class-definition time, so cache the column keys per class with functools.cache and iterate the cached tuple. Behavior is unchanged. Fixes #3653.
This commit is contained in:
parent
3e055d8f43
commit
8cde305fe4
@ -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})"
|
||||
|
||||
46
backend/tests/test_base_to_dict.py
Normal file
46
backend/tests/test_base_to_dict.py
Normal file
@ -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
|
||||
Loading…
x
Reference in New Issue
Block a user