feat(auth): add allow_sign_up environment variable for OpenID Connect

src/phoenix/config.py

@@ -577,6 +577,7 @@ class OAuth2ClientConfig:
     client_id: str
     client_secret: str
     oidc_config_url: str
+    allow_sign_up: bool
 
     @classmethod
     def from_env(cls, idp_name: str) -> "OAuth2ClientConfig":
@@ -608,6 +609,7 @@ def from_env(cls, idp_name: str) -> "OAuth2ClientConfig":
                 f"An OpenID Connect configuration URL must be set for the {idp_name} OAuth2 IDP "
                 f"via the {oidc_config_url_env_var} environment variable"
             )
+        allow_sign_up = get_env_oauth2_allow_sign_up(idp_name)
         parsed_oidc_config_url = urlparse(oidc_config_url)
         is_local_oidc_config_url = parsed_oidc_config_url.hostname in ("localhost", "127.0.0.1")
         if parsed_oidc_config_url.scheme != "https" and not is_local_oidc_config_url:
@@ -624,24 +626,79 @@ def from_env(cls, idp_name: str) -> "OAuth2ClientConfig":
             client_id=client_id,
             client_secret=client_secret,
             oidc_config_url=oidc_config_url,
+            allow_sign_up=allow_sign_up,
         )
 
 
 def get_env_oauth2_settings() -> list[OAuth2ClientConfig]:
     """
-    Get OAuth2 settings from environment variables.
-    """
+    Retrieves and validates OAuth2/OpenID Connect (OIDC) identity provider configurations from environment variables.
+
+    This function scans the environment for OAuth2 configuration variables and returns a list of
+    configured identity providers. It supports multiple identity providers simultaneously.
+
+    Environment Variable Pattern:
+        PHOENIX_OAUTH2_{IDP_NAME}_{CONFIG_TYPE}
+
+    Required Environment Variables for each IDP:
+        - PHOENIX_OAUTH2_{IDP_NAME}_CLIENT_ID: The OAuth2 client ID issued by the identity provider
+        - PHOENIX_OAUTH2_{IDP_NAME}_CLIENT_SECRET: The OAuth2 client secret issued by the identity provider
+        - PHOENIX_OAUTH2_{IDP_NAME}_OIDC_CONFIG_URL: The OpenID Connect configuration URL (must be HTTPS)
 
