saksham
1a4c80958f
FastAPI + Celery + Next.js + Postgres/Redis app with company monitoring, source collection, LLM-based change analysis, enrichment, and account security (Turnstile, escalating lockout, email verification).
226 lines
8.8 KiB
Python
226 lines
8.8 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from sqlalchemy import delete, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.db.base import ensure_aware_utc
|
|
from app.models.company import Company
|
|
from app.models.enums import SourceStatus, SourceType
|
|
from app.models.snapshot import Snapshot
|
|
from app.models.source import Source
|
|
from app.models.source_document import SourceDocument
|
|
|
|
|
|
def _is_due(value: datetime | None, now: datetime) -> bool:
|
|
# SQLite silently drops tzinfo on read-back (Postgres does not) - any
|
|
# value read from the DB must go through ensure_aware_utc before being
|
|
# compared against an aware `now` in Python, or this comparison would
|
|
# raise on SQLite while working fine on Postgres. See db/base.py.
|
|
return value is not None and ensure_aware_utc(value) <= now
|
|
|
|
|
|
def _is_source_due(source: Source, now: datetime, company_next_run: datetime | None) -> bool:
|
|
"""A source with its own frequency override uses its own next_check
|
|
(due immediately if never computed yet); one with no override rides
|
|
the company's own next_run clock instead."""
|
|
if source.frequency_type is not None:
|
|
return source.next_check is None or _is_due(source.next_check, now)
|
|
return _is_due(company_next_run, now)
|
|
|
|
|
|
class SourceRepository:
|
|
def __init__(self, db: AsyncSession) -> None:
|
|
self.db = db
|
|
|
|
async def list_for_company(self, company_id: uuid.UUID) -> list[Source]:
|
|
result = await self.db.execute(
|
|
select(Source).where(Source.company_id == company_id).order_by(Source.created_at)
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
async def list_due_for_company(
|
|
self, company_id: uuid.UUID, now: datetime, company_next_run: datetime | None
|
|
) -> list[Source]:
|
|
"""Active sources that are due for a check right now. A source with
|
|
its own frequency override (frequency_type is not None) uses its
|
|
own next_check, and is treated as due if it's never been computed
|
|
yet (a brand-new override should get its first check immediately,
|
|
same as any brand-new source always has). A source with no
|
|
override rides the company's own next_run clock instead - this is
|
|
what keeps "no override configured" behaviorally identical to how
|
|
every source worked before per-source scheduling existed. Small
|
|
per-company table (sources per company is always small for this
|
|
app), so filtering in Python after one plain SELECT is simpler than
|
|
expressing the OR-with-NULL-fallback logic as SQL."""
|
|
sources = await self._list_active_for_company(company_id)
|
|
return [s for s in sources if _is_source_due(s, now, company_next_run)]
|
|
|
|
async def company_has_due_work(
|
|
self, company_id: uuid.UUID, now: datetime, company_next_run: datetime | None
|
|
) -> bool:
|
|
"""Whether the scheduler should enqueue a run for this company: a
|
|
company with zero sources yet has nothing to check per-source - it
|
|
stays gated purely by its own next_run (the company's first-ever
|
|
run, which is what triggers source discovery), same as before
|
|
per-source scheduling existed. A company with sources is due if any
|
|
active one is (see list_due_for_company)."""
|
|
sources = await self._list_active_for_company(company_id)
|
|
if not sources:
|
|
return _is_due(company_next_run, now)
|
|
return any(_is_source_due(s, now, company_next_run) for s in sources)
|
|
|
|
async def _list_active_for_company(self, company_id: uuid.UUID) -> list[Source]:
|
|
result = await self.db.execute(
|
|
select(Source).where(Source.company_id == company_id, Source.active.is_(True))
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
async def get_for_company(self, source_id: uuid.UUID, company_id: uuid.UUID) -> Source | None:
|
|
result = await self.db.execute(
|
|
select(Source).where(Source.id == source_id, Source.company_id == company_id)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
async def get_for_user(self, source_id: uuid.UUID, user_id: uuid.UUID) -> Source | None:
|
|
"""Ownership-checked lookup that doesn't require the caller to
|
|
already know the company_id (matches the spec's `/sources/{id}`
|
|
routes, which aren't nested under `/companies/{company_id}`)."""
|
|
result = await self.db.execute(
|
|
select(Source)
|
|
.join(Company, Company.id == Source.company_id)
|
|
.where(Source.id == source_id, Company.user_id == user_id)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
async def create(
|
|
self,
|
|
*,
|
|
company_id: uuid.UUID,
|
|
source_type: SourceType,
|
|
name: str,
|
|
base_url: str | None,
|
|
configuration_metadata: dict[str, Any] | None = None,
|
|
trust_score: float = 0.7,
|
|
) -> Source:
|
|
source = Source(
|
|
company_id=company_id,
|
|
source_type=source_type,
|
|
name=name,
|
|
base_url=base_url,
|
|
configuration_metadata=configuration_metadata or {},
|
|
trust_score=trust_score,
|
|
)
|
|
self.db.add(source)
|
|
await self.db.flush()
|
|
return source
|
|
|
|
async def delete(self, source: Source) -> None:
|
|
await self.db.delete(source)
|
|
await self.db.flush()
|
|
|
|
async def mark_checked(
|
|
self,
|
|
source: Source,
|
|
*,
|
|
status: SourceStatus,
|
|
checked_at: datetime,
|
|
success: bool,
|
|
) -> None:
|
|
source.status = status
|
|
source.last_checked = checked_at
|
|
if success:
|
|
source.last_successful_check = checked_at
|
|
source.failure_count = 0
|
|
else:
|
|
source.failure_count += 1
|
|
await self.db.flush()
|
|
|
|
|
|
class SourceDocumentRepository:
|
|
def __init__(self, db: AsyncSession) -> None:
|
|
self.db = db
|
|
|
|
async def exists_with_hash(self, source_id: uuid.UUID, content_hash: str) -> bool:
|
|
result = await self.db.execute(
|
|
select(SourceDocument.id).where(
|
|
SourceDocument.source_id == source_id,
|
|
SourceDocument.content_hash == content_hash,
|
|
)
|
|
)
|
|
return result.scalar_one_or_none() is not None
|
|
|
|
async def create(self, **kwargs: Any) -> SourceDocument:
|
|
document = SourceDocument(**kwargs)
|
|
self.db.add(document)
|
|
await self.db.flush()
|
|
return document
|
|
|
|
async def latest_for_source(
|
|
self, source_id: uuid.UUID, limit: int = 50
|
|
) -> list[SourceDocument]:
|
|
result = await self.db.execute(
|
|
select(SourceDocument)
|
|
.where(SourceDocument.source_id == source_id)
|
|
.order_by(SourceDocument.retrieved_date.desc())
|
|
.limit(limit)
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
async def delete_older_than(self, cutoff: datetime) -> int:
|
|
"""Data-retention purge target (DATA_RETENTION_DAYS). Only
|
|
SourceDocument is in scope - nothing else has a foreign key onto it
|
|
(see KNOWN_LIMITATIONS.md), so this can't cascade-delete a Snapshot,
|
|
DetectedChange, Alert, or Report a user might still want to see."""
|
|
result = await self.db.execute(
|
|
delete(SourceDocument).where(SourceDocument.retrieved_date < cutoff)
|
|
)
|
|
return result.rowcount or 0
|
|
|
|
|
|
class SnapshotRepository:
|
|
def __init__(self, db: AsyncSession) -> None:
|
|
self.db = db
|
|
|
|
async def list_for_company(self, company_id: uuid.UUID, limit: int = 50) -> list[Snapshot]:
|
|
result = await self.db.execute(
|
|
select(Snapshot)
|
|
.where(Snapshot.company_id == company_id)
|
|
.order_by(Snapshot.created_at.desc())
|
|
.limit(limit)
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
async def latest_for_source(self, source_id: uuid.UUID) -> Snapshot | None:
|
|
result = await self.db.execute(
|
|
select(Snapshot)
|
|
.where(Snapshot.source_id == source_id)
|
|
.order_by(Snapshot.created_at.desc())
|
|
.limit(1)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
async def create(self, **kwargs: Any) -> Snapshot:
|
|
snapshot = Snapshot(**kwargs)
|
|
self.db.add(snapshot)
|
|
await self.db.flush()
|
|
return snapshot
|
|
|
|
async def get_previous(self, source_id: uuid.UUID, before: Snapshot) -> Snapshot | None:
|
|
"""The snapshot immediately preceding `before` for this source -
|
|
what change detection diffs the new snapshot against."""
|
|
result = await self.db.execute(
|
|
select(Snapshot)
|
|
.where(
|
|
Snapshot.source_id == source_id,
|
|
Snapshot.id != before.id,
|
|
Snapshot.created_at <= before.created_at,
|
|
)
|
|
.order_by(Snapshot.created_at.desc())
|
|
.limit(1)
|
|
)
|
|
return result.scalar_one_or_none()
|