"""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