app/social/database.py

from __future__ import annotations

import contextvars
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

from sqlalchemy import event, inspect, text
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)

from app.analytics import models as analytics_models  # noqa: F401
from app.core.config import Settings
from app.core.database_url import normalize_async_database_url
from app.social.models import SocialBase

_ANALYTICS_MODELS_REGISTERED = analytics_models

REQUIRED_SOCIAL_TABLES = frozenset(
    {
        "social_accounts",
        "social_account_tokens",
        "social_account_capabilities",
        "media_variants",
        "social_media_assets",
        "social_campaigns",
        "social_posts",
        "social_post_targets",
        "social_post_media",
        "social_schedules",
        "social_jobs",
        "social_job_attempts",
        "oauth_states",
        "social_webhook_events",
        "social_post_metrics",
        "social_audit_events",
        "social_publishing_batches",
        "social_publishing_batch_items",
        "analytics_sync_runs",
        "analytics_metric_snapshots",
        "analytics_post_metrics",
        "analytics_platform_metrics",
    }
)

REQUIRED_SOCIAL_COLUMNS: dict[str, frozenset[str]] = {
    "social_media_assets": frozenset({"canonical_asset_id"}),
    "social_posts": frozenset(
        {"project_id", "canonical_caption", "canonical_hashtags", "revision"}
    ),
    "social_schedules": frozenset({"revision"}),
    "social_post_targets": frozenset({"cancellation_requested_at"}),
    "social_jobs": frozenset({"provider_state_encrypted", "cancellation_requested_at"}),
    "oauth_states": frozenset({"requested_account_type", "requested_scopes"}),
    "analytics_sync_runs": frozenset(
        {"workspace_id", "idempotency_key", "status", "date_from", "date_to"}
    ),
}

_trusted_worker_context: contextvars.ContextVar[bool] = contextvars.ContextVar(
    "trusted_social_worker_context", default=False
)


