"""switch tenant RLS from GUC to pg_session_jwt + authenticated role

Replaces the `current_setting('app.tenant_id', true)` policies (introduced
in 0023 / 0025) with `auth.session()->>'activeRestaurantId'` policies that
consume the Better Auth JWT loaded into `pg_session_jwt`. This makes RLS
actually enforce when the backend connects as Neon's `authenticated` role
(NOBYPASSRLS), which is now the path for every user-facing FastAPI request.

This migration assumes:

1. The `pg_session_jwt` extension is available on the cluster (Neon installs
   it project-wide when "Auth (Neon RLS)" is enabled — see
   `backend/OPERATIONS.md` → "Neon RLS Authorize").
2. Better Auth's `definePayload` emits an `activeRestaurantId` claim
   (handled by `frontend/src/lib/auth.config.ts` from this same change set).
3. The `authenticated` role exists on the branch — also a side-effect of
   enabling Neon RLS in the Console.

The migration is idempotent on policy DDL (drop-if-exists then recreate) so
it can re-run against branches that already saw an earlier RLS rollout.

Revision ID: 0030_rls_via_pg_session_jwt
Revises: 0029_create_invitation_token
Create Date: 2026-05-25 18:30:00.000000
"""

from collections.abc import Sequence

from alembic import op

revision: str = "0030_rls_via_pg_session_jwt"
down_revision: str | None = "0029_create_invitation_token"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


# Every table that carries a restaurant_id and must enforce tenant isolation.
# Keep in lockstep with `openspec/changes/row-level-security/specs/tenant-rls/spec.md`.
TENANT_TABLES: tuple[str, ...] = (
    "zone",
    "floor_table",
    "reservation",
    "customer",
    "order",
    "menu_item",
    "service_block",
    "service_block_override",
    "table_combination",
    "conversation",
    "knowledge_document",
    "faq_entry",
    "common_question",
    "chair",
    "order_item",
    "message",
    "service_block_zones",
    "combined_chair_config",
    "notification",
    "whatsapp_account",
    "whatsapp_template",
)


def upgrade() -> None:
    # 1. Ensure the JWT-session extension is present. Idempotent.
    op.execute("CREATE EXTENSION IF NOT EXISTS pg_session_jwt")

    # 2. Replace each table's tenant_isolation policy with the JWT-based form
    #    and keep FORCE RLS on (the authenticated role is NOBYPASSRLS, so
    #    FORCE only affects the owner role's queries — fine, we want owner
    #    bypass to remain explicit on its own pool).
    for table in TENANT_TABLES:
        op.execute(f'ALTER TABLE "{table}" ENABLE ROW LEVEL SECURITY')
        op.execute(f'ALTER TABLE "{table}" FORCE ROW LEVEL SECURITY')
        op.execute(f'DROP POLICY IF EXISTS tenant_isolation ON "{table}"')
        op.execute(
            f'CREATE POLICY tenant_isolation ON "{table}" '
            f"FOR ALL "
            f"USING (restaurant_id = (auth.session() ->> 'activeRestaurantId')) "
            f"WITH CHECK (restaurant_id = (auth.session() ->> 'activeRestaurantId'))"
        )

    # 3. Grant DML on every tenant table to the `authenticated` role so the
    #    runtime can actually issue queries. The RLS policy still filters
    #    rows; without the grant the role would get permission-denied errors
    #    instead of fail-closed (empty result) behavior.
    #
    #    These grants are no-ops on branches where the `authenticated` role
    #    has not been provisioned yet (Neon RLS not enabled). We guard with
    #    a DO block so the migration succeeds on branches in transition.
    op.execute(
        """
        DO $$
        BEGIN
            IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'authenticated') THEN
                GRANT USAGE ON SCHEMA public TO authenticated;
                GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public
                    TO authenticated;
                GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO authenticated;
            END IF;
        END $$;
        """
    )

    # 4. Future tables created by neondb_owner inherit the same grants so a
    #    new tenant table doesn't accidentally fail-closed for the wrong
    #    reason (permission denied vs. RLS empty result).
    op.execute(
        """
        DO $$
        BEGIN
            IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'authenticated') THEN
                ALTER DEFAULT PRIVILEGES FOR ROLE neondb_owner IN SCHEMA public
                    GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO authenticated;
                ALTER DEFAULT PRIVILEGES FOR ROLE neondb_owner IN SCHEMA public
                    GRANT USAGE, SELECT ON SEQUENCES TO authenticated;
            END IF;
        END $$;
        """
    )


def downgrade() -> None:
    # Restore the GUC-based policies. We do not drop pg_session_jwt because
    # other components on the cluster may rely on it (it's a project-level
    # Neon feature).
    for table in TENANT_TABLES:
        op.execute(f'DROP POLICY IF EXISTS tenant_isolation ON "{table}"')
        op.execute(
            f'CREATE POLICY tenant_isolation ON "{table}" '
            f"FOR ALL "
            f"USING (restaurant_id = current_setting('app.tenant_id', true)) "
            f"WITH CHECK (restaurant_id = current_setting('app.tenant_id', true))"
        )

    op.execute(
        """
        DO $$
        BEGIN
            IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'authenticated') THEN
                ALTER DEFAULT PRIVILEGES FOR ROLE neondb_owner IN SCHEMA public
                    REVOKE SELECT, INSERT, UPDATE, DELETE ON TABLES FROM authenticated;
                ALTER DEFAULT PRIVILEGES FOR ROLE neondb_owner IN SCHEMA public
                    REVOKE USAGE, SELECT ON SEQUENCES FROM authenticated;
                REVOKE SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public
                    FROM authenticated;
                REVOKE USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public FROM authenticated;
                REVOKE USAGE ON SCHEMA public FROM authenticated;
            END IF;
        END $$;
        """
    )