+    Optional Environment Variables:
+        - PHOENIX_OAUTH2_{IDP_NAME}_DISPLAY_NAME: A user-friendly name for the identity provider
+        - PHOENIX_OAUTH2_{IDP_NAME}_ALLOW_SIGN_UP: Whether to allow new user registration (defaults to True)
+        When set to False, the system will check if the user exists in the database by their email address.
+        If the user does not exist or has a password set, they will be redirected to the login page with
+        an error message.
+
+    Returns:
+        list[OAuth2ClientConfig]: A list of configured OAuth2 identity providers, sorted alphabetically by IDP name.
+            Each OAuth2ClientConfig contains the validated configuration for one identity provider.
+
+    Raises:
+        ValueError: If required environment variables are missing or invalid.
+            Specifically, if the OIDC configuration URL is not HTTPS (except for localhost).
+
+    Example:
+        To configure Google as an identity provider, set these environment variables:
+        PHOENIX_OAUTH2_GOOGLE_CLIENT_ID=your_client_id
+        PHOENIX_OAUTH2_GOOGLE_CLIENT_SECRET=your_client_secret
+        PHOENIX_OAUTH2_GOOGLE_OIDC_CONFIG_URL=https://accounts.google.com/.well-known/openid-configuration
+        PHOENIX_OAUTH2_GOOGLE_DISPLAY_NAME=Google (optional)
+        PHOENIX_OAUTH2_GOOGLE_ALLOW_SIGN_UP=true (optional, defaults to true)
+    """  # noqa: E501
     idp_names = set()
     pattern = re.compile(
-        r"^PHOENIX_OAUTH2_(\w+)_(DISPLAY_NAME|CLIENT_ID|CLIENT_SECRET|OIDC_CONFIG_URL)$"
+        r"^PHOENIX_OAUTH2_(\w+)_(DISPLAY_NAME|CLIENT_ID|CLIENT_SECRET|OIDC_CONFIG_URL|ALLOW_SIGN_UP)$"
     )
     for env_var in os.environ:
         if (match := pattern.match(env_var)) is not None and (idp_name := match.group(1).lower()):
             idp_names.add(idp_name)
     return [OAuth2ClientConfig.from_env(idp_name) for idp_name in sorted(idp_names)]
 
 
+def get_env_oauth2_allow_sign_up(idp_name: str) -> bool:
+    """Retrieves the allow_sign_up setting for a specific OAuth2 identity provider.
+
+    This function determines whether new user registration is allowed for the specified identity provider.
+    When set to False, the system will check if the user exists in the database by their email address.
+    If the user does not exist or has a password set, they will be redirected to the login page with
+    an error message.
+
+    Parameters:
+        idp_name (str): The name of the identity provider (e.g., 'google', 'aws_cognito', 'microsoft_entra_id')
+
+    Returns:
+        bool: True if new user registration is allowed (default), False otherwise
+
+    Environment Variable:
+        PHOENIX_OAUTH2_{IDP_NAME}_ALLOW_SIGN_UP: Controls whether new user registration is allowed (defaults to True if not set)
+    """  # noqa: E501
+    env_var = f"PHOENIX_OAUTH2_{idp_name}_ALLOW_SIGN_UP".upper()
+    return _bool_val(env_var, True)
+
+
 PHOENIX_DIR = Path(__file__).resolve().parent
 # Server config
 SERVER_DIR = PHOENIX_DIR / "server"

src/phoenix/db/constants.py

@@ -0,0 +1,2 @@
+TBD_OAUTH2_CLIENT_ID_PREFIX = "TBD_OAUTH2_CLIENT_ID_"
+TBD_OAUTH2_USER_ID_PREFIX = "TBD_OAUTH2_USER_ID_"

src/phoenix/server/api/mutations/user_mutations.py

@@ -24,6 +24,7 @@
     validate_password_format,
 )
 from phoenix.db import enums, models
+from phoenix.db.constants import TBD_OAUTH2_CLIENT_ID_PREFIX, TBD_OAUTH2_USER_ID_PREFIX
 from phoenix.server.api.auth import IsAdmin, IsLocked, IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import Conflict, NotFound, Unauthorized
@@ -93,16 +94,25 @@ async def create_user(
         input: CreateUserInput,
     ) -> UserMutationPayload:
         validate_email_format(email := input.email)
-        validate_password_format(password := input.password)
-        salt = secrets.token_bytes(DEFAULT_SECRET_LENGTH)
-        password_hash = await info.context.hash_password(password, salt)
-        user = models.User(
-            reset_password=True,
-            username=input.username,
-            email=email,
-            password_hash=password_hash,
-            password_salt=salt,
-        )
+        if input.password:
+            validate_password_format(password := input.password)
+            salt = secrets.token_bytes(DEFAULT_SECRET_LENGTH)
+            password_hash = await info.context.hash_password(password, salt)
+            user = models.User(
+                reset_password=True,
+                username=input.username,
+                email=email,
+                password_hash=password_hash,
+                password_salt=salt,
+            )
+        else:
+            user = models.User(
+                reset_password=False,
+                username=input.username,
+                email=email,
+                oauth2_client_id=f"{TBD_OAUTH2_CLIENT_ID_PREFIX}{secrets.token_hex(8)}",
+                oauth2_user_id=f"{TBD_OAUTH2_USER_ID_PREFIX}{secrets.token_hex(8)}",
+            )
         async with AsyncExitStack() as stack:
             session = await stack.enter_async_context(info.context.db())
             user_role_id = await session.scalar(_select_role_id_by_name(input.role.value))

src/phoenix/server/api/routers/oauth2.py

@@ -34,6 +34,7 @@
 )
 from phoenix.config import get_env_disable_rate_limit
 from phoenix.db import models
+from phoenix.db.constants import TBD_OAUTH2_CLIENT_ID_PREFIX, TBD_OAUTH2_USER_ID_PREFIX
 from phoenix.db.enums import UserRole
 from phoenix.server.bearer_auth import create_access_and_refresh_tokens
 from phoenix.server.oauth2 import OAuth2Client
@@ -169,12 +170,13 @@ async def create_tokens(
     user_info = _parse_user_info(user_info)
     try:
         async with request.app.state.db() as session:
-            user = await _ensure_user_exists_and_is_up_to_date(
+            user = await _process_oauth2_user(
                 session,
                 oauth2_client_id=str(oauth2_client.client_id),
                 user_info=user_info,
+                allow_sign_up=oauth2_client.allow_sign_up,
             )
-    except EmailAlreadyInUse as error:
+    except (EmailAlreadyInUse, SignInNotAllowed) as error:
         return _redirect_to_login(request=request, error=str(error))
     access_token, refresh_token = await create_access_and_refresh_tokens(
         user=user,
@@ -198,13 +200,21 @@ async def create_tokens(
     return response
 
 
-@dataclass
+@dataclass(frozen=True)
 class UserInfo:
     idp_user_id: str
     email: str
     username: Optional[str]
     profile_picture_url: Optional[str]
 
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "idp_user_id", self.idp_user_id.strip())
+        object.__setattr__(self, "email", self.email.strip())
+        if username := self.username:
+            object.__setattr__(self, "username", username.strip())
+        if profile_picture_url := self.profile_picture_url:
+            object.__setattr__(self, "profile_picture_url", profile_picture_url.strip())
+
 
 def _validate_token_data(token_data: dict[str, Any]) -> None:
     """
