company-ai-platform/app/modules/personalization/services/erasure.py

from collections.abc import Callable, Mapping, Sequence
from datetime import timedelta
from hashlib import sha256
from hmac import compare_digest
from secrets import token_hex
from typing import Any
from uuid import uuid4

from fastapi import HTTPException, status
from sqlalchemy import delete, select
from sqlalchemy.orm import Session

from app.core.utils.time import utc_now
from app.modules.ai_memory.models import AIMemoryEntry
from app.modules.business.models import MarketWatchlist
from app.modules.personalization.constants import ERASURE_CONFIRMATION_TTL_MINUTES
from app.modules.personalization.models import (
    AIConversation,
    AIConversationMessage,
    PersonalDataErasureRequest,
    UserPreference,
)
from app.modules.personalization.schemas import ErasureConfirmation, ErasureResult

ErasureHook = Callable[
    [Session, int, str],
    int | Mapping[str, int] | None,
]


class PersonalDataErasureService:
    """Issue one-time confirmations and erase owner data in one transaction."""

    def __init__(self, db: Session):
        self.db = db

    def request_confirmation(
        self,
        owner_id: int,
        *,
        ttl_minutes: int = ERASURE_CONFIRMATION_TTL_MINUTES,
    ) -> ErasureConfirmation:
        _validate_owner(owner_id)
        if ttl_minutes <= 0 or ttl_minutes > 60:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail="Invalid erasure confirmation lifetime",
            )
        confirmation_code = token_hex(4).upper()
        now = utc_now()
        expires_at = now + timedelta(minutes=ttl_minutes)
        record = self.db.execute(
            select(PersonalDataErasureRequest).where(
                PersonalDataErasureRequest.owner_id == owner_id
            )
        ).scalar_one_or_none()
        if record is None:
            record = PersonalDataErasureRequest(
                owner_id=owner_id,
                token_hash=_token_hash(owner_id, confirmation_code),
                expires_at=expires_at,
            )
            self.db.add(record)
        else:
            record.token_hash = _token_hash(owner_id, confirmation_code)
            record.expires_at = expires_at
            record.updated_at = now
        self.db.commit()
        return ErasureConfirmation(
            confirmation_code=confirmation_code,
            expires_at=expires_at,
        )

    def confirm_and_erase(
        self,
        owner_id: int,
        confirmation_code: str,
        *,
        before_hooks: Sequence[ErasureHook] = (),
        extra_hooks: Sequence[ErasureHook] = (),
    ) -> ErasureResult:
        """Erase core personal tables and run integration hooks before one commit."""

        _validate_owner(owner_id)
        request = self.db.execute(
            select(PersonalDataErasureRequest)
            .where(PersonalDataErasureRequest.owner_id == owner_id)
            .with_for_update()
        ).scalar_one_or_none()
        if request is None or request.expires_at <= utc_now():
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail="Erasure confirmation is invalid or expired",
            )
        supplied_hash = _token_hash(owner_id, confirmation_code.strip().upper())
        if not compare_digest(supplied_hash, request.token_hash):
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail="Erasure confirmation is invalid or expired",
            )

        anonymous_id = f"anonymous-{uuid4().hex}"
        deleted: dict[str, int] = {}
        extra: dict[str, Any] = {}
        try:
            for index, hook in enumerate(before_hooks):
                _merge_hook_result(
                    hook(self.db, owner_id, anonymous_id),
                    index=index,
                    prefix="before",
                    deleted=deleted,
                    extra=extra,
                )
            conversation_ids = select(AIConversation.id).where(
                AIConversation.owner_id == owner_id
            )
            deleted["conversation_messages"] = _row_count(
                self.db.execute(
                    delete(AIConversationMessage).where(
                        AIConversationMessage.conversation_id.in_(conversation_ids)
                    )
                ).rowcount
            )
            deleted["conversations"] = _row_count(
                self.db.execute(
                    delete(AIConversation).where(AIConversation.owner_id == owner_id)
                ).rowcount
            )
            deleted["preferences"] = _row_count(
                self.db.execute(
                    delete(UserPreference).where(UserPreference.owner_id == owner_id)
                ).rowcount
            )
            deleted["ai_memory"] = _row_count(
                self.db.execute(
                    delete(AIMemoryEntry).where(AIMemoryEntry.owner_id == owner_id)
                ).rowcount
            )
            deleted["watchlist"] = _row_count(
                self.db.execute(
                    delete(MarketWatchlist).where(MarketWatchlist.owner_id == owner_id)
                ).rowcount
            )
            self.db.delete(request)

            for index, hook in enumerate(extra_hooks):
                _merge_hook_result(
                    hook(self.db, owner_id, anonymous_id),
                    index=index,
                    prefix="extra",
                    deleted=deleted,
                    extra=extra,
                )
            self.db.commit()
        except Exception:
            self.db.rollback()
            raise
        return ErasureResult(
            anonymous_id=anonymous_id,
            deleted=deleted,
            extra=extra,
        )

    def purge_expired_confirmations(self) -> int:
        result = self.db.execute(
            delete(PersonalDataErasureRequest).where(
                PersonalDataErasureRequest.expires_at <= utc_now()
            )
        )
        count = _row_count(result.rowcount)
        self.db.commit()
        return count


def _token_hash(owner_id: int, confirmation_code: str) -> str:
    return sha256(f"{owner_id}\0{confirmation_code}".encode("utf-8")).hexdigest()


def _validate_owner(owner_id: int) -> None:
    if owner_id <= 0:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Valid erasure owner is required",
        )


def _row_count(value: int | None) -> int:
    return max(0, int(value or 0))


def _merge_hook_result(
    hook_result: int | Mapping[str, int] | None,
    *,
    index: int,
    prefix: str,
    deleted: dict[str, int],
    extra: dict[str, Any],
) -> None:
    if isinstance(hook_result, Mapping):
        for key, value in hook_result.items():
            deleted[str(key)] = int(value)
    elif isinstance(hook_result, int):
        deleted[f"{prefix}_{index}"] = hook_result
    elif hook_result is not None:
        extra[f"{prefix}_{index}"] = hook_result