110 lines
3.2 KiB
Python
110 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import Any
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.db.models import CrawlRun, RunStatus
|
|
from app.logging_utils import structured_log
|
|
from app.services.ingestion.enrichment import EnrichmentRunner
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_background_tasks: set[asyncio.Task[Any]] = set()
|
|
|
|
|
|
async def background_enrich(
|
|
session_factory: Any,
|
|
enrichment_runner: EnrichmentRunner,
|
|
*,
|
|
run_id: int,
|
|
intended_final_status: RunStatus,
|
|
openalex_api_key: str | None = None,
|
|
) -> None:
|
|
try:
|
|
async with session_factory() as session:
|
|
await enrichment_runner.enrich_pending_publications(
|
|
session,
|
|
run_id=run_id,
|
|
openalex_api_key=openalex_api_key,
|
|
)
|
|
run = await session.get(CrawlRun, run_id)
|
|
if run is not None and run.status == RunStatus.RESOLVING:
|
|
run.status = intended_final_status
|
|
await session.commit()
|
|
structured_log(
|
|
logger,
|
|
"info",
|
|
"ingestion.background_enrichment_complete",
|
|
run_id=run_id,
|
|
final_status=str(intended_final_status),
|
|
)
|
|
except Exception:
|
|
structured_log(
|
|
logger,
|
|
"exception",
|
|
"ingestion.background_enrichment_failed",
|
|
run_id=run_id,
|
|
)
|
|
try:
|
|
async with session_factory() as fallback_session:
|
|
run = await fallback_session.get(CrawlRun, run_id)
|
|
if run is not None and run.status == RunStatus.RESOLVING:
|
|
run.status = intended_final_status
|
|
await fallback_session.commit()
|
|
except Exception:
|
|
structured_log(
|
|
logger,
|
|
"exception",
|
|
"ingestion.background_enrichment_fallback_failed",
|
|
run_id=run_id,
|
|
)
|
|
|
|
|
|
async def inline_enrich_and_finalize(
|
|
db_session: AsyncSession,
|
|
enrichment_runner: EnrichmentRunner,
|
|
*,
|
|
run: CrawlRun,
|
|
user_settings: Any,
|
|
intended_final_status: RunStatus,
|
|
) -> None:
|
|
try:
|
|
await enrichment_runner.enrich_pending_publications(
|
|
db_session,
|
|
run_id=run.id,
|
|
openalex_api_key=getattr(user_settings, "openalex_api_key", None),
|
|
)
|
|
except Exception:
|
|
structured_log(
|
|
logger,
|
|
"exception",
|
|
"ingestion.enrichment_failed",
|
|
run_id=run.id,
|
|
)
|
|
if run.status == RunStatus.RESOLVING:
|
|
run.status = intended_final_status
|
|
await db_session.commit()
|
|
|
|
|
|
def spawn_background_enrichment_task(
|
|
session_factory: Any,
|
|
enrichment_runner: EnrichmentRunner,
|
|
*,
|
|
run_id: int,
|
|
intended_final_status: RunStatus,
|
|
openalex_api_key: str | None,
|
|
) -> None:
|
|
task = asyncio.create_task(
|
|
background_enrich(
|
|
session_factory,
|
|
enrichment_runner,
|
|
run_id=run_id,
|
|
intended_final_status=intended_final_status,
|
|
openalex_api_key=openalex_api_key,
|
|
)
|
|
)
|
|
_background_tasks.add(task)
|
|
task.add_done_callback(_background_tasks.discard)
|