@@ -235,18 +245,135 @@ def _parse_user_info(user_info: dict[str, Any]) -> UserInfo:
     )
 
 
-async def _ensure_user_exists_and_is_up_to_date(
-    session: AsyncSession, /, *, oauth2_client_id: str, user_info: UserInfo
+async def _process_oauth2_user(
+    session: AsyncSession,
+    /,
+    *,
+    oauth2_client_id: str,
+    user_info: UserInfo,
+    allow_sign_up: bool,
 ) -> models.User:
+    """
+    Processes an OAuth2 user, either signing in an existing user or creating/updating one.
+
+    This function handles two main scenarios based on the allow_sign_up parameter:
+    1. When sign-up is not allowed (allow_sign_up=False):
+       - Checks if the user exists and can sign in with the given OAuth2 credentials
+       - Updates placeholder OAuth2 credentials if needed (e.g., temporary IDs)
+       - If the user doesn't exist or has a password set, raises SignInNotAllowed
+    2. When sign-up is allowed (allow_sign_up=True):
+       - Finds the user by OAuth2 credentials (client_id and user_id)
+       - Creates a new user if one doesn't exist, with default member role
+       - Updates the user's email if it has changed
+       - Handles username conflicts by adding a random suffix if needed
+
+    The allow_sign_up parameter is typically controlled by the PHOENIX_OAUTH2_{IDP_NAME}_ALLOW_SIGN_UP
+    environment variable for the specific identity provider.
+
+    Args:
+        session: The database session
+        oauth2_client_id: The ID of the OAuth2 client
+        user_info: User information from the OAuth2 provider
+        allow_sign_up: Whether to allow creating new users
+
+    Returns:
+        The user object
+
+    Raises:
+        SignInNotAllowed: When sign-in is not allowed for the user (user doesn't exist or has a password)
+        EmailAlreadyInUse: When the email is already in use by another account
+    """  # noqa: E501
+    if not allow_sign_up:
+        return await _sign_in_existing_oauth2_user(
+            session,
+            oauth2_client_id=oauth2_client_id,
+            user_info=user_info,
+        )
+    return await _create_or_update_user(
+        session,
+        oauth2_client_id=oauth2_client_id,
+        user_info=user_info,
+    )
+
+
+async def _sign_in_existing_oauth2_user(
+    session: AsyncSession,
+    /,
+    *,
+    oauth2_client_id: str,
+    user_info: UserInfo,
+) -> models.User:
+    """
+    Signs in an existing user with OAuth2 credentials.
+
+    Args:
+        session: The database session
+        oauth2_client_id: The ID of the OAuth2 client
+        user_info: User information from the OAuth2 provider
+
+    Returns:
+        The signed-in user
+
+    Raises:
+        SignInNotAllowed: When sign-in is not allowed for the user
+    """
+    email = user_info.email
+    stmt = select(models.User).filter_by(email=email).options(joinedload(models.User.role))
+    user = await session.scalar(stmt)
+    if (
+        user is None
+        or user.password_hash is not None
+        or user.oauth2_client_id is None
+        or user.oauth2_user_id is None
+        or (
+            user_info.idp_user_id != user.oauth2_user_id
+            and not user.oauth2_user_id.startswith(TBD_OAUTH2_USER_ID_PREFIX)
+        )
+        or (
+            oauth2_client_id != user.oauth2_client_id
+            and not user.oauth2_client_id.startswith(TBD_OAUTH2_CLIENT_ID_PREFIX)
+        )
+    ):
+        raise SignInNotAllowed(f"Sign in is not allowed for {email}.")
+    if user.oauth2_client_id.startswith(TBD_OAUTH2_CLIENT_ID_PREFIX):
+        user.oauth2_client_id = oauth2_client_id
+    if user.oauth2_user_id.startswith(TBD_OAUTH2_USER_ID_PREFIX):
+        user.oauth2_user_id = user_info.idp_user_id
+    if user in session.dirty:
+        await session.flush()
+    return user
+
+
+async def _create_or_update_user(
+    session: AsyncSession,
+    /,
+    *,
+    oauth2_client_id: str,
+    user_info: UserInfo,
+) -> models.User:
+    """
+    Creates a new user or updates an existing one with OAuth2 credentials.
+
+    Args:
+        session: The database session
+        oauth2_client_id: The ID of the OAuth2 client
+        user_info: User information from the OAuth2 provider
+
+    Returns:
+        The created or updated user
+
+    Raises:
+        EmailAlreadyInUse: When the email is already in use by another account
+    """
     user = await _get_user(
         session,
         oauth2_client_id=oauth2_client_id,
         idp_user_id=user_info.idp_user_id,
     )
     if user is None:
-        user = await _create_user(session, oauth2_client_id=oauth2_client_id, user_info=user_info)
-    elif user.email != user_info.email:
-        user = await _update_user_email(session, user_id=user.id, email=user_info.email)
+        return await _create_user(session, oauth2_client_id=oauth2_client_id, user_info=user_info)
+    if user.email != user_info.email:
+        return await _update_user_email(session, user_id=user.id, email=user_info.email)
     return user
 
 
@@ -366,6 +493,10 @@ class EmailAlreadyInUse(Exception):
     pass
 
 
+class SignInNotAllowed(Exception):
+    pass
+
+
 def _redirect_to_login(*, request: Request, error: str) -> RedirectResponse:
     """
     Creates a RedirectResponse to the login page to display an error message.

src/phoenix/server/oauth2.py

@@ -19,8 +19,13 @@ class OAuth2Client(AsyncOAuth2Mixin, AsyncOpenIDMixin, BaseApp):  # type:ignore[
 
     client_cls = AsyncHttpxOAuth2Client
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
+    def __init__(self, *args: Any, allow_sign_up: bool, **kwargs: Any) -> None:
         super().__init__(framework=None, *args, **kwargs)
+        self._allow_sign_up = allow_sign_up
+
+    @property
+    def allow_sign_up(self) -> bool:
+        return self._allow_sign_up
 
 
 class OAuth2Clients:
@@ -35,6 +40,7 @@ def add_client(self, config: OAuth2ClientConfig) -> None:
             client_secret=config.client_secret,
             server_metadata_url=config.oidc_config_url,
             client_kwargs={"scope": "openid email profile"},
+            allow_sign_up=config.allow_sign_up,
         )
         assert isinstance(client, OAuth2Client)
         self._clients[config.idp_name] = client

tests/integration/auth/conftest.py

@@ -71,6 +71,19 @@ def _app(
             f"PHOENIX_OAUTH2_{_oidc_server}_OIDC_CONFIG_URL".upper(),
             f"{_oidc_server.base_url}/.well-known/openid-configuration",
         ),
+        (
+            f"PHOENIX_OAUTH2_{_oidc_server}_NO_SIGN_UP_CLIENT_ID".upper(),
+            _oidc_server.client_id,
+        ),
+        (
+            f"PHOENIX_OAUTH2_{_oidc_server}_NO_SIGN_UP_CLIENT_SECRET".upper(),
+            _oidc_server.client_secret,
+        ),
+        (
+            f"PHOENIX_OAUTH2_{_oidc_server}_NO_SIGN_UP_OIDC_CONFIG_URL".upper(),
+            f"{_oidc_server.base_url}/.well-known/openid-configuration",
+        ),
+        (f"PHOENIX_OAUTH2_{_oidc_server}_NO_SIGN_UP_ALLOW_SIGN_UP".upper(), "false"),
     )
     with ExitStack() as stack:
         stack.enter_context(mock.patch.dict(os.environ, values))