saas_backend/alembic/versions/d1a3c5e7f508_identity_providers.py

"""Signing in with somebody else's identity provider.

Until now "SSO" here meant the platform signing users *into modules* — outbound.
This is the direction enterprise customers mean: a workspace points at its own
Azure AD, Okta or Google, and its people sign in there rather than holding a
password on this platform.

Three tables, all workspace-scoped and all under row-level security. A provider's
configuration is a workspace's own business, and `user_identities` says which of
their staff exist at which external directory.

Revision ID: d1a3c5e7f508
Revises: c9e2a4b6d407
"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects.postgresql import UUID

revision: str = "d1a3c5e7f508"
down_revision: Union[str, Sequence[str], None] = "c9e2a4b6d407"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

SCOPED = ("identity_providers", "user_identities", "sso_login_states")


def upgrade() -> None:
    op.create_table(
        "identity_providers",
        sa.Column("id", UUID(as_uuid=True), primary_key=True,
                  server_default=sa.text("gen_random_uuid()")),
        sa.Column("tenant_id", UUID(as_uuid=True),
                  sa.ForeignKey("tenants.id", ondelete="CASCADE"), nullable=False),
        sa.Column("kind", sa.String(10), nullable=False, server_default="OIDC"),
        sa.Column("name", sa.String(150), nullable=False),
        sa.Column("slug", sa.String(100), nullable=False),
        sa.Column("enabled", sa.Boolean(), nullable=False, server_default=sa.false()),
        sa.Column("issuer", sa.String(500), nullable=True),
        sa.Column("client_id", sa.String(255), nullable=True),
        sa.Column("client_secret_enc", sa.Text(), nullable=True),
        sa.Column("scopes", sa.String(500), nullable=False,
                  server_default="openid email profile"),
        sa.Column("authorization_endpoint", sa.String(500), nullable=True),
        sa.Column("token_endpoint", sa.String(500), nullable=True),
        sa.Column("jwks_uri", sa.String(500), nullable=True),
        sa.Column("discovered_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("allowed_domains", sa.Text(), nullable=True),
        sa.Column("jit_provisioning", sa.Boolean(), nullable=False,
                  server_default=sa.true()),
        sa.Column("default_role_id", UUID(as_uuid=True),
                  sa.ForeignKey("roles.id", ondelete="SET NULL"), nullable=True),
        sa.Column("link_existing_by_email", sa.Boolean(), nullable=False,
                  server_default=sa.false()),
        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False,
                  server_default=sa.func.now()),
        sa.Column("updated_at", sa.DateTime(timezone=True),
                  server_default=sa.func.now()),
        sa.UniqueConstraint("tenant_id", "slug", name="uq_identity_provider_slug"),
    )
    op.create_index("ix_identity_providers_tenant", "identity_providers", ["tenant_id"])

    op.create_table(
        "user_identities",
        sa.Column("id", UUID(as_uuid=True), primary_key=True,
                  server_default=sa.text("gen_random_uuid()")),
        sa.Column("tenant_id", UUID(as_uuid=True),
                  sa.ForeignKey("tenants.id", ondelete="CASCADE"), nullable=False),
        sa.Column("user_id", UUID(as_uuid=True),
                  sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False),
        sa.Column("provider_id", UUID(as_uuid=True),
                  sa.ForeignKey("identity_providers.id", ondelete="CASCADE"),
                  nullable=False),
        sa.Column("subject", sa.String(255), nullable=False),
        sa.Column("last_login_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False,
                  server_default=sa.func.now()),
        sa.UniqueConstraint("provider_id", "subject", name="uq_user_identity_subject"),
    )
    op.create_index("ix_user_identities_tenant", "user_identities", ["tenant_id"])
    op.create_index("ix_user_identities_user", "user_identities", ["user_id"])

    op.create_table(
        "sso_login_states",
        sa.Column("id", UUID(as_uuid=True), primary_key=True,
                  server_default=sa.text("gen_random_uuid()")),
        sa.Column("tenant_id", UUID(as_uuid=True),
                  sa.ForeignKey("tenants.id", ondelete="CASCADE"), nullable=False),
        sa.Column("provider_id", UUID(as_uuid=True),
                  sa.ForeignKey("identity_providers.id", ondelete="CASCADE"),
                  nullable=False),
        sa.Column("state", sa.String(128), nullable=False, unique=True),
        sa.Column("nonce", sa.String(128), nullable=False),
        sa.Column("code_verifier", sa.String(256), nullable=False),
        sa.Column("redirect_to", sa.Text(), nullable=True),
        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False,
                  server_default=sa.func.now()),
        sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
    )
    op.create_index("ix_sso_login_states_state", "sso_login_states", ["state"])
    op.create_index("ix_sso_login_states_expiry", "sso_login_states", ["expires_at"])

    for table in SCOPED:
        op.execute(f"ALTER TABLE {table} ENABLE ROW LEVEL SECURITY")
        op.execute(f"ALTER TABLE {table} FORCE ROW LEVEL SECURITY")
        op.execute(
            f"""
            CREATE POLICY tenant_isolation ON {table}
            USING (
                current_setting('app.bypass_rls', true) = 'on'
                OR tenant_id = NULLIF(current_setting('app.tenant_id', true), '')::uuid
            )
            WITH CHECK (
                current_setting('app.bypass_rls', true) = 'on'
                OR tenant_id = NULLIF(current_setting('app.tenant_id', true), '')::uuid
            )
            """
        )


def downgrade() -> None:
    for table in SCOPED:
        op.execute(f"DROP POLICY IF EXISTS tenant_isolation ON {table}")
    op.drop_index("ix_sso_login_states_expiry", table_name="sso_login_states")
    op.drop_index("ix_sso_login_states_state", table_name="sso_login_states")
    op.drop_table("sso_login_states")
    op.drop_index("ix_user_identities_user", table_name="user_identities")
    op.drop_index("ix_user_identities_tenant", table_name="user_identities")
    op.drop_table("user_identities")
    op.drop_index("ix_identity_providers_tenant", table_name="identity_providers")
    op.drop_table("identity_providers")