Skip to content
Open
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
29 changes: 29 additions & 0 deletions app/core/db/crud.py
Original file line number Diff line number Diff line change
Expand Up @@ -1007,6 +1007,35 @@ async def count(
result = await session.scalar(query)
return result or 0

@classmethod
async def exists(
cls,
session: AsyncSession,
*whereclause: ColumnExpressionArgument[bool],
**equal_filters: Any,
) -> bool:
"""Return whether any record matches the query.

Emit a short-circuiting ``SELECT EXISTS (...)`` so the database can
stop at the first matching row. Filter arguments match :meth:`count`
and are applied through :meth:`_filter_query` (``None`` equal-filter
values are skipped — same behaviour as ``count``).

:param session: The SQLAlchemy asynchronous session to use for database
operations.
:type session: AsyncSession
:param whereclause: SQL expressions for the ``where`` clause of the query.
:type whereclause: ColumnExpressionArgument[bool]
:param equal_filters: Keyword arguments representing column names and their
respective filter values.
:type equal_filters: Any
:return: ``True`` when at least one matching row exists.
:rtype: bool
"""
inner = cls._filter_query(select(cls.Model), *whereclause, **equal_filters)
result = await session.scalar(select(inner.exists()))
return bool(result)


class BaseSQLModelManager(BaseManager):
"""Manage database operations for a BaseSQLModel-based model.
Expand Down
7 changes: 2 additions & 5 deletions app/core/settings_override/api/routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1088,11 +1088,8 @@ async def delete_setting(
settings_cls, proxy = _resolve(setting_class)
key = canonical_override_key(settings_cls, key)
field_meta = _field_meta_or_404(settings_cls, key)
has_override_row = (
await SettingsOverrideManager.count(
session, setting_class=setting_class, key=key
)
> 0
has_override_row = await SettingsOverrideManager.exists(
session, setting_class=setting_class, key=key
)
_assert_key_deletable(
settings_cls, field_meta, has_override_row=has_override_row
Expand Down
2 changes: 1 addition & 1 deletion app/sep/app_drain.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ async def finalize_drain_if_complete(session: AsyncSession, app_key: str) -> boo
!= AppLifecycleEnum.DISABLING
):
return False
if await AppRunningTaskManager.count(session, app_key=app_key):
if await AppRunningTaskManager.exists(session, app_key=app_key):
return False
result = await AppStateManager.update_where(
session,
Expand Down
28 changes: 2 additions & 26 deletions app/tasks/crud.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from collections.abc import AsyncGenerator, Mapping, Sequence
from datetime import datetime

from sqlalchemy import CursorResult, delete, func, literal, or_, update
from sqlalchemy import CursorResult, delete, func, or_, update
from sqlalchemy.orm import aliased
from sqlalchemy.sql.elements import ColumnElement
from sqlmodel import and_, col, select
Expand Down Expand Up @@ -802,30 +802,6 @@ async def delete_aged_batch(
result = await cls.delete_where(session, col(TaskHistoryLog.id).in_(doomed))
return result.rowcount

@classmethod
async def exists_for_task(cls, session: AsyncSession, task_history_id: int) -> bool:
"""Return ``True`` when at least one chunk exists for the task history.

Uses a ``SELECT 1 ... LIMIT 1`` short-circuit query so the database
can stop scanning as soon as it finds the first matching row instead
of counting every chunk.

:param session: The SQLAlchemy asynchronous session to use for query
execution.
:type session: AsyncSession
:param task_history_id: The ``TaskHistory`` identifier.
:type task_history_id: int
:return: Whether any chunk rows exist for the task history.
:rtype: bool
"""
query = (
select(literal(1))
.where(col(TaskHistoryLog.task_history_id) == task_history_id)
.limit(1)
)
result = await cls._exec(session, query)
return result.first() is not None

@classmethod
async def ids_with_chunks(
cls,
Expand All @@ -836,7 +812,7 @@ async def ids_with_chunks(

Emit a single ``SELECT DISTINCT task_history_id FROM taskhistory_log
WHERE task_history_id IN (:ids)`` so list endpoints avoid an N+1
:meth:`exists_for_task` call per paginated row. Return an empty set
:meth:`exists` call per paginated row. Return an empty set
for empty input without emitting any SQL -- an empty ``IN ()``
predicate triggers a SQLAlchemy warning and is a no-op anyway.

Expand Down
4 changes: 3 additions & 1 deletion app/tasks/logs/log_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -415,7 +415,9 @@ async def iter_task_history_logs(
offsets: defaultdict[str, dict[str | TaskLogType, int]] = defaultdict(
dict, start_offsets or {}
)
has_chunks = await TaskHistoryLogManager.exists_for_task(session, task_history.id)
has_chunks = await TaskHistoryLogManager.exists(
session, task_history_id=task_history.id
)
if tail_lines is not None and tail_lines > 0:
if has_chunks:
tail_offsets = await compute_tail_offsets_from_chunks(
Expand Down
10 changes: 6 additions & 4 deletions app/tasks/routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -360,7 +360,7 @@ async def _populate_has_logs(
"""Set ``has_logs`` on each history using chunk-store + legacy fallback.

Read the chunk store in one batched query so list endpoints avoid an
N+1 :meth:`TaskHistoryLogManager.exists_for_task` call per row, then
N+1 :meth:`TaskHistoryLogManager.exists` call per row, then
OR the result with :func:`has_legacy_logs` so legacy rows keep
rendering the **View Logs** button until the backfill lands.

Expand Down Expand Up @@ -465,7 +465,9 @@ async def retrieve_task_history(
logger.debug("Requesting task history %s", task_history.id)
_set_has_logs(
task_history,
value=await TaskHistoryLogManager.exists_for_task(session, task_history.id)
value=await TaskHistoryLogManager.exists(
session, task_history_id=task_history.id
)
or has_legacy_logs(task_history),
)
return task_history
Expand Down Expand Up @@ -605,7 +607,7 @@ async def stop_task_history(
stopped = await executor.stop_task(session, task_history)
_set_has_logs(
stopped,
value=await TaskHistoryLogManager.exists_for_task(session, stopped.id)
value=await TaskHistoryLogManager.exists(session, task_history_id=stopped.id)
or has_legacy_logs(stopped),
)
return stopped
Expand Down Expand Up @@ -692,7 +694,7 @@ async def sync_task_history(
await maybe_record_run(synced.id, executor)
_set_has_logs(
synced,
value=await TaskHistoryLogManager.exists_for_task(session, synced.id)
value=await TaskHistoryLogManager.exists(session, task_history_id=synced.id)
or has_legacy_logs(synced),
)
return synced
Expand Down
52 changes: 52 additions & 0 deletions tests/app/core/db/test_crud.py
Original file line number Diff line number Diff line change
Expand Up @@ -1021,3 +1021,55 @@ async def _seed_nullable_parent_items(session: AsyncSession) -> None:
)
await _create_item(session, name="b1", created_at=base_time + timedelta(minutes=2))
await _create_item(session, name="b2", created_at=base_time + timedelta(minutes=3))


class TestExists:
"""Cover ``BaseManager.exists`` short-circuit existence checks."""

@pytest.mark.asyncio
async def test_exists_false_on_empty_table(self, session: AsyncSession) -> None:
"""Return ``False`` when the table has no rows."""
assert await UniqueKeyManager.exists(session) is False

@pytest.mark.asyncio
async def test_exists_true_with_no_filters(self, session: AsyncSession) -> None:
"""Return ``True`` when any row exists and no filters are applied."""
await UniqueKeyManager.get_or_create(
session,
UniqueKeyModel(key="alpha", label="a"),
filter_include={"key"},
)
assert await UniqueKeyManager.exists(session) is True

@pytest.mark.asyncio
async def test_exists_with_equal_filters(self, session: AsyncSession) -> None:
"""Apply keyword equal-filters the same way ``count`` does."""
await UniqueKeyManager.get_or_create(
session,
UniqueKeyModel(key="alpha", label="a"),
filter_include={"key"},
)
assert await UniqueKeyManager.exists(session, key="alpha") is True
assert await UniqueKeyManager.exists(session, key="missing") is False

@pytest.mark.asyncio
async def test_exists_with_whereclause(self, session: AsyncSession) -> None:
"""Apply a positional ``whereclause`` expression."""
await UniqueKeyManager.get_or_create(
session,
UniqueKeyModel(key="alpha", label="keep"),
filter_include={"key"},
)
await UniqueKeyManager.get_or_create(
session,
UniqueKeyModel(key="beta", label="drop"),
filter_include={"key"},
)
assert (
await UniqueKeyManager.exists(session, col(UniqueKeyModel.label) == "keep")
is True
)
assert (
await UniqueKeyManager.exists(session, col(UniqueKeyModel.label) == "gone")
is False
)
12 changes: 6 additions & 6 deletions tests/app/tasks/connectivity/test_routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -354,8 +354,8 @@ async def sync_task_history(
assert data["success"] is True
assert data["error"] is None
assert call_count["n"] >= MIN_POLL_ITERATIONS
assert await TaskHistoryLogManager.exists_for_task(
session, data["task_history_id"]
assert await TaskHistoryLogManager.exists(
session, task_history_id=data["task_history_id"]
)

async def test_provisioning_latency_does_not_false_negative_over_http(
Expand Down Expand Up @@ -481,8 +481,8 @@ async def sync_task_history(
# Provisioning spanned more polls than the connect budget alone permits,
# yet the check still succeeded — the budgets are independent.
assert call_count["n"] > connect_budget // POLL_INTERVAL
assert await TaskHistoryLogManager.exists_for_task(
session, data["task_history_id"]
assert await TaskHistoryLogManager.exists(
session, task_history_id=data["task_history_id"]
)

async def test_timeout_surfaces_partial_logs_and_id_over_http(
Expand Down Expand Up @@ -599,6 +599,6 @@ async def sync_task_history(
assert "timed out" in data["error"]
assert "installing deps..." in data["error"]
assert data["task_history_id"] is not None
assert await TaskHistoryLogManager.exists_for_task(
session, data["task_history_id"]
assert await TaskHistoryLogManager.exists(
session, task_history_id=data["task_history_id"]
)
8 changes: 4 additions & 4 deletions tests/app/tasks/connectivity/test_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -323,8 +323,8 @@ async def sync_task_history(

assert result.success is True
assert result.error is None
assert await TaskHistoryLogManager.exists_for_task(
session, result.task_history_id
assert await TaskHistoryLogManager.exists(
session, task_history_id=result.task_history_id
)

async def test_terminal_run_records_nothing(
Expand Down Expand Up @@ -976,8 +976,8 @@ async def sync_task_history(
)
assert result.success is True
assert result.error is None
assert await TaskHistoryLogManager.exists_for_task(
session, result.task_history_id
assert await TaskHistoryLogManager.exists(
session, task_history_id=result.task_history_id
)

async def test_provisioning_phase_does_not_consume_connect_budget(
Expand Down
4 changes: 3 additions & 1 deletion tests/app/tasks/test_routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -2982,7 +2982,9 @@ async def fake_sync(
call_kwargs = mock_executor.sync_task_history.await_args.kwargs
assert "writer_session" in call_kwargs
assert call_kwargs["writer_session"] is not None
assert await TaskHistoryLogManager.exists_for_task(session, saved_history.id)
assert await TaskHistoryLogManager.exists(
session, task_history_id=saved_history.id
)

async def test_sync_hands_the_run_result_to_the_recorder(
self,
Expand Down
Loading