CIAgent/apps/api/app/services/provisioning_service.py

"""Pre-provisions a not-yet-existing account: the moment someone verifies
an email address an admin has queued up (see PendingProvisioning), they
get a deep copy of another user's per-user API keys and companies -
including every tab's worth of data (sources, source documents, monitoring
runs, reports, snapshots, detected changes, enrichment) - plus an email
notification destination for their own address linked to the copied
companies, and optionally an admin promotion.

Applied at email-verification time (see auth_service.verify_email), not
raw registration - registering an email doesn't prove ownership of it,
verifying does.
"""

from __future__ import annotations

import uuid
from typing import Any

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.errors import ConflictError, NotFoundError
from app.models.company import Company, CompanyAlias, Competitor
from app.models.company_enrichment import CompanyEnrichment
from app.models.detected_change import DetectedChange
from app.models.enums import NotificationType, SeverityLevel
from app.models.monitor_configuration import MonitorConfiguration
from app.models.monitoring_run import MonitoringRun
from app.models.notification_destination import (
    NotificationDestination,
    NotificationDestinationCompany,
)
from app.models.pending_provisioning import PendingProvisioning
from app.models.report import Report
from app.models.snapshot import Snapshot
from app.models.source import Source
from app.models.source_document import SourceDocument
from app.models.user import User
from app.models.user_api_key import UserApiKey
from app.repositories.pending_provisioning_repository import PendingProvisioningRepository
from app.repositories.user_repository import UserRepository


def _copy_row[ModelT](
    model_cls: type[ModelT], source_row: Any, overrides: dict[str, Any]
) -> ModelT:
    """A new ORM instance of the same class, with every column value copied
    from `source_row` except whatever `overrides` replaces (always at least
    `id`, plus any foreign keys that need remapping to the new owner's
    copies of their parents). Generic over every table involved so adding a
    new company-related table later doesn't require touching this file."""
    data = {
        col.name: getattr(source_row, col.name) for col in model_cls.__table__.columns  # type: ignore[attr-defined]
    }
    data.update(overrides)
    return model_cls(**data)


async def _clone_company(
    db: AsyncSession, source_company_id: uuid.UUID, target_user_id: uuid.UUID
) -> Company:
    id_map: dict[uuid.UUID, uuid.UUID] = {}

    def new_id(old_id: uuid.UUID) -> uuid.UUID:
        if old_id not in id_map:
            id_map[old_id] = uuid.uuid4()
        return id_map[old_id]

    company = await db.get(Company, source_company_id)
    assert company is not None

    new_company = _copy_row(Company, company, {"id": new_id(company.id), "user_id": target_user_id})
    db.add(new_company)
    await db.flush()

    for alias in (
        await db.execute(select(CompanyAlias).where(CompanyAlias.company_id == source_company_id))
    ).scalars():
        db.add(_copy_row(CompanyAlias, alias, {"id": uuid.uuid4(), "company_id": new_company.id}))

    for competitor in (
        await db.execute(select(Competitor).where(Competitor.company_id == source_company_id))
    ).scalars():
        db.add(
            _copy_row(Competitor, competitor, {"id": uuid.uuid4(), "company_id": new_company.id})
        )

    config = (
        await db.execute(
            select(MonitorConfiguration).where(MonitorConfiguration.company_id == source_company_id)
        )
    ).scalar_one_or_none()
    if config is not None:
        db.add(
            _copy_row(
                MonitorConfiguration, config, {"id": uuid.uuid4(), "company_id": new_company.id}
            )
        )

    enrichment = (
        await db.execute(
            select(CompanyEnrichment).where(CompanyEnrichment.company_id == source_company_id)
        )
    ).scalar_one_or_none()
    if enrichment is not None:
        db.add(
            _copy_row(
                CompanyEnrichment, enrichment, {"id": uuid.uuid4(), "company_id": new_company.id}
            )
        )

    sources = (
        (await db.execute(select(Source).where(Source.company_id == source_company_id)))
        .scalars()
        .all()
    )
    for source in sources:
        db.add(_copy_row(Source, source, {"id": new_id(source.id), "company_id": new_company.id}))
    await db.flush()

    for doc in (
        await db.execute(
            select(SourceDocument).where(SourceDocument.company_id == source_company_id)
        )
    ).scalars():
        db.add(
            _copy_row(
                SourceDocument,
                doc,
                {
                    "id": uuid.uuid4(),
                    "company_id": new_company.id,
                    "source_id": new_id(doc.source_id),
                },
            )
        )

    runs = (
        (
            await db.execute(
                select(MonitoringRun).where(MonitoringRun.company_id == source_company_id)
            )
        )
        .scalars()
        .all()
    )
    for run in runs:
        db.add(_copy_row(MonitoringRun, run, {"id": new_id(run.id), "company_id": new_company.id}))
    await db.flush()

    for report in (
        await db.execute(select(Report).where(Report.company_id == source_company_id))
    ).scalars():
        db.add(
            _copy_row(
                Report,
                report,
                {
                    "id": uuid.uuid4(),
                    "company_id": new_company.id,
                    "monitoring_run_id": (
                        new_id(report.monitoring_run_id) if report.monitoring_run_id else None
                    ),
                },
            )
        )

    snapshots = (
        (await db.execute(select(Snapshot).where(Snapshot.company_id == source_company_id)))
        .scalars()
        .all()
    )
    for snapshot in snapshots:
        db.add(
            _copy_row(
                Snapshot,
                snapshot,
                {
                    "id": new_id(snapshot.id),
                    "company_id": new_company.id,
                    "source_id": new_id(snapshot.source_id) if snapshot.source_id else None,
                    "monitoring_run_id": (
                        new_id(snapshot.monitoring_run_id) if snapshot.monitoring_run_id else None
                    ),
                },
            )
        )
    await db.flush()

    for change in (
        await db.execute(
            select(DetectedChange).where(DetectedChange.company_id == source_company_id)
        )
    ).scalars():
        db.add(
            _copy_row(
                DetectedChange,
                change,
                {
                    "id": uuid.uuid4(),
                    "company_id": new_company.id,
                    "source_id": new_id(change.source_id) if change.source_id else None,
                    "monitoring_run_id": (
                        new_id(change.monitoring_run_id) if change.monitoring_run_id else None
                    ),
                    "previous_snapshot_id": (
                        new_id(change.previous_snapshot_id) if change.previous_snapshot_id else None
                    ),
                    "current_snapshot_id": (
                        new_id(change.current_snapshot_id) if change.current_snapshot_id else None
                    ),
                },
            )
        )

    return new_company


