move sso logic to postgresé

Author: leoguillaume

api/alembic/versions/2026_06_26_1257-f94a77198389_add_sso_access.py

@@ -0,0 +1,45 @@
+"""add sso access
+
+Revision ID: f94a77198389
+Revises: 5138304df5f5
+Create Date: 2026-06-26 12:57:42.469421
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'f94a77198389'
+down_revision: Union[str, None] = '5138304df5f5'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('sso_policy_rule',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('type', sa.Enum('EMAIL', 'ORGANIZATION', 'ROLE', name='ssoaccessruletype'), nullable=False),
+    sa.Column('value', sa.String(), nullable=True),
+    sa.Column('role_id', sa.Integer(), nullable=True),
+    sa.Column('organization_id', sa.Integer(), nullable=True),
+    sa.Column('created', sa.DateTime(timezone=True), nullable=False),
+    sa.Column('updated', sa.DateTime(timezone=True), nullable=False),
+    sa.ForeignKeyConstraint(['organization_id'], ['organization.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['role_id'], ['role.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id'),
+    sa.UniqueConstraint('type', 'value', name='unique_sso_policy_type_value')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('sso_policy_rule')
+    op.execute("DROP TYPE IF EXISTS ssoaccessruletype")
+    # ### end Alembic commands ###

api/dependencies.py

@@ -7,7 +7,7 @@
 from redis.asyncio import Redis
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from api.domain.auth import AuthSsoSessionValidator
+from api.domain.auth import AuthSsoSessionValidator, SsoPolicyRepository
 from api.domain.key import KeyEncoder, KeyRepository
 from api.domain.model import ModelEnvironmentalImpactsComputer, ModelTokenizer
 from api.domain.organization import OrganizationRepository
@@ -37,6 +37,7 @@
     PostgresProviderRepository,
     PostgresRolesRepository,
     PostgresRouterRepository,
+    PostgresSsoPolicyRepository,
     PostgresUserRepository,
 )
 from api.infrastructure.redis import RedisProviderLoadBalancer, RedisProviderMetricsLogger, RedisRouterRateLimiter
@@ -52,6 +53,7 @@
 )
 from api.use_cases.admin.roles import CreateRoleUseCase, DeleteRoleUseCase, GetRolesUseCase, GetRoleUseCase, UpdateRoleUseCase
 from api.use_cases.admin.routers import CreateRouterUseCase, DeleteRouterUseCase, GetOneRouterUseCase, GetRoutersUseCase, UpdateRouterUseCase
+from api.use_cases.admin.sso import GetSsoPolicyUseCase, UpdateSsoPolicyUseCase
 from api.use_cases.admin.users import CreateUserUseCase, DeleteUserUseCase, GetOneUserUseCase, GetUsersUseCase
 from api.use_cases.auth import AuthLoginUseCase, AuthSsoLoginUseCase
 from api.use_cases.health import GetHealthModelsUseCase
@@ -163,6 +165,10 @@ def _router_repository(session: AsyncSession) -> PostgresRouterRepository:
     return PostgresRouterRepository(postgres_session=session, app_title=configuration.settings.app_title)
 
 
+def _sso_policy_repository(session: AsyncSession) -> SsoPolicyRepository:
+    return PostgresSsoPolicyRepository(postgres_session=session)
+
+
 def _limit_repository(session: AsyncSession) -> LimitRepository:
     return PostgresLimitRepository(postgres_session=session)
 
@@ -195,11 +201,10 @@ def auth_sso_login_use_case_factory(
 ) -> AuthSsoLoginUseCase:
     return AuthSsoLoginUseCase(
         key_repository=_key_repository(key_encoder=key_encoder, session=postgres_session),
-        organization_repository=_organization_repository(session=postgres_session),
         user_repository=_user_repository(session=postgres_session),
+        sso_policy_repository=_sso_policy_repository(session=postgres_session),
         auth_sso_session_validator=_auth_sso_session_validator(),
         auth_login_type=configuration.settings.auth_login_type,
-        auth_sso_default_role_id=configuration.settings.auth_sso_default_role_id,
         auth_login_session_duration=configuration.settings.auth_login_session_duration,
     )
 
@@ -234,23 +239,39 @@ def get_model_use_case_factory(postgres_session: AsyncSession = Depends(get_post
     return GetModelUseCase(router_repository=_router_repository(postgres_session))
 
 
-# user use cases
-def create_user_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> CreateUserUseCase:
-    return CreateUserUseCase(user_repository=_user_repository(postgres_session), user_password_encoder=_user_password_encoder())
+# provider use cases
+def create_provider_use_case_factory(
+    postgres_session: AsyncSession = Depends(get_postgres_session),
+    provider_client: ProviderClient = Depends(_provider_client),
+) -> CreateProviderUseCase:
+    return CreateProviderUseCase(
+        router_repository=_router_repository(postgres_session),
+        provider_repository=_provider_repository(postgres_session),
+        provider_gateway=_provider_gateway(provider_client=provider_client, provider_adapter_builder=HttpProviderAdapterBuilder()),
+    )
 
 
-def get_one_user_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetOneUserUseCase:
-    return GetOneUserUseCase(user_repository=_user_repository(postgres_session))
+def update_provider_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> UpdateProviderUseCase:
+    return UpdateProviderUseCase(
+        router_repository=_router_repository(postgres_session),
+        provider_repository=_provider_repository(postgres_session),
+    )
 
 
-def get_users_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetUsersUseCase:
-    return GetUsersUseCase(user_repository=_user_repository(postgres_session))
+def delete_provider_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> DeleteProviderUseCase:
+    return DeleteProviderUseCase(
+        provider_repository=_provider_repository(postgres_session),
+    )
 
 
-def delete_user_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> DeleteUserUseCase:
-    return DeleteUserUseCase(
-        user_repository=_user_repository(postgres_session),
-        router_repository=_router_repository(postgres_session),
+def get_one_provider_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetOneProviderUseCase:
+    return GetOneProviderUseCase(
+        provider_repository=_provider_repository(postgres_session),
+    )
+
+
+def get_providers_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetProvidersUseCase:
+    return GetProvidersUseCase(
         provider_repository=_provider_repository(postgres_session),
     )
 
@@ -345,38 +366,33 @@ def update_router_use_case_factory(postgres_session: AsyncSession = Depends(get_
     )
 
 
-# provider use cases
-def create_provider_use_case_factory(
-    postgres_session: AsyncSession = Depends(get_postgres_session),
-    provider_client: ProviderClient = Depends(_provider_client),
-) -> CreateProviderUseCase:
-    return CreateProviderUseCase(
-        router_repository=_router_repository(postgres_session),
-        provider_repository=_provider_repository(postgres_session),
-        provider_gateway=_provider_gateway(provider_client=provider_client, provider_adapter_builder=HttpProviderAdapterBuilder()),
+# sso policy use cases
+def get_sso_policy_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetSsoPolicyUseCase:
+    return GetSsoPolicyUseCase(
+        sso_policy_repository=_sso_policy_repository(session=postgres_session),
     )
 
 
-def update_provider_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> UpdateProviderUseCase:
-    return UpdateProviderUseCase(
-        router_repository=_router_repository(postgres_session),
-        provider_repository=_provider_repository(postgres_session),
-    )
+def update_sso_policy_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> UpdateSsoPolicyUseCase:
+    return UpdateSsoPolicyUseCase(sso_policy_repository=_sso_policy_repository(session=postgres_session))
 
 
-def delete_provider_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> DeleteProviderUseCase:
-    return DeleteProviderUseCase(
-        provider_repository=_provider_repository(postgres_session),
-    )
+# user use cases
+def create_user_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> CreateUserUseCase:
+    return CreateUserUseCase(user_repository=_user_repository(postgres_session), user_password_encoder=_user_password_encoder())
 
 
-def get_one_provider_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetOneProviderUseCase:
-    return GetOneProviderUseCase(
-        provider_repository=_provider_repository(postgres_session),
-    )
+def get_one_user_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetOneUserUseCase:
+    return GetOneUserUseCase(user_repository=_user_repository(postgres_session))
 
 
-def get_providers_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetProvidersUseCase:
-    return GetProvidersUseCase(
+def get_users_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetUsersUseCase:
+    return GetUsersUseCase(user_repository=_user_repository(postgres_session))
+
+
+def delete_user_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> DeleteUserUseCase:
+    return DeleteUserUseCase(
+        user_repository=_user_repository(postgres_session),
+        router_repository=_router_repository(postgres_session),
         provider_repository=_provider_repository(postgres_session),
     )

api/domain/auth/__init__.py

@@ -1,3 +1,5 @@
-from ._authssosessionvalidator import AuthSsoSessionValidator, SsoSessionClaims
+from ._authssosessionvalidator import AuthSsoSessionValidator
+from ._ssopolicyrepository import SsoPolicyRepository
+from .entities import SsoPolicy
 
-__all__ = ["AuthSsoSessionValidator", "SsoSessionClaims"]
+__all__ = ["AuthSsoSessionValidator", "SsoPolicy", "SsoPolicyRepository"]

api/domain/auth/_authssosessionvalidator.py

@@ -1,16 +1,12 @@
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
 
-from api.domain.auth.errors import InvalidOidcTokenError, SsoProviderNotAvailableError
-
-
-@dataclass(frozen=True)
-class SsoSessionClaims:
-    email: str
-    user: str | None = None
+from api.domain.auth.errors import SsoInvalidSessionError, SsoProviderNotAvailableError
 
 
 class AuthSsoSessionValidator(ABC):
     @abstractmethod
-    async def validate_session(self, session_cookie: str) -> SsoSessionClaims | InvalidOidcTokenError | SsoProviderNotAvailableError:
+    async def validate_session(self, session_cookie: str) -> str | SsoInvalidSessionError | SsoProviderNotAvailableError:
+        """
+        Validate the session cookie by calling the authentication service and return the email of the user.
+        """
         pass

api/domain/auth/_ssopolicyrepository.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+
+from api.domain.auth.entities import NewSsoPolicy, SsoPolicy
+from api.domain.auth.errors import SsoPolicyRuleAlreadyExistsError
+from api.domain.organization.errors import OrganizationNotFoundError
+from api.domain.role.errors import RoleNotFoundError
+
+
+class SsoPolicyRepository(ABC):
+    @abstractmethod
+    async def get_policy(self) -> SsoPolicy:
+        pass
+
+    @abstractmethod
+    async def replace_policy(
+        self, policy: NewSsoPolicy
+    ) -> SsoPolicy | RoleNotFoundError | OrganizationNotFoundError | SsoPolicyRuleAlreadyExistsError:
+        pass

api/domain/auth/entities.py

@@ -1,6 +1,8 @@
+from datetime import datetime
 from enum import StrEnum
+from typing import Literal
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, Field
 
 
 class SsoAccessRuleType(StrEnum):
@@ -9,25 +11,91 @@ class SsoAccessRuleType(StrEnum):
     ROLE = "role"
 
 
-class SsoAccessRule(BaseModel):
-    model_config = ConfigDict(from_attributes=True)
+class NewSsoPolicyEmailRule(BaseModel):
+    type: Literal[SsoAccessRuleType.EMAIL]
+    value: str
+
+
+class NewSsoPolicyOrganizationRule(BaseModel):
+    type: Literal[SsoAccessRuleType.ORGANIZATION]
+    value: str | None = None
+    organization_id: int
+
+
+class NewSsoPolicyRoleRule(BaseModel):
+    type: Literal[SsoAccessRuleType.ROLE]
+    value: str | None = None
+    role_id: int
 
-    id: int | None = None
+
+type NewSsoPolicyRule = NewSsoPolicyEmailRule | NewSsoPolicyOrganizationRule | NewSsoPolicyRoleRule
+
+
+class NewSsoPolicy(BaseModel):
+    rules: list[NewSsoPolicyRule] = Field(default_factory=list)
+
+
+class SsoPolicyRuleBase(BaseModel):
+    id: int
     type: SsoAccessRuleType
+    created: datetime
+    updated: datetime
+
+
+class SsoPolicyEmailRule(SsoPolicyRuleBase):
+    type: Literal[SsoAccessRuleType.EMAIL]
     value: str
 
 
-class SsoRoleMapping(BaseModel):
-    model_config = ConfigDict(from_attributes=True)
+class SsoPolicyOrganizationRule(SsoPolicyRuleBase):
+    type: Literal[SsoAccessRuleType.ORGANIZATION]
+    value: str | None = None
+    organization_id: int
+
 
-    id: int | None = None
-    organization_name: str
-    oidc_role_name: str
+class SsoPolicyRoleRule(SsoPolicyRuleBase):
+    type: Literal[SsoAccessRuleType.ROLE]
+    value: str | None = None
     role_id: int
 
 
+type SsoPolicyRule = SsoPolicyEmailRule | SsoPolicyOrganizationRule | SsoPolicyRoleRule
+
+
 class SsoPolicy(BaseModel):
-    allowed_emails: list[str] = Field(default_factory=list)
-    allowed_organizations: list[str] = Field(default_factory=list)
-    allowed_roles: list[str] = Field(default_factory=list)
-    role_mappings: list[SsoRoleMapping] = Field(default_factory=list)
+    rules: list[SsoPolicyRule] = Field(default_factory=list)
+
+    def is_allowed(self, email: str, organization: str | None, roles: list[str]) -> bool:
+        checks: list[bool] = []
+        allowed_emails = [rule.value for rule in self.rules if rule.type == SsoAccessRuleType.EMAIL]
+        allowed_organizations = [rule.value for rule in self.rules if rule.type == SsoAccessRuleType.ORGANIZATION and rule.value is not None]
+        allowed_roles = [rule.value for rule in self.rules if rule.type == SsoAccessRuleType.ROLE and rule.value is not None]
+
+        if allowed_roles:
+            checks.append(any(role in allowed_roles for role in roles))
+
+        if allowed_emails:
+            checks.append(any(email.endswith(allowed_email) for allowed_email in allowed_emails))
+
+        if allowed_organizations:
+            checks.append(organization is not None and organization in allowed_organizations)
+
+        return True if not checks else any(checks)
+
+    def get_matching_organization_rule(self, organization: str | None) -> SsoPolicyOrganizationRule | None:
+        for rule in self.rules:
+            if rule.type != SsoAccessRuleType.ORGANIZATION:
+                continue
+            if organization is None and rule.value is None:  # default organization rule
+                return rule
+            if organization is not None and organization == rule.value:
+                return rule
+
+    def get_matching_role_rule(self, roles: list[str]) -> SsoPolicyRoleRule | None:
+        for rule in self.rules:
+            if rule.type != SsoAccessRuleType.ROLE:
+                continue
+            if not roles and rule.value is None:  # default role rule
+                return rule
+            if roles and rule.value in roles:
+                return rule