First product
This commit is contained in:
parent
778da9e2fc
commit
4433d7d2c4
157 changed files with 23975 additions and 0 deletions
464
app/services/runs.py
Normal file
464
app/services/runs.py
Normal file
|
|
@ -0,0 +1,464 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import and_, case, func, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.db.models import (
|
||||
CrawlRun,
|
||||
IngestionQueueItem,
|
||||
RunStatus,
|
||||
RunTriggerType,
|
||||
ScholarProfile,
|
||||
)
|
||||
from app.services import continuation_queue as queue_service
|
||||
|
||||
QUEUE_STATUS_QUEUED = "queued"
|
||||
QUEUE_STATUS_RETRYING = "retrying"
|
||||
QUEUE_STATUS_DROPPED = "dropped"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueueListItem:
|
||||
id: int
|
||||
scholar_profile_id: int
|
||||
scholar_label: str
|
||||
status: str
|
||||
reason: str
|
||||
dropped_reason: str | None
|
||||
attempt_count: int
|
||||
resume_cstart: int
|
||||
next_attempt_dt: datetime | None
|
||||
updated_at: datetime
|
||||
last_error: str | None
|
||||
last_run_id: int | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueueClearResult:
|
||||
queue_item_id: int
|
||||
previous_status: str
|
||||
|
||||
|
||||
class QueueTransitionError(RuntimeError):
|
||||
def __init__(self, *, code: str, message: str, current_status: str) -> None:
|
||||
super().__init__(message)
|
||||
self.code = code
|
||||
self.message = message
|
||||
self.current_status = current_status
|
||||
|
||||
|
||||
def _safe_int(value: object, default: int = 0) -> int:
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _summary_dict(error_log: object) -> dict[str, Any]:
|
||||
if not isinstance(error_log, dict):
|
||||
return {}
|
||||
summary = error_log.get("summary")
|
||||
if not isinstance(summary, dict):
|
||||
return {}
|
||||
return summary
|
||||
|
||||
|
||||
def extract_run_summary(error_log: object) -> dict[str, Any]:
|
||||
summary = _summary_dict(error_log)
|
||||
return {
|
||||
"succeeded_count": _safe_int(summary.get("succeeded_count", 0)),
|
||||
"failed_count": _safe_int(summary.get("failed_count", 0)),
|
||||
"partial_count": _safe_int(summary.get("partial_count", 0)),
|
||||
"failed_state_counts": {
|
||||
str(key): _safe_int(value, 0)
|
||||
for key, value in summary.get("failed_state_counts", {}).items()
|
||||
if isinstance(key, str)
|
||||
}
|
||||
if isinstance(summary.get("failed_state_counts"), dict)
|
||||
else {},
|
||||
"failed_reason_counts": {
|
||||
str(key): _safe_int(value, 0)
|
||||
for key, value in summary.get("failed_reason_counts", {}).items()
|
||||
if isinstance(key, str)
|
||||
}
|
||||
if isinstance(summary.get("failed_reason_counts"), dict)
|
||||
else {},
|
||||
}
|
||||
|
||||
|
||||
async def list_recent_runs_for_user(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
limit: int = 20,
|
||||
) -> list[CrawlRun]:
|
||||
result = await db_session.execute(
|
||||
select(CrawlRun)
|
||||
.where(CrawlRun.user_id == user_id)
|
||||
.order_by(CrawlRun.start_dt.desc(), CrawlRun.id.desc())
|
||||
.limit(limit)
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
|
||||
|
||||
async def list_runs_for_user(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
limit: int = 100,
|
||||
failed_only: bool = False,
|
||||
) -> list[CrawlRun]:
|
||||
stmt = (
|
||||
select(CrawlRun)
|
||||
.where(CrawlRun.user_id == user_id)
|
||||
.order_by(CrawlRun.start_dt.desc(), CrawlRun.id.desc())
|
||||
.limit(limit)
|
||||
)
|
||||
if failed_only:
|
||||
stmt = stmt.where(
|
||||
CrawlRun.status.in_([RunStatus.FAILED, RunStatus.PARTIAL_FAILURE])
|
||||
)
|
||||
result = await db_session.execute(stmt)
|
||||
return list(result.scalars().all())
|
||||
|
||||
|
||||
async def get_run_for_user(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
run_id: int,
|
||||
) -> CrawlRun | None:
|
||||
result = await db_session.execute(
|
||||
select(CrawlRun).where(
|
||||
CrawlRun.user_id == user_id,
|
||||
CrawlRun.id == run_id,
|
||||
)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def get_manual_run_by_idempotency_key(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
idempotency_key: str,
|
||||
) -> CrawlRun | None:
|
||||
result = await db_session.execute(
|
||||
select(CrawlRun)
|
||||
.where(
|
||||
CrawlRun.user_id == user_id,
|
||||
CrawlRun.trigger_type == RunTriggerType.MANUAL,
|
||||
CrawlRun.error_log["meta"]["idempotency_key"].astext == idempotency_key,
|
||||
)
|
||||
.order_by(CrawlRun.start_dt.desc(), CrawlRun.id.desc())
|
||||
.limit(1)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def set_manual_run_idempotency_key(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
run_id: int,
|
||||
idempotency_key: str,
|
||||
) -> bool:
|
||||
run = await get_run_for_user(
|
||||
db_session,
|
||||
user_id=user_id,
|
||||
run_id=run_id,
|
||||
)
|
||||
if run is None:
|
||||
return False
|
||||
if run.trigger_type != RunTriggerType.MANUAL:
|
||||
return False
|
||||
|
||||
payload = dict(run.error_log) if isinstance(run.error_log, dict) else {}
|
||||
meta = payload.get("meta")
|
||||
meta_dict = dict(meta) if isinstance(meta, dict) else {}
|
||||
meta_dict["idempotency_key"] = idempotency_key
|
||||
payload["meta"] = meta_dict
|
||||
run.error_log = payload
|
||||
await db_session.commit()
|
||||
return True
|
||||
|
||||
|
||||
async def list_queue_items_for_user(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
limit: int = 200,
|
||||
) -> list[QueueListItem]:
|
||||
result = await db_session.execute(
|
||||
select(
|
||||
IngestionQueueItem.id,
|
||||
IngestionQueueItem.scholar_profile_id,
|
||||
ScholarProfile.display_name,
|
||||
ScholarProfile.scholar_id,
|
||||
IngestionQueueItem.status,
|
||||
IngestionQueueItem.reason,
|
||||
IngestionQueueItem.dropped_reason,
|
||||
IngestionQueueItem.attempt_count,
|
||||
IngestionQueueItem.resume_cstart,
|
||||
IngestionQueueItem.next_attempt_dt,
|
||||
IngestionQueueItem.updated_at,
|
||||
IngestionQueueItem.last_error,
|
||||
IngestionQueueItem.last_run_id,
|
||||
)
|
||||
.join(
|
||||
ScholarProfile,
|
||||
and_(
|
||||
ScholarProfile.id == IngestionQueueItem.scholar_profile_id,
|
||||
ScholarProfile.user_id == IngestionQueueItem.user_id,
|
||||
),
|
||||
)
|
||||
.where(IngestionQueueItem.user_id == user_id)
|
||||
.order_by(
|
||||
case((IngestionQueueItem.status == "dropped", 1), else_=0).asc(),
|
||||
IngestionQueueItem.next_attempt_dt.asc(),
|
||||
IngestionQueueItem.id.asc(),
|
||||
)
|
||||
.limit(limit)
|
||||
)
|
||||
|
||||
items: list[QueueListItem] = []
|
||||
for row in result.all():
|
||||
(
|
||||
item_id,
|
||||
scholar_profile_id,
|
||||
display_name,
|
||||
scholar_id,
|
||||
status,
|
||||
reason,
|
||||
dropped_reason,
|
||||
attempt_count,
|
||||
resume_cstart,
|
||||
next_attempt_dt,
|
||||
updated_at,
|
||||
last_error,
|
||||
last_run_id,
|
||||
) = row
|
||||
items.append(
|
||||
QueueListItem(
|
||||
id=int(item_id),
|
||||
scholar_profile_id=int(scholar_profile_id),
|
||||
scholar_label=(display_name or scholar_id),
|
||||
status=str(status),
|
||||
reason=str(reason),
|
||||
dropped_reason=dropped_reason,
|
||||
attempt_count=int(attempt_count or 0),
|
||||
resume_cstart=int(resume_cstart or 0),
|
||||
next_attempt_dt=next_attempt_dt,
|
||||
updated_at=updated_at,
|
||||
last_error=last_error,
|
||||
last_run_id=int(last_run_id) if last_run_id is not None else None,
|
||||
)
|
||||
)
|
||||
return items
|
||||
|
||||
|
||||
async def get_queue_item_for_user(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
queue_item_id: int,
|
||||
) -> QueueListItem | None:
|
||||
result = await db_session.execute(
|
||||
select(
|
||||
IngestionQueueItem.id,
|
||||
IngestionQueueItem.scholar_profile_id,
|
||||
ScholarProfile.display_name,
|
||||
ScholarProfile.scholar_id,
|
||||
IngestionQueueItem.status,
|
||||
IngestionQueueItem.reason,
|
||||
IngestionQueueItem.dropped_reason,
|
||||
IngestionQueueItem.attempt_count,
|
||||
IngestionQueueItem.resume_cstart,
|
||||
IngestionQueueItem.next_attempt_dt,
|
||||
IngestionQueueItem.updated_at,
|
||||
IngestionQueueItem.last_error,
|
||||
IngestionQueueItem.last_run_id,
|
||||
)
|
||||
.join(
|
||||
ScholarProfile,
|
||||
and_(
|
||||
ScholarProfile.id == IngestionQueueItem.scholar_profile_id,
|
||||
ScholarProfile.user_id == IngestionQueueItem.user_id,
|
||||
),
|
||||
)
|
||||
.where(
|
||||
IngestionQueueItem.user_id == user_id,
|
||||
IngestionQueueItem.id == queue_item_id,
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
row = result.one_or_none()
|
||||
if row is None:
|
||||
return None
|
||||
(
|
||||
item_id,
|
||||
scholar_profile_id,
|
||||
display_name,
|
||||
scholar_id,
|
||||
status,
|
||||
reason,
|
||||
dropped_reason,
|
||||
attempt_count,
|
||||
resume_cstart,
|
||||
next_attempt_dt,
|
||||
updated_at,
|
||||
last_error,
|
||||
last_run_id,
|
||||
) = row
|
||||
return QueueListItem(
|
||||
id=int(item_id),
|
||||
scholar_profile_id=int(scholar_profile_id),
|
||||
scholar_label=(display_name or scholar_id),
|
||||
status=str(status),
|
||||
reason=str(reason),
|
||||
dropped_reason=dropped_reason,
|
||||
attempt_count=int(attempt_count or 0),
|
||||
resume_cstart=int(resume_cstart or 0),
|
||||
next_attempt_dt=next_attempt_dt,
|
||||
updated_at=updated_at,
|
||||
last_error=last_error,
|
||||
last_run_id=int(last_run_id) if last_run_id is not None else None,
|
||||
)
|
||||
|
||||
|
||||
async def retry_queue_item_for_user(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
queue_item_id: int,
|
||||
) -> QueueListItem | None:
|
||||
result = await db_session.execute(
|
||||
select(IngestionQueueItem).where(
|
||||
IngestionQueueItem.id == queue_item_id,
|
||||
IngestionQueueItem.user_id == user_id,
|
||||
)
|
||||
)
|
||||
item = result.scalar_one_or_none()
|
||||
if item is None:
|
||||
return None
|
||||
if item.status == QUEUE_STATUS_QUEUED:
|
||||
raise QueueTransitionError(
|
||||
code="queue_item_already_queued",
|
||||
message="Queue item is already queued.",
|
||||
current_status=item.status,
|
||||
)
|
||||
if item.status == QUEUE_STATUS_RETRYING:
|
||||
raise QueueTransitionError(
|
||||
code="queue_item_retrying",
|
||||
message="Queue item is currently retrying.",
|
||||
current_status=item.status,
|
||||
)
|
||||
|
||||
await queue_service.mark_queued_now(
|
||||
db_session,
|
||||
job_id=item.id,
|
||||
reason="manual_retry",
|
||||
reset_attempt_count=(item.status == QUEUE_STATUS_DROPPED),
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
return await get_queue_item_for_user(
|
||||
db_session,
|
||||
user_id=user_id,
|
||||
queue_item_id=queue_item_id,
|
||||
)
|
||||
|
||||
|
||||
async def drop_queue_item_for_user(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
queue_item_id: int,
|
||||
) -> QueueListItem | None:
|
||||
result = await db_session.execute(
|
||||
select(IngestionQueueItem).where(
|
||||
IngestionQueueItem.id == queue_item_id,
|
||||
IngestionQueueItem.user_id == user_id,
|
||||
)
|
||||
)
|
||||
item = result.scalar_one_or_none()
|
||||
if item is None:
|
||||
return None
|
||||
if item.status == QUEUE_STATUS_DROPPED:
|
||||
raise QueueTransitionError(
|
||||
code="queue_item_already_dropped",
|
||||
message="Queue item is already dropped.",
|
||||
current_status=item.status,
|
||||
)
|
||||
|
||||
await queue_service.mark_dropped(
|
||||
db_session,
|
||||
job_id=int(item.id),
|
||||
reason="manual_drop",
|
||||
)
|
||||
await db_session.commit()
|
||||
return await get_queue_item_for_user(
|
||||
db_session,
|
||||
user_id=user_id,
|
||||
queue_item_id=queue_item_id,
|
||||
)
|
||||
|
||||
|
||||
async def clear_queue_item_for_user(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
queue_item_id: int,
|
||||
) -> QueueClearResult | None:
|
||||
result = await db_session.execute(
|
||||
select(IngestionQueueItem).where(
|
||||
IngestionQueueItem.id == queue_item_id,
|
||||
IngestionQueueItem.user_id == user_id,
|
||||
)
|
||||
)
|
||||
item = result.scalar_one_or_none()
|
||||
if item is None:
|
||||
return None
|
||||
if item.status != QUEUE_STATUS_DROPPED:
|
||||
raise QueueTransitionError(
|
||||
code="queue_item_not_dropped",
|
||||
message="Queue item can only be cleared after it is dropped.",
|
||||
current_status=item.status,
|
||||
)
|
||||
|
||||
item_id = int(item.id)
|
||||
previous_status = str(item.status)
|
||||
deleted = await queue_service.delete_job_by_id(
|
||||
db_session,
|
||||
job_id=item_id,
|
||||
)
|
||||
await db_session.commit()
|
||||
if not deleted:
|
||||
return None
|
||||
return QueueClearResult(
|
||||
queue_item_id=item_id,
|
||||
previous_status=previous_status,
|
||||
)
|
||||
|
||||
|
||||
async def queue_status_counts_for_user(
|
||||
db_session: AsyncSession,
|
||||
*,
|
||||
user_id: int,
|
||||
) -> dict[str, int]:
|
||||
result = await db_session.execute(
|
||||
select(
|
||||
IngestionQueueItem.status,
|
||||
func.count(IngestionQueueItem.id),
|
||||
)
|
||||
.where(IngestionQueueItem.user_id == user_id)
|
||||
.group_by(IngestionQueueItem.status)
|
||||
)
|
||||
counts: dict[str, int] = {"queued": 0, "retrying": 0, "dropped": 0}
|
||||
for status, count in result.all():
|
||||
counts[str(status)] = int(count or 0)
|
||||
return counts
|
||||
Loading…
Add table
Add a link
Reference in a new issue