async def _copy_api_keys(
    db: AsyncSession, source_user_id: uuid.UUID, target_user_id: uuid.UUID
) -> None:
    keys = (
        (await db.execute(select(UserApiKey).where(UserApiKey.user_id == source_user_id)))
        .scalars()
        .all()
    )
    for key in keys:
        db.add(_copy_row(UserApiKey, key, {"id": uuid.uuid4(), "user_id": target_user_id}))


async def create_pending(
    db: AsyncSession, email: str, source_user_id: uuid.UUID, make_admin: bool
) -> PendingProvisioning:
    if await UserRepository(db).get_by_email(email) is not None:
        raise ConflictError("An account with this email already exists")
    if await PendingProvisioningRepository(db).get_by_email(email) is not None:
        raise ConflictError("This email is already queued for provisioning")
    record = await PendingProvisioningRepository(db).create(email, source_user_id, make_admin)
    await db.commit()
    return record


async def list_pending(db: AsyncSession) -> list[PendingProvisioning]:
    return await PendingProvisioningRepository(db).list_all()


async def delete_pending(db: AsyncSession, pending_id: uuid.UUID) -> None:
    repo = PendingProvisioningRepository(db)
    match = next((p for p in await repo.list_all() if p.id == pending_id), None)
    if match is None:
        raise NotFoundError("Pending provisioning entry not found")
    await repo.delete(match)
    await db.commit()


async def apply_if_pending(db: AsyncSession, new_user: User) -> bool:
    """Called right after a fresh account's email gets verified. Returns
    True if a pending provisioning entry matched and was applied (and
    consumed - a pending entry only ever fires once)."""
    pending = await PendingProvisioningRepository(db).get_by_email(new_user.email)
    if pending is None:
        return False

    await _copy_api_keys(db, pending.source_user_id, new_user.id)

    source_companies = (
        (await db.execute(select(Company).where(Company.user_id == pending.source_user_id)))
        .scalars()
        .all()
    )
    new_company_ids = [
        (await _clone_company(db, company.id, new_user.id)).id for company in source_companies
    ]

    if new_company_ids:
        destination = NotificationDestination(
            user_id=new_user.id,
            type=NotificationType.EMAIL,
            destination_value=new_user.email,
            verified=False,
            enabled=True,
            minimum_severity=SeverityLevel.MEDIUM,
        )
        db.add(destination)
        await db.flush()
        for company_id in new_company_ids:
            db.add(
                NotificationDestinationCompany(destination_id=destination.id, company_id=company_id)
            )

    if pending.make_admin:
        new_user.is_admin = True

    await PendingProvisioningRepository(db).delete(pending)
    return True