Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions src/backend/base/langflow/services/telemetry_writer/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -752,13 +752,13 @@ async def _flush(self, tx_batch: list[dict], vb_batch: list[dict]) -> None:
return
async with self._session_maker() as session:
if tx_batch:
await session.execute(TransactionTable.__table__.insert(), tx_batch)
await session.exec(TransactionTable.__table__.insert(), params=tx_batch)
for row in tx_batch:
flow_id = row.get("flow_id")
if flow_id is not None:
self._dirty_tx_flows.add(str(flow_id))
if vb_batch:
await session.execute(VertexBuildTable.__table__.insert(), vb_batch)
await session.exec(VertexBuildTable.__table__.insert(), params=vb_batch)
for row in vb_batch:
flow_id = row.get("flow_id")
if flow_id is not None:
Expand Down Expand Up @@ -817,7 +817,7 @@ def _as_uuid(value: str) -> UUID:
.order_by(col(TransactionTable.timestamp).desc())
.limit(max_transactions)
)
await session.execute(
await session.exec(
delete(TransactionTable).where(
TransactionTable.flow_id == flow_id,
col(TransactionTable.id).not_in(keep_subq),
Expand All @@ -827,7 +827,7 @@ def _as_uuid(value: str) -> UUID:
for flow_id in vb_flows:
vertex_ids = (
(
await session.execute(
await session.exec(
select(VertexBuildTable.id).where(VertexBuildTable.flow_id == flow_id).distinct()
)
)
Expand All @@ -847,7 +847,7 @@ def _as_uuid(value: str) -> UUID:
)
.limit(max_per_vertex)
)
await session.execute(
await session.exec(
delete(VertexBuildTable).where(
VertexBuildTable.flow_id == flow_id,
VertexBuildTable.id == vertex_id,
Expand All @@ -863,7 +863,7 @@ def _as_uuid(value: str) -> UUID:
)
.limit(max_vertex_builds)
)
await session.execute(
await session.exec(
delete(VertexBuildTable).where(col(VertexBuildTable.build_id).not_in(keep_global_subq))
)
await session.commit()
Expand Down
25 changes: 21 additions & 4 deletions src/backend/tests/unit/services/telemetry_writer/test_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import asyncio
import shutil
import tempfile
import warnings
from pathlib import Path
from uuid import uuid4

Expand All @@ -27,6 +28,7 @@
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import StaticPool
from sqlmodel import SQLModel
from sqlmodel.ext.asyncio.session import AsyncSession as SQLModelAsyncSession


class _FakeSettings:
Expand Down Expand Up @@ -104,7 +106,7 @@ async def writer_with_engine():
await conn.run_sync(SQLModel.metadata.create_all)
writer = _build_writer()
writer._engine = engine
writer._session_maker = async_sessionmaker(engine, expire_on_commit=False)
writer._session_maker = async_sessionmaker(engine, class_=SQLModelAsyncSession, expire_on_commit=False)
writer._started = True
writer._shutdown_event = asyncio.Event()
try:
Expand Down Expand Up @@ -342,6 +344,21 @@ async def test_flush_inserts_transactions_and_vertex_builds(writer_with_engine)
assert str(flow_id) in writer._dirty_vb_flows


async def test_flush_and_retention_do_not_call_deprecated_session_execute(writer_with_engine) -> None:
writer, _ = writer_with_engine
writer.settings_service.settings.max_transactions_to_keep = 1
writer.settings_service.settings.max_vertex_builds_to_keep = 1
flow_id = uuid4()

tx_batch = [_make_transaction_row(flow_id) for _ in range(3)]
vb_batch = [_make_vertex_build_row(flow_id, vertex_id="v1") for _ in range(3)]

with warnings.catch_warnings():
warnings.filterwarnings("error", message="(?s).*session\\.exec\\(\\).*", category=DeprecationWarning)
await writer._flush(tx_batch, vb_batch)
await writer._run_retention_pass()


async def test_retention_sweep_caps_transactions_per_flow(writer_with_engine) -> None:
writer, engine = writer_with_engine
writer.settings_service.settings.max_transactions_to_keep = 3
Expand Down Expand Up @@ -682,7 +699,7 @@ async def test_retention_failure_preserves_dirty_flows(writer_with_engine) -> No
writer._dirty_tx_flows.add(str(tx_flow))
writer._dirty_vb_flows.add(str(vb_flow))

# Wrap the real session_maker so commit raises but execute() works
# Wrap the real session_maker so commit raises but exec() works
# normally — this drives the retention pass through every query and only
# fails at the commit boundary, which is exactly the scenario the snapshot/
# restore logic guards against.
Expand Down Expand Up @@ -729,7 +746,7 @@ async def test_in_flight_batch_returned_on_cancel(writer_with_engine) -> None:
writer.enqueue_transaction(_make_transaction_row(flow_id))
assert len(writer._tx_buffer) == 20

# Replace the session_maker with one that hangs forever inside execute,
# Replace the session_maker with one that hangs forever inside exec,
# so the writer is guaranteed to be blocked mid-flush when we cancel.
hang = asyncio.Event() # never set

Expand All @@ -743,7 +760,7 @@ async def __aenter__(self):
async def __aexit__(self, *_):
return False

async def execute(self, *_):
async def exec(self, *_, **__):
await hang.wait()

async def commit(self):
Expand Down
Loading