class SocialDatabase:
    """Social persistence with migration-only production schema changes."""

    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self.database_url = normalize_async_database_url(settings.resolved_social_database_url)
        self.engine: AsyncEngine = create_async_engine(self.database_url, pool_pre_ping=True)
        if self.database_url.startswith("sqlite"):
            event.listen(self.engine.sync_engine, "connect", self._configure_sqlite)
        self.session_factory = async_sessionmaker(
            self.engine, expire_on_commit=False, class_=AsyncSession
        )
        self.worker_database_url = normalize_async_database_url(settings.social_worker_database_url)
        self.worker_engine: AsyncEngine | None = None
        self.worker_session_factory: async_sessionmaker[AsyncSession] | None = None
        if self.worker_database_url:
            self.worker_engine = create_async_engine(self.worker_database_url, pool_pre_ping=True)
            if self.worker_database_url.startswith("sqlite"):
                event.listen(self.worker_engine.sync_engine, "connect", self._configure_sqlite)
            self.worker_session_factory = async_sessionmaker(
                self.worker_engine, expire_on_commit=False, class_=AsyncSession
            )

    @staticmethod
    def _configure_sqlite(dbapi_connection: object, _record: object) -> None:
        cursor = dbapi_connection.cursor()  # type: ignore[attr-defined]
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.execute("PRAGMA busy_timeout=5000")
        cursor.close()

    async def initialize(self) -> None:
        if self.settings.social_auto_migrate:
            async with self.engine.begin() as connection:
                await connection.run_sync(SocialBase.metadata.create_all)

    @property
    def is_postgres(self) -> bool:
        return self.database_url.startswith(("postgresql", "postgres"))

    async def verify_execution_boundaries(self) -> None:
        """Fail closed when a Postgres deployment cannot enforce RLS safely."""

        if not self.is_postgres or not self.settings.social_enforce_rls:
            return
        tenant_role = self.settings.social_tenant_database_role.strip()
        if not tenant_role:
            raise RuntimeError(
                "SOCIAL_TENANT_DATABASE_ROLE is required when SOCIAL_ENFORCE_RLS is enabled."
            )
        tenant = await self._role_attributes(self.engine)
        if tenant["role"] != tenant_role:
            raise RuntimeError(
                "SOCIAL_DATABASE_URL is not connected as SOCIAL_TENANT_DATABASE_ROLE."
            )
        if tenant["bypass_rls"] or tenant["superuser"]:
            raise RuntimeError(
                "SOCIAL_DATABASE_URL must use a non-privileged tenant role, never a service role."
            )
        # Startup also adopts historic API-key tenant rows through this
        # boundary, so every enabled PostgreSQL social deployment needs it,
        # even if the scheduler is temporarily disabled.
        if not self.worker_engine:
            raise RuntimeError(
                "SOCIAL_WORKER_DATABASE_URL is required for PostgreSQL social access."
            )
        worker_role = self.settings.social_worker_database_role.strip()
        if not worker_role:
            raise RuntimeError("SOCIAL_WORKER_DATABASE_ROLE is required for trusted worker access.")
        worker = await self._role_attributes(self.worker_engine)
        if worker["role"] != worker_role or not worker["bypass_rls"]:
            raise RuntimeError(
                "SOCIAL_WORKER_DATABASE_URL must use the configured BYPASSRLS worker role."
            )

    @staticmethod
    async def _role_attributes(engine: AsyncEngine) -> dict[str, object]:
        async with engine.connect() as connection:
            row = (
                (
                    await connection.execute(
                        text(
                            "select current_user as role, r.rolbypassrls as bypass_rls, r.rolsuper as superuser "
                            "from pg_roles r where r.rolname = current_user"
                        )
                    )
                )
                .mappings()
                .one_or_none()
            )
        if row is None:
            raise RuntimeError("Unable to verify the active PostgreSQL database role.")
        return dict(row)

    async def schema_ready(self) -> bool:
        """Check the complete Phase 1 schema without changing the database."""
        async with self.engine.connect() as connection:
            tables, columns = await connection.run_sync(self._schema_snapshot)
        return REQUIRED_SOCIAL_TABLES.issubset(tables) and all(
            required.issubset(columns.get(table, set()))
            for table, required in REQUIRED_SOCIAL_COLUMNS.items()
        )

    async def missing_tables(self) -> list[str]:
        """Return absent required tables for an actionable startup warning."""
        async with self.engine.connect() as connection:
            tables, columns = await connection.run_sync(self._schema_snapshot)
        missing = list(REQUIRED_SOCIAL_TABLES - tables)
        for table, required in REQUIRED_SOCIAL_COLUMNS.items():
            missing.extend(f"{table}.{column}" for column in required - columns.get(table, set()))
        return sorted(missing)

    async def adopt_legacy_workspace(
        self, *, legacy_workspace_id: str, workspace_id: str, user_id: str
    ) -> int:
        """Move historic API-key tenant rows to the authoritative workspace."""
        if legacy_workspace_id == workspace_id:
            raise RuntimeError("Legacy and authoritative workspace IDs must differ.")
        workspace_tables = (
            "social_accounts",
            "media_variants",
            "social_media_assets",
            "social_campaigns",
            "social_posts",
            "social_jobs",
            "social_webhook_events",
            "social_audit_events",
            "oauth_states",
            "social_publishing_batches",
        )
        changed = 0
        async with self.worker_session() as session:
            for table in workspace_tables:
                result = await session.execute(
                    text(
                        f"update {table} set workspace_id = :workspace_id "
                        "where workspace_id = :legacy_workspace_id"
                    ),
                    {"workspace_id": workspace_id, "legacy_workspace_id": legacy_workspace_id},
                )
                changed += max(0, int(result.rowcount or 0))
            for table, column in (("social_posts", "created_by"), ("oauth_states", "user_id")):
                result = await session.execute(
                    text(
                        f"update {table} set {column} = :user_id "
                        f"where {column} = :legacy_workspace_id"
                    ),
                    {"user_id": user_id, "legacy_workspace_id": legacy_workspace_id},
                )
                changed += max(0, int(result.rowcount or 0))
            await session.commit()
        return changed

    @staticmethod
    def _schema_snapshot(connection: object) -> tuple[set[str], dict[str, set[str]]]:
        inspector = inspect(connection)
        tables = set(inspector.get_table_names())
        columns = {
            table: {column["name"] for column in inspector.get_columns(table)}
            for table in REQUIRED_SOCIAL_COLUMNS
            if table in tables
        }
        return tables, columns

    async def close(self) -> None:
        await self.engine.dispose()
        if self.worker_engine is not None:
            await self.worker_engine.dispose()

    @asynccontextmanager
    async def session(self, workspace_id: str | None = None) -> AsyncIterator[AsyncSession]:
        """Open a tenant session; no-context access is OAuth-state compatibility only."""
        if workspace_id is None:
            async with self.oauth_session() as session:
                yield session
            return
        context = (
            self.worker_tenant_session(workspace_id)
            if _trusted_worker_context.get()
            else self.tenant_session(workspace_id)
        )
        async with context as session:
            yield session

    @asynccontextmanager
    async def tenant_session(self, workspace_id: str) -> AsyncIterator[AsyncSession]:
        if not workspace_id:
            raise RuntimeError("A tenant session requires an authoritative workspace ID.")
        async with self.session_factory() as session:
            if self.is_postgres:
                # RLS policies read this transaction-local tenant identity.
                await session.execute(
                    text("select set_config('app.workspace_id', :workspace_id, true)"),
                    {"workspace_id": workspace_id},
                )
            yield session

    @asynccontextmanager
    async def oauth_session(self) -> AsyncIterator[AsyncSession]:
        """The sole non-tenant API session, for random single-use OAuth state."""
        async with self.session_factory() as session:
            yield session

    @asynccontextmanager
    async def worker_session(self) -> AsyncIterator[AsyncSession]:
        """Backend-only cross-workspace session for scheduler, jobs, and Vault."""
        if self.is_postgres:
            if self.worker_session_factory is None:
                raise RuntimeError("Trusted worker database access is not configured.")
            async with self.worker_session_factory() as session:
                yield session
            return
        # SQLite has no RLS. It remains supported for local/unit-test use only.
        async with self.session_factory() as session:
            yield session

    @asynccontextmanager
    async def worker_tenant_session(self, workspace_id: str) -> AsyncIterator[AsyncSession]:
        """Trusted worker session annotated with the job's tenant for auditability."""
        if not workspace_id:
            raise RuntimeError("A worker tenant session requires a workspace ID.")
        if not self.is_postgres:
            async with self.tenant_session(workspace_id) as session:
                yield session
            return
        if self.worker_session_factory is None:
            raise RuntimeError("Trusted worker database access is not configured.")
        async with self.worker_session_factory() as session:
            await session.execute(
                text("select set_config('app.workspace_id', :workspace_id, true)"),
                {"workspace_id": workspace_id},
            )
            yield session

    @asynccontextmanager
    async def worker_boundary(self) -> AsyncIterator[None]:
        """Mark a scheduler/publisher call tree as trusted worker execution."""
        if self.is_postgres and self.worker_session_factory is None:
            raise RuntimeError("Trusted worker database access is not configured.")
        token = _trusted_worker_context.set(True)
        try:
            yield
        finally:
            _trusted_worker_context.reset(token)