diff --git a/app/core/db/crud.py b/app/core/db/crud.py index 5032a9ac3a..611e0e0d9c 100644 --- a/app/core/db/crud.py +++ b/app/core/db/crud.py @@ -1007,6 +1007,31 @@ 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. + :param whereclause: SQL expressions for the ``where`` clause of the query. + :param equal_filters: Keyword arguments representing column names and their + respective filter values. + :return: ``True`` when at least one matching row exists. + """ + 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. diff --git a/app/core/settings_override/api/routes.py b/app/core/settings_override/api/routes.py index 8ec6880d98..33c5831431 100644 --- a/app/core/settings_override/api/routes.py +++ b/app/core/settings_override/api/routes.py @@ -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 diff --git a/app/sep/app_drain.py b/app/sep/app_drain.py index 8b5a271381..fb91dd4204 100644 --- a/app/sep/app_drain.py +++ b/app/sep/app_drain.py @@ -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, diff --git a/app/tasks/crud.py b/app/tasks/crud.py index 5b8893c73a..c93876cf65 100644 --- a/app/tasks/crud.py +++ b/app/tasks/crud.py @@ -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 @@ -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, @@ -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. diff --git a/app/tasks/logs/log_reader.py b/app/tasks/logs/log_reader.py index fe93d8cf95..1f97606a73 100644 --- a/app/tasks/logs/log_reader.py +++ b/app/tasks/logs/log_reader.py @@ -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( diff --git a/app/tasks/routes.py b/app/tasks/routes.py index b5d7544394..59abb75a32 100644 --- a/app/tasks/routes.py +++ b/app/tasks/routes.py @@ -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. @@ -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 @@ -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 @@ -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 diff --git a/tests/app/core/db/test_crud.py b/tests/app/core/db/test_crud.py index d2ca8f9278..29eb22d56e 100644 --- a/tests/app/core/db/test_crud.py +++ b/tests/app/core/db/test_crud.py @@ -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 + ) diff --git a/tests/app/tasks/connectivity/test_routes.py b/tests/app/tasks/connectivity/test_routes.py index 099af640cb..3ca148725f 100644 --- a/tests/app/tasks/connectivity/test_routes.py +++ b/tests/app/tasks/connectivity/test_routes.py @@ -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( @@ -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( @@ -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"] ) diff --git a/tests/app/tasks/connectivity/test_service.py b/tests/app/tasks/connectivity/test_service.py index be5d8ee915..cad7743fa6 100644 --- a/tests/app/tasks/connectivity/test_service.py +++ b/tests/app/tasks/connectivity/test_service.py @@ -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( @@ -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( diff --git a/tests/app/tasks/test_routes.py b/tests/app/tasks/test_routes.py index 06814a84e3..dd65accec4 100644 --- a/tests/app/tasks/test_routes.py +++ b/tests/app/tasks/test_routes.py @@ -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,