from __future__ import annotations import uuid from datetime import datetime from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload from app.models.company import Company, CompanyAlias, Competitor from app.models.monitor_configuration import MonitorConfiguration def _with_relations(stmt): return stmt.options( selectinload(Company.aliases), selectinload(Company.competitors), selectinload(Company.monitor_configuration), selectinload(Company.enrichment), # Loaded eagerly so ORM-level cascade-delete (SQLite doesn't enforce # FK ON DELETE CASCADE without a pragma this app doesn't set) can # actually see the children to remove when a Company is deleted. selectinload(Company.notification_links), ) class CompanyRepository: def __init__(self, db: AsyncSession) -> None: self.db = db async def get_for_user(self, company_id: uuid.UUID, user_id: uuid.UUID) -> Company | None: stmt = _with_relations( select(Company).where(Company.id == company_id, Company.user_id == user_id) ) result = await self.db.execute(stmt) return result.scalar_one_or_none() async def get_by_id(self, company_id: uuid.UUID) -> Company | None: """Unscoped lookup for trusted internal callers (Celery tasks) that already have the company_id from a source they control - not for anything reachable from an HTTP request without an ownership check.""" stmt = _with_relations(select(Company).where(Company.id == company_id)) result = await self.db.execute(stmt) return result.scalar_one_or_none() async def list_for_user(self, user_id: uuid.UUID) -> list[Company]: stmt = _with_relations( select(Company).where(Company.user_id == user_id).order_by(Company.created_at.desc()) ) result = await self.db.execute(stmt) return list(result.scalars().all()) async def count_for_user(self, user_id: uuid.UUID) -> int: result = await self.db.execute( select(func.count()).select_from(Company).where(Company.user_id == user_id) ) return int(result.scalar_one()) async def slug_exists_for_user(self, user_id: uuid.UUID, slug: str) -> bool: result = await self.db.execute( select(func.count()) .select_from(Company) .where(Company.user_id == user_id, Company.slug == slug) ) return int(result.scalar_one()) > 0 async def name_exists_for_user(self, user_id: uuid.UUID, name: str) -> bool: result = await self.db.execute( select(func.count()) .select_from(Company) .where(Company.user_id == user_id, func.lower(Company.name) == name.lower()) ) return int(result.scalar_one()) > 0 async def create( self, *, user_id: uuid.UUID, name: str, slug: str, official_website: str | None, description: str | None, monitoring_focus: str | None, industry: str | None, country: str | None, region: str | None, headquarters: str | None = None, public_identifiers: dict[str, str] | None = None, alias_names: list[str], competitor_names: list[str], ) -> Company: company = Company( user_id=user_id, name=name, slug=slug, official_website=official_website, description=description, monitoring_focus=monitoring_focus, industry=industry, country=country, region=region, headquarters=headquarters, public_identifiers=public_identifiers or {}, ) company.aliases = [CompanyAlias(alias=a) for a in alias_names if a.strip()] company.competitors = [Competitor(name=c) for c in competitor_names if c.strip()] self.db.add(company) await self.db.flush() return company async def delete(self, company: Company) -> None: await self.db.delete(company) await self.db.flush() async def replace_aliases(self, company: Company, alias_names: list[str]) -> None: company.aliases = [CompanyAlias(alias=a) for a in alias_names if a.strip()] async def replace_competitors(self, company: Company, competitor_names: list[str]) -> None: company.competitors = [Competitor(name=c) for c in competitor_names if c.strip()] class MonitorConfigurationRepository: def __init__(self, db: AsyncSession) -> None: self.db = db async def create_default( self, *, company_id: uuid.UUID, frequency_type, interval_minutes: int | None, cron_expression: str | None, timezone: str, severity_threshold, next_run, ) -> MonitorConfiguration: config = MonitorConfiguration( company_id=company_id, frequency_type=frequency_type, interval_minutes=interval_minutes, cron_expression=cron_expression, timezone=timezone, severity_threshold=severity_threshold, next_run=next_run, ) self.db.add(config) await self.db.flush() return config async def list_due(self, now: datetime) -> list[MonitorConfiguration]: """Enabled schedules whose next_run has arrived - what Celery Beat's sync_schedules task polls instead of requiring a static per-company beat_schedule entry (see ARCHITECTURE.md).""" result = await self.db.execute( select(MonitorConfiguration).where( MonitorConfiguration.enabled.is_(True), MonitorConfiguration.next_run.is_not(None), MonitorConfiguration.next_run <= now, ) ) return list(result.scalars().all()) async def list_enabled(self) -> list[MonitorConfiguration]: """Every enabled schedule regardless of next_run - since a company can now also become due purely via a per-source frequency override even when its own default next_run isn't due yet (see SourceRepository.list_due_for_company), the scheduler needs this broader set to check, not just the ones already due on the company-level clock.""" result = await self.db.execute( select(MonitorConfiguration).where(MonitorConfiguration.enabled.is_(True)) ) return list(result.scalars().all())