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).
58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.company_enrichment import CompanyEnrichment
|
|
from app.models.enums import EnrichmentStatus
|
|
|
|
|
|
class CompanyEnrichmentRepository:
|
|
def __init__(self, db: AsyncSession) -> None:
|
|
self.db = db
|
|
|
|
async def get_for_company(self, company_id: uuid.UUID) -> CompanyEnrichment | None:
|
|
result = await self.db.execute(
|
|
select(CompanyEnrichment).where(CompanyEnrichment.company_id == company_id)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
async def upsert(
|
|
self,
|
|
company_id: uuid.UUID,
|
|
*,
|
|
status: EnrichmentStatus,
|
|
data: dict[str, Any],
|
|
errors: dict[str, str],
|
|
credits_spent: int | None,
|
|
fetched_at: datetime,
|
|
) -> CompanyEnrichment:
|
|
"""One row per company (unique FK) - the second-ever call for a
|
|
company would only happen via a manual re-run, never automatically
|
|
(enrichment fires once, at onboarding - see enrichment_service)."""
|
|
existing = await self.get_for_company(company_id)
|
|
if existing is not None:
|
|
existing.status = status
|
|
existing.data = data
|
|
existing.errors = errors
|
|
existing.credits_spent = credits_spent
|
|
existing.fetched_at = fetched_at
|
|
await self.db.flush()
|
|
return existing
|
|
|
|
enrichment = CompanyEnrichment(
|
|
company_id=company_id,
|
|
status=status,
|
|
data=data,
|
|
errors=errors,
|
|
credits_spent=credits_spent,
|
|
fetched_at=fetched_at,
|
|
)
|
|
self.db.add(enrichment)
|
|
await self.db.flush()
|
|
return enrichment
|