Merge pull request #45 from JustinZeus/bugfix/broken-import

openalex cooldown

.claude/settings.json

@@ -23,7 +23,8 @@
       "Bash(gh run list --branch openalex-size-reduction-scholar-hydration --limit 3 --json databaseId,status,conclusion,name,event)",
       "Bash(gh run list --limit 10 --json databaseId,status,conclusion,name,event,headBranch)",
       "Bash(gh run view 22554585841 --log-failed)",
-      "Bash(gh run rerun 22554585841 --failed)"
+      "Bash(gh run rerun 22554585841 --failed)",
+      "Bash(bash scripts/premerge.sh)"
     ],
     "additionalDirectories": [
       "/home/jv/src/personal/playground/scholar_scraper/scholar-scraper/frontend/src"

.env.example

@@ -28,6 +28,7 @@ DATABASE_POOL_MODE=auto
 DATABASE_POOL_SIZE=5
 DATABASE_POOL_MAX_OVERFLOW=10
 DATABASE_POOL_TIMEOUT_SECONDS=30
+DATABASE_RESERVED_API_CONNECTIONS=3
 
 # ------------------------------
 # Frontend Dev Overrides

app/api/routers/runs.py

@@ -221,11 +221,11 @@ def _spawn_background_execution(
     user_settings: Any,
     idempotency_key: str | None,
 ) -> None:
-    from app.db.session import get_session_factory
+    from app.db.background_session import background_session
 
     task = asyncio.create_task(
         ingest_service.execute_run(
-            session_factory=get_session_factory(),
+            session_factory=background_session,
             run_id=run.id,
             user_id=current_user.id,
             scholars=scholars,

app/db/background_session.py

@@ -0,0 +1,57 @@
+"""Semaphore-gated DB sessions for background tasks.
+
+Background tasks (ingestion, PDF resolution, scheduler) acquire a semaphore
+permit before checking out a database connection.  The semaphore cap is
+``pool_size + max_overflow - reserved_for_api``, which guarantees that at
+least *reserved_for_api* connections remain available for API request
+handlers at all times.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.db.session import get_session_factory
+from app.logging_utils import structured_log
+from app.settings import settings
+
+logger = logging.getLogger(__name__)
+
+_semaphore: asyncio.Semaphore | None = None
+
+
+def _build_semaphore() -> asyncio.Semaphore:
+    pool_capacity = max(1, settings.database_pool_size) + max(0, settings.database_pool_max_overflow)
+    reserved = max(0, settings.database_reserved_api_connections)
+    limit = max(1, pool_capacity - reserved)
+    structured_log(
+        logger,
+        "info",
+        "db.background_semaphore_initialized",
+        pool_capacity=pool_capacity,
+        reserved_for_api=reserved,
+        background_limit=limit,
+    )
+    return asyncio.Semaphore(limit)
+
+
+def get_background_semaphore() -> asyncio.Semaphore:
+    global _semaphore
+    if _semaphore is None:
+        _semaphore = _build_semaphore()
+    return _semaphore
+
+
+@asynccontextmanager
+async def background_session() -> AsyncIterator[AsyncSession]:
+    """Yield a DB session after acquiring the background semaphore."""
+    semaphore = get_background_semaphore()
+    async with semaphore:
+        session_factory = get_session_factory()
+        async with session_factory() as session:
+            yield session

app/services/ingestion/queue_runner.py

@@ -5,8 +5,8 @@ from datetime import UTC, datetime
 
 from sqlalchemy import select
 
+from app.db.background_session import background_session
 from app.db.models import QueueItemStatus, RunTriggerType, ScholarProfile, UserSetting
-from app.db.session import get_session_factory
 from app.logging_utils import structured_log
 from app.services.ingestion import queue as queue_service
 from app.services.ingestion.application import (
@@ -57,8 +57,7 @@ class QueueJobRunner:
 
     async def drain_continuation_queue(self) -> None:
         now = datetime.now(UTC)
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        async with background_session() as session:
             jobs = await queue_service.list_due_jobs(
                 session,
                 now=now,
@@ -73,8 +72,7 @@ class QueueJobRunner:
     ) -> bool:
         if job.attempt_count < self._continuation_max_attempts:
             return False
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        async with background_session() as session:
             dropped = await queue_service.mark_dropped(
                 session,
                 job_id=job.id,
@@ -98,8 +96,7 @@ class QueueJobRunner:
         self,
         job: queue_service.ContinuationQueueJob,
     ) -> bool:
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        async with background_session() as session:
             queue_item = await queue_service.mark_retrying(session, job_id=job.id)
             await session.commit()
         if queue_item is None:
@@ -110,8 +107,7 @@ class QueueJobRunner:
         self,
         job: queue_service.ContinuationQueueJob,
     ) -> bool:
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        async with background_session() as session:
             scholar_result = await session.execute(
                 select(ScholarProfile.id).where(
                     ScholarProfile.user_id == job.user_id,
@@ -140,8 +136,7 @@ class QueueJobRunner:
         return False
 
     async def _reschedule_queue_job_lock_active(self, job: queue_service.ContinuationQueueJob) -> None:
-        session_factory = get_session_factory()
-        async with session_factory() as recovery_session:
+        async with background_session() as recovery_session:
             await queue_service.reschedule_job(
                 recovery_session,
                 job_id=job.id,
@@ -167,8 +162,7 @@ class QueueJobRunner:
             self._tick_seconds,
             int(exc.safety_state.get("cooldown_remaining_seconds") or 0),
         )
-        session_factory = get_session_factory()
-        async with session_factory() as recovery_session:
+        async with background_session() as recovery_session:
             await queue_service.reschedule_job(
                 recovery_session,
                 job_id=job.id,
@@ -193,8 +187,7 @@ class QueueJobRunner:
         *,
         exc: Exception,
     ) -> None:
-        session_factory = get_session_factory()
-        async with session_factory() as recovery_session:
+        async with background_session() as recovery_session:
             queue_item = await queue_service.increment_attempt_count(recovery_session, job_id=job.id)
             if queue_item is None:
                 await recovery_session.commit()
@@ -238,8 +231,7 @@ class QueueJobRunner:
         )
 
     async def _load_request_delay_for_user(self, user_id: int, *, floor: int) -> int:
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        async with background_session() as session:
             result = await session.execute(
                 select(UserSetting.request_delay_seconds).where(UserSetting.user_id == user_id)
             )
@@ -253,8 +245,7 @@ class QueueJobRunner:
         request_delay_floor: int,
     ):
         request_delay_seconds = await self._load_request_delay_for_user(job.user_id, floor=request_delay_floor)
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        async with background_session() as session:
             ingestion = ScholarIngestionService(source=self._source)
             try:
                 return await ingestion.run_for_user(
@@ -366,8 +357,7 @@ class QueueJobRunner:
         )
 
     async def _finalize_queue_job_after_run(self, job: queue_service.ContinuationQueueJob, run_summary) -> None:
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        async with background_session() as session:
             if int(run_summary.failed_count) <= 0:
                 await self._finalize_successful_queue_job(session, job, run_summary)
             else:

app/services/ingestion/scheduler.py

@@ -8,13 +8,13 @@ from typing import Any
 
 from sqlalchemy import select
 
+from app.db.background_session import background_session
 from app.db.models import (
     CrawlRun,
     RunTriggerType,
     User,
     UserSetting,
 )
-from app.db.session import get_session_factory
 from app.logging_utils import structured_log
 from app.services.ingestion.application import (
     RunAlreadyInProgressError,
@@ -148,8 +148,7 @@ class SchedulerService:
             await self._run_candidate(candidate)
 
     async def _load_candidate_rows(self) -> list[Any]:
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        async with background_session() as session:
             result = await session.execute(
                 select(
                     UserSetting.user_id,
@@ -204,8 +203,7 @@ class SchedulerService:
         return candidates
 
     async def _is_due(self, candidate: _AutoRunCandidate, *, now: datetime) -> bool:
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        async with background_session() as session:
             result = await session.execute(
                 select(CrawlRun.start_dt)
                 .where(
@@ -227,8 +225,7 @@ class SchedulerService:
         *,
         candidate: _AutoRunCandidate,
     ):
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        async with background_session() as session:
             ingestion = ScholarIngestionService(source=self._source)
             try:
                 return await ingestion.run_for_user(
@@ -289,9 +286,12 @@ class SchedulerService:
 
     async def _drain_pdf_queue(self) -> None:
         from app.services.publications.pdf_queue import drain_ready_jobs
+        from app.services.publications.pdf_queue_resolution import is_budget_cooldown_active
 
-        session_factory = get_session_factory()
-        async with session_factory() as session:
+        if is_budget_cooldown_active():
+            return
+
+        async with background_session() as session:
             try:
                 processed = await drain_ready_jobs(
                     session,

app/services/publications/pdf_queue_resolution.py

@@ -2,9 +2,10 @@ from __future__ import annotations
 
 import asyncio
 import logging
+from datetime import UTC, datetime, timedelta
 
+from app.db.background_session import background_session
 from app.db.models import Publication, PublicationPdfJob
-from app.db.session import get_session_factory
 from app.logging_utils import structured_log
 from app.services.publication_identifiers import application as identifier_service
 from app.services.publications.pdf_queue_common import (
@@ -29,8 +30,20 @@ PDF_EVENT_ATTEMPT_STARTED = "attempt_started"
 PDF_EVENT_RESOLVED = "resolved"
 PDF_EVENT_FAILED = "failed"
 
+_BUDGET_COOLDOWN_MINUTES = 15
+
 logger = logging.getLogger(__name__)
 _scheduled_tasks: set[asyncio.Task[None]] = set()
+_budget_cooldown_until: datetime | None = None
+
+
+def is_budget_cooldown_active() -> bool:
+    return _budget_cooldown_until is not None and datetime.now(UTC) < _budget_cooldown_until
+
+
+def _enter_budget_cooldown() -> None:
+    global _budget_cooldown_until
+    _budget_cooldown_until = datetime.now(UTC) + timedelta(minutes=_BUDGET_COOLDOWN_MINUTES)
 
 
 async def _mark_attempt_started(
@@ -38,8 +51,7 @@ async def _mark_attempt_started(
     publication_id: int,
     user_id: int,
 ) -> None:
-    session_factory = get_session_factory()
-    async with session_factory() as db_session:
+    async with background_session() as db_session:
         job = await db_session.get(PublicationPdfJob, publication_id)
         if job is None:
             job = queued_job(publication_id=publication_id, user_id=user_id)
@@ -125,8 +137,7 @@ async def _persist_outcome(
     user_id: int,
     outcome: OaResolutionOutcome,
 ) -> None:
-    session_factory = get_session_factory()
-    async with session_factory() as db_session:
+    async with background_session() as db_session:
         publication = await db_session.get(Publication, publication_id)
         job = await db_session.get(PublicationPdfJob, publication_id)
         if publication is None or job is None:
@@ -207,8 +218,7 @@ async def _run_resolution_task(
 
     openalex_api_key: str | None = None
     try:
-        session_factory = get_session_factory()
-        async with session_factory() as key_session:
+        async with background_session() as key_session:
             user_settings = await user_settings_service.get_or_create_settings(key_session, user_id=user_id)
             openalex_api_key = getattr(user_settings, "openalex_api_key", None) or settings.openalex_api_key
     except Exception:
@@ -233,11 +243,13 @@ async def _run_resolution_task(
                     detail="arXiv temporarily disabled for remaining batch after rate limit",
                 )
         except OpenAlexBudgetExhaustedError:
+            _enter_budget_cooldown()
             structured_log(
                 logger,
                 "warning",
                 "pdf_queue.budget_exhausted",
                 detail="Stopping PDF resolution batch — OpenAlex daily budget exhausted",
+                cooldown_minutes=_BUDGET_COOLDOWN_MINUTES,
             )
             break
         except Exception:

app/settings.py

@@ -275,6 +275,8 @@ class Settings:
     crossref_max_lookups_per_request: int = _env_int("CROSSREF_MAX_LOOKUPS_PER_REQUEST", 8)
 
     openalex_api_key: str | None = os.getenv("OPENALEX_API_KEY")
+    database_reserved_api_connections: int = _env_int("DATABASE_RESERVED_API_CONNECTIONS", 3)
+
     crossref_api_token: str | None = os.getenv("CROSSREF_API_TOKEN")
     crossref_api_mailto: str | None = os.getenv("CROSSREF_API_MAILTO")
 

frontend/src/features/publications/composables/usePublicationData.ts

@@ -277,12 +277,16 @@ export function usePublicationData() {
 
   // --- Run-triggered refresh watcher ---
 
+  let previousRunStatusKey: string | null = null;
+
   watch(
     () => runStatus.latestRun,
-    async (nextRun, previousRun) => {
-      const nextRunId = nextRun && (nextRun.status === "running" || nextRun.status === "resolving") ? nextRun.id : null;
-      const previousRunId = previousRun && (previousRun.status === "running" || previousRun.status === "resolving") ? previousRun.id : null;
-      if (nextRunId === null || nextRunId === previousRunId) return;
+    async (nextRun) => {
+      const nextStatus = nextRun ? `${nextRun.id}:${nextRun.status}` : null;
+      if (nextStatus === previousRunStatusKey) return;
+      previousRunStatusKey = nextStatus;
+      const isActive = nextRun && (nextRun.status === "running" || nextRun.status === "resolving");
+      if (!isActive) return;
       resetPageAndSnapshot();
       await loadPublications();
     },

frontend/src/pages/PublicationsPage.vue

@@ -1,5 +1,5 @@
 <script setup lang="ts">
-import { computed, onMounted, onUnmounted, ref, watch } from "vue";
+import { computed, onMounted, ref, watch } from "vue";
 
 import AppPage from "@/components/layout/AppPage.vue";
 import AsyncStateGate from "@/components/patterns/AsyncStateGate.vue";
@@ -37,23 +37,6 @@ const selectedPublicationKeys = ref<Set<string>>(new Set());
 const retryingPublicationKeys = ref<Set<string>>(new Set());
 const favoriteUpdatingKeys = ref<Set<string>>(new Set());
 
-const PUBLICATIONS_RUN_STATUS_SYNC_INTERVAL_MS = 5000;
-let runStatusSyncTimer: ReturnType<typeof setInterval> | null = null;
-
-function startRunStatusSyncLoop(): void {
-  if (runStatusSyncTimer !== null) return;
-  runStatusSyncTimer = setInterval(() => {
-    if (runStatus.isRunActive) return;
-    void runStatus.syncLatest();
-  }, PUBLICATIONS_RUN_STATUS_SYNC_INTERVAL_MS);
-}
-
-function stopRunStatusSyncLoop(): void {
-  if (runStatusSyncTimer === null) return;
-  clearInterval(runStatusSyncTimer);
-  runStatusSyncTimer = null;
-}
-
 // --- Computed ---
 
 const selectedCount = computed(() => selectedPublicationKeys.value.size);
@@ -358,13 +341,8 @@ async function onStartRun(): Promise<void> {
 
 onMounted(() => {
   pub.syncFiltersFromRoute();
-  startRunStatusSyncLoop();
   void Promise.all([pub.loadScholarFilters(), pub.loadPublications(), runStatus.syncLatest()]);
 });
-
-onUnmounted(() => {
-  stopRunStatusSyncLoop();
-});
 </script>
 
 <template>

frontend/src/stores/run_status.ts

@@ -9,7 +9,7 @@ import {
 import { type PublicationItem } from "@/features/publications";
 import { ApiRequestError } from "@/lib/api/errors";
 
-export const RUN_STATUS_POLL_INTERVAL_MS = 5000;
+export const RUN_STATUS_POLL_INTERVAL_MS = 15000;
 export const RUN_STATUS_STARTING_PHASE_MS = 1500;
 
 export type StartManualCheckResult =

tests/unit/test_publication_pdf_queue_policy.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 from datetime import UTC, datetime, timedelta
 from types import SimpleNamespace
 from typing import Any
@@ -238,8 +239,10 @@ async def test_run_resolution_task_disables_arxiv_for_remaining_batch(
     first = _row()
     second: Any = SimpleNamespace(**{**first.__dict__, "publication_id": 2})
 
-    def _raise_session_factory_error():
+    @contextlib.asynccontextmanager
+    async def _raise_background_session_error():
         raise RuntimeError("skip user settings lookup in test")
+        yield  # pragma: no cover
 
     async def _fake_resolve_publication_row(
         *,
@@ -253,7 +256,7 @@ async def test_run_resolution_task_disables_arxiv_for_remaining_batch(
         calls.append((int(row.publication_id), bool(allow_arxiv_lookup)))
         return row.publication_id == 1
 
-    monkeypatch.setattr(pdf_queue_resolution, "get_session_factory", _raise_session_factory_error)
+    monkeypatch.setattr(pdf_queue_resolution, "background_session", _raise_background_session_error)
     monkeypatch.setattr(pdf_queue_resolution, "_resolve_publication_row", _fake_resolve_publication_row)
 
     await pdf_queue_resolution._run_resolution_task(