arize-phoenix 4.36.0__py3-none-any.whl → 5.1.0__py3-none-any.whl

phoenix/server/api/routers/oauth2.py

@@ -0,0 +1,456 @@
+import re
+from dataclasses import dataclass
+from datetime import timedelta
+from random import randrange
+from typing import Any, Dict, Optional, Tuple, TypedDict
+from urllib.parse import unquote, urlparse
+
+from authlib.common.security import generate_token
+from authlib.integrations.starlette_client import OAuthError
+from authlib.jose import jwt
+from authlib.jose.errors import JoseError
+from fastapi import APIRouter, Cookie, Depends, Path, Query, Request
+from sqlalchemy import Boolean, and_, case, cast, func, insert, or_, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
+from sqlean.dbapi2 import IntegrityError  # type: ignore[import-untyped]
+from starlette.datastructures import URL
+from starlette.responses import RedirectResponse
+from starlette.routing import Router
+from starlette.status import HTTP_302_FOUND
+from typing_extensions import Annotated, NotRequired, TypeGuard
+
+from phoenix.auth import (
+    DEFAULT_OAUTH2_LOGIN_EXPIRY_MINUTES,
+    PHOENIX_OAUTH2_NONCE_COOKIE_NAME,
+    PHOENIX_OAUTH2_STATE_COOKIE_NAME,
+    delete_oauth2_nonce_cookie,
+    delete_oauth2_state_cookie,
+    set_access_token_cookie,
+    set_oauth2_nonce_cookie,
+    set_oauth2_state_cookie,
+    set_refresh_token_cookie,
+)
+from phoenix.config import get_env_disable_rate_limit
+from phoenix.db import models
+from phoenix.db.enums import UserRole
+from phoenix.server.bearer_auth import create_access_and_refresh_tokens
+from phoenix.server.oauth2 import OAuth2Client
+from phoenix.server.rate_limiters import (
+    ServerRateLimiter,
+    fastapi_ip_rate_limiter,
+    fastapi_route_rate_limiter,
+)
+from phoenix.server.types import TokenStore
+
+_LOWERCASE_ALPHANUMS_AND_UNDERSCORES = r"[a-z0-9_]+"
+
+login_rate_limiter = fastapi_ip_rate_limiter(
+    ServerRateLimiter(
+        per_second_rate_limit=0.2,
+        enforcement_window_seconds=30,
+        partition_seconds=60,
+        active_partitions=2,
+    ),
+)
+
+create_tokens_rate_limiter = fastapi_route_rate_limiter(
+    ServerRateLimiter(
+        per_second_rate_limit=0.5,
+        enforcement_window_seconds=30,
+        partition_seconds=60,
+        active_partitions=2,
+    )
+)
+
+router = APIRouter(
+    prefix="/oauth2",
+    include_in_schema=False,
+)
+
+if not get_env_disable_rate_limit():
+    login_dependencies = [Depends(login_rate_limiter)]
+    create_tokens_dependencies = [Depends(create_tokens_rate_limiter)]
+else:
+    login_dependencies = []
+    create_tokens_dependencies = []
+
+
+@router.post("/{idp_name}/login", dependencies=login_dependencies)
+async def login(
+    request: Request,
+    idp_name: Annotated[str, Path(min_length=1, pattern=_LOWERCASE_ALPHANUMS_AND_UNDERSCORES)],
+    return_url: Optional[str] = Query(default=None, alias="returnUrl"),
+) -> RedirectResponse:
+    secret = request.app.state.get_secret()
+    if not isinstance(
+        oauth2_client := request.app.state.oauth2_clients.get_client(idp_name), OAuth2Client
+    ):
+        return _redirect_to_login(error=f"Unknown IDP: {idp_name}.")
+    origin_url = _get_origin_url(request)
+    authorization_url_data = await oauth2_client.create_authorization_url(
+        redirect_uri=_get_create_tokens_endpoint(
+            request=request, origin_url=origin_url, idp_name=idp_name
+        ),
+        state=_generate_state_for_oauth2_authorization_code_flow(
+            secret=secret, origin_url=origin_url, return_url=return_url
+        ),
+    )
+    assert isinstance(authorization_url := authorization_url_data.get("url"), str)
+    assert isinstance(state := authorization_url_data.get("state"), str)
+    assert isinstance(nonce := authorization_url_data.get("nonce"), str)
+    response = RedirectResponse(url=authorization_url, status_code=HTTP_302_FOUND)
+    response = set_oauth2_state_cookie(
+        response=response,
+        state=state,
+        max_age=timedelta(minutes=DEFAULT_OAUTH2_LOGIN_EXPIRY_MINUTES),
+    )
+    response = set_oauth2_nonce_cookie(
+        response=response,
+        nonce=nonce,
+        max_age=timedelta(minutes=DEFAULT_OAUTH2_LOGIN_EXPIRY_MINUTES),
+    )
+    return response
+
+
+@router.get("/{idp_name}/tokens", dependencies=create_tokens_dependencies)
+async def create_tokens(
+    request: Request,
+    idp_name: Annotated[str, Path(min_length=1, pattern=_LOWERCASE_ALPHANUMS_AND_UNDERSCORES)],
+    state: str = Query(),
+    authorization_code: str = Query(alias="code"),
+    stored_state: str = Cookie(alias=PHOENIX_OAUTH2_STATE_COOKIE_NAME),
+    stored_nonce: str = Cookie(alias=PHOENIX_OAUTH2_NONCE_COOKIE_NAME),
+) -> RedirectResponse:
+    secret = request.app.state.get_secret()
+    if state != stored_state:
+        return _redirect_to_login(error=_INVALID_OAUTH2_STATE_MESSAGE)
+    try:
+        payload = _parse_state_payload(secret=secret, state=state)
+    except JoseError:
+        return _redirect_to_login(error=_INVALID_OAUTH2_STATE_MESSAGE)
+    if (return_url := payload.get("return_url")) is not None and not _is_relative_url(
+        unquote(return_url)
+    ):
+        return _redirect_to_login(error="Attempting login with unsafe return URL.")
+    assert isinstance(access_token_expiry := request.app.state.access_token_expiry, timedelta)
+    assert isinstance(refresh_token_expiry := request.app.state.refresh_token_expiry, timedelta)
+    token_store: TokenStore = request.app.state.get_token_store()
+    if not isinstance(
+        oauth2_client := request.app.state.oauth2_clients.get_client(idp_name), OAuth2Client
+    ):
+        return _redirect_to_login(error=f"Unknown IDP: {idp_name}.")
+    try:
+        token_data = await oauth2_client.fetch_access_token(
+            state=state,
+            code=authorization_code,
+            redirect_uri=_get_create_tokens_endpoint(
+                request=request, origin_url=payload["origin_url"], idp_name=idp_name
+            ),
+        )
+    except OAuthError as error:
+        return _redirect_to_login(error=str(error))
+    _validate_token_data(token_data)
+    if "id_token" not in token_data:
+        return _redirect_to_login(
+            error=f"OAuth2 IDP {idp_name} does not appear to support OpenID Connect."
+        )
+    user_info = await oauth2_client.parse_id_token(token_data, nonce=stored_nonce)
+    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(
+                session,
+                oauth2_client_id=str(oauth2_client.client_id),
+                user_info=user_info,
+            )
+    except EmailAlreadyInUse as error:
+        return _redirect_to_login(error=str(error))
+    access_token, refresh_token = await create_access_and_refresh_tokens(
+        user=user,
+        token_store=token_store,
+        access_token_expiry=access_token_expiry,
+        refresh_token_expiry=refresh_token_expiry,
+    )
+    response = RedirectResponse(url=return_url or "/", status_code=HTTP_302_FOUND)
+    response = set_access_token_cookie(
+        response=response, access_token=access_token, max_age=access_token_expiry
+    )
+    response = set_refresh_token_cookie(
+        response=response, refresh_token=refresh_token, max_age=refresh_token_expiry
+    )
+    response = delete_oauth2_state_cookie(response)
+    response = delete_oauth2_nonce_cookie(response)
+    return response
+
+
+@dataclass
+class UserInfo:
+    idp_user_id: str
+    email: str
+    username: Optional[str]
+    profile_picture_url: Optional[str]
+
+
+def _validate_token_data(token_data: Dict[str, Any]) -> None:
+    """
+    Performs basic validations on the token data returned by the IDP.
+    """
+    assert isinstance(token_data.get("access_token"), str)
+    assert isinstance(token_type := token_data.get("token_type"), str)
+    assert token_type.lower() == "bearer"
+
+
+def _parse_user_info(user_info: Dict[str, Any]) -> UserInfo:
+    """
+    Parses user info from the IDP's ID token.
+    """
+    assert isinstance(subject := user_info.get("sub"), (str, int))
+    idp_user_id = str(subject)
+    assert isinstance(email := user_info.get("email"), str)
+    assert isinstance(username := user_info.get("name"), str) or username is None
+    assert (
+        isinstance(profile_picture_url := user_info.get("picture"), str)
+        or profile_picture_url is None
+    )
+    return UserInfo(
+        idp_user_id=idp_user_id,
+        email=email,
+        username=username,
+        profile_picture_url=profile_picture_url,
+    )
+
+
+async def _ensure_user_exists_and_is_up_to_date(
+    session: AsyncSession, /, *, oauth2_client_id: str, user_info: UserInfo
+) -> models.User:
+    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 user
+
+
+async def _get_user(
+    session: AsyncSession, /, *, oauth2_client_id: str, idp_user_id: str
+) -> Optional[models.User]:
+    """
+    Retrieves the user uniquely identified by the given OAuth2 client ID and IDP
+    user ID.
+    """
+    user = await session.scalar(
+        select(models.User)
+        .where(
+            and_(
+                models.User.oauth2_client_id == oauth2_client_id,
+                models.User.oauth2_user_id == idp_user_id,
+            )
+        )
+        .options(joinedload(models.User.role))
+    )
+    return user
+
+
+async def _create_user(
+    session: AsyncSession,
+    /,
+    *,
+    oauth2_client_id: str,
+    user_info: UserInfo,
+) -> models.User:
+    """
+    Creates a new user with the user info from the IDP.
+    """
+    email_exists, username_exists = await _email_and_username_exist(
+        session,
+        email=(email := user_info.email),
+        username=(username := user_info.username),
+    )
+    if email_exists:
+        raise EmailAlreadyInUse(f"An account for {email} is already in use.")
+    member_role_id = (
+        select(models.UserRole.id)
+        .where(models.UserRole.name == UserRole.MEMBER.value)
+        .scalar_subquery()
+    )
+    user_id = await session.scalar(
+        insert(models.User)
+        .returning(models.User.id)
+        .values(
+            user_role_id=member_role_id,
+            oauth2_client_id=oauth2_client_id,
+            oauth2_user_id=user_info.idp_user_id,
+            username=_with_random_suffix(username) if username and username_exists else username,
+            email=email,
+            profile_picture_url=user_info.profile_picture_url,
+            reset_password=False,
+        )
+    )
+    assert isinstance(user_id, int)
+    user = await session.scalar(
+        select(models.User).where(models.User.id == user_id).options(joinedload(models.User.role))
+    )  # query user again for joined load
+    assert isinstance(user, models.User)
+    return user
+
+
+async def _update_user_email(session: AsyncSession, /, *, user_id: int, email: str) -> models.User:
+    """
+    Updates an existing user's email.
+    """
+    try:
+        await session.execute(
+            update(models.User)
+            .where(models.User.id == user_id)
+            .values(email=email)
+            .options(joinedload(models.User.role))
+        )
+    except IntegrityError:
+        raise EmailAlreadyInUse(f"An account for {email} is already in use.")
+    user = await session.scalar(
+        select(models.User).where(models.User.id == user_id).options(joinedload(models.User.role))
+    )  # query user again for joined load
+    assert isinstance(user, models.User)
+    return user
+
+
+async def _email_and_username_exist(
+    session: AsyncSession, /, *, email: str, username: Optional[str]
+) -> Tuple[bool, bool]:
+    """
+    Checks whether the email and username are already in use.
+    """
+    [(email_exists, username_exists)] = (
+        await session.execute(
+            select(
+                cast(
+                    func.coalesce(
+                        func.max(case((models.User.email == email, 1), else_=0)),
+                        0,
+                    ),
+                    Boolean,
+                ).label("email_exists"),
+                cast(
+                    func.coalesce(
+                        func.max(case((models.User.username == username, 1), else_=0)),
+                        0,
+                    ),
+                    Boolean,
+                ).label("username_exists"),
+            ).where(or_(models.User.email == email, models.User.username == username))
+        )
+    ).all()
+    return email_exists, username_exists
+
+
+class EmailAlreadyInUse(Exception):
+    pass
+
+
+def _redirect_to_login(*, error: str) -> RedirectResponse:
+    """
+    Creates a RedirectResponse to the login page to display an error message.
+    """
+    url = URL("/login").include_query_params(error=error)
+    response = RedirectResponse(url=url)
+    response = delete_oauth2_state_cookie(response)
+    response = delete_oauth2_nonce_cookie(response)
+    return response
+
+
+def _get_create_tokens_endpoint(*, request: Request, origin_url: str, idp_name: str) -> str:
+    """
+    Gets the endpoint for create tokens route.
+    """
+    router: Router = request.scope["router"]
+    url_path = router.url_path_for(create_tokens.__name__, idp_name=idp_name)
+    return str(url_path.make_absolute_url(base_url=origin_url))
+
+
+def _generate_state_for_oauth2_authorization_code_flow(
+    *, secret: str, origin_url: str, return_url: Optional[str]
+) -> str:
+    """
+    Generates a JWT whose payload contains both an OAuth2 state (generated using
+    the `authlib` default algorithm) and a return URL. This allows us to pass
+    the return URL to the OAuth2 authorization server via the `state` query
+    parameter and have it returned to us in the callback without needing to
+    maintain state.
+    """
+    header = {"alg": _JWT_ALGORITHM}
+    payload = _OAuth2StatePayload(
+        random=generate_token(),
+        origin_url=origin_url,
+    )
+    if return_url is not None:
+        payload["return_url"] = return_url
+    jwt_bytes: bytes = jwt.encode(header=header, payload=payload, key=secret)
+    return jwt_bytes.decode()
+
+
+class _OAuth2StatePayload(TypedDict):
+    """
+    Represents the OAuth2 state payload.
+    """
+
+    random: str
+    origin_url: str
+    return_url: NotRequired[str]
+
+
+def _parse_state_payload(*, secret: str, state: str) -> _OAuth2StatePayload:
+    """
+    Validates the JWT signature and parses the return URL from the OAuth2 state.
+    """
+    payload = jwt.decode(s=state, key=secret)
+    if _is_oauth2_state_payload(payload):
+        return payload
+    raise ValueError("Invalid OAuth2 state payload.")
+
+
+def _is_relative_url(url: str) -> bool:
+    """
+    Determines whether the URL is relative.
+    """
+    return bool(_RELATIVE_URL_PATTERN.match(url))
+
+
+def _with_random_suffix(string: str) -> str:
+    """
+    Appends a random suffix.
+    """
+    return f"{string}-{randrange(10_000, 100_000)}"
+
+
+def _get_origin_url(request: Request) -> str:
+    """
+    Infers the origin URL from the request.
+    """
+    if (referer := request.headers.get("referer")) is None:
+        return str(request.base_url)
+    parsed_url = urlparse(referer)
+    return f"{parsed_url.scheme}://{parsed_url.netloc}"
+
+
+def _is_oauth2_state_payload(maybe_state_payload: Any) -> TypeGuard[_OAuth2StatePayload]:
+    """
+    Determines whether the given object is an OAuth2 state payload.
+    """
+
+    return (
+        isinstance(maybe_state_payload, dict)
+        and {"random", "origin_url"}.issubset((keys := set(maybe_state_payload.keys())))
+        and keys.issubset({"random", "origin_url", "return_url"})
+    )
+
+
+_JWT_ALGORITHM = "HS256"
+_INVALID_OAUTH2_STATE_MESSAGE = (
+    "Received invalid state parameter during OAuth2 authorization code flow for IDP {idp_name}."
+)
+_RELATIVE_URL_PATTERN = re.compile(r"^/($|\w)")

phoenix/server/api/routers/v1/__init__.py

@@ -1,6 +1,9 @@
 from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi.security import APIKeyHeader
 from starlette.status import HTTP_403_FORBIDDEN
 
+from phoenix.server.bearer_auth import is_authenticated
+
 from .datasets import router as datasets_router
 from .evaluations import router as evaluations_router
 from .experiment_evaluations import router as experiment_evaluations_router
@@ -24,19 +27,38 @@ async def prevent_access_in_read_only_mode(request: Request) -> None:
         )
 
 
-router = APIRouter(
-    prefix="/v1",
-    dependencies=[Depends(prevent_access_in_read_only_mode)],
-    responses=add_errors_to_responses(
-        [
-            HTTP_403_FORBIDDEN  # adds a 403 response to each route in the generated OpenAPI schema
-        ]
-    ),
-)
-router.include_router(datasets_router)
-router.include_router(experiments_router)
-router.include_router(experiment_runs_router)
-router.include_router(experiment_evaluations_router)
-router.include_router(traces_router)
-router.include_router(spans_router)
-router.include_router(evaluations_router)
+def create_v1_router(authentication_enabled: bool) -> APIRouter:
+    """
+    Instantiates the v1 REST API router.
+    """
+    dependencies = [Depends(prevent_access_in_read_only_mode)]
+    if authentication_enabled:
+        dependencies.append(
+            Depends(
+                APIKeyHeader(
+                    name="Authorization",
+                    scheme_name="Bearer",
+                    auto_error=False,
+                    description="Enter `Bearer` followed by a space and then the token.",
+                )
+            )
+        )
+        dependencies.append(Depends(is_authenticated))
+
+    router = APIRouter(
+        prefix="/v1",
+        dependencies=dependencies,
+        responses=add_errors_to_responses(
+            [
+                HTTP_403_FORBIDDEN  # adds a 403 response to routes in the generated OpenAPI schema
+            ]
+        ),
+    )
+    router.include_router(datasets_router)
+    router.include_router(experiments_router)
+    router.include_router(experiment_runs_router)
+    router.include_router(experiment_evaluations_router)
+    router.include_router(traces_router)
+    router.include_router(spans_router)
+    router.include_router(evaluations_router)
+    return router

phoenix/server/api/types/ApiKey.py

@@ -3,6 +3,8 @@ from typing import Optional
 
 import strawberry
 
+from phoenix.db.models import ApiKey as ORMApiKey
+
 
 @strawberry.interface
 class ApiKey:
@@ -14,3 +16,12 @@ class ApiKey:
     expires_at: Optional[datetime] = strawberry.field(
         description="The date and time the API key will expire."
     )
+
+
+def to_gql_api_key(api_key: ORMApiKey) -> ApiKey:
+    return ApiKey(
+        name=api_key.name,
+        description=api_key.description,
+        created_at=api_key.created_at,
+        expires_at=api_key.expires_at,
+    )

phoenix/server/api/types/AuthMethod.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+import strawberry
+
+
+@strawberry.enum
+class AuthMethod(Enum):
+    LOCAL = "LOCAL"
+    OAUTH2 = "OAUTH2"

phoenix/server/api/types/User.py

@@ -1,16 +1,60 @@
 from datetime import datetime
-from typing import Optional
+from typing import List, Optional
 
 import strawberry
+from sqlalchemy import select
+from strawberry import Private
 from strawberry.relay import Node, NodeID
+from strawberry.types import Info
 
-from .UserRole import UserRole
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import NotFound
+from phoenix.server.api.types.AuthMethod import AuthMethod
+from phoenix.server.api.types.UserApiKey import UserApiKey, to_gql_api_key
+
+from .UserRole import UserRole, to_gql_user_role
 
 
 @strawberry.type
 class User(Node):
     id_attr: NodeID[int]
+    password_needs_reset: bool
     email: str
-    username: Optional[str]
+    username: str
+    profile_picture_url: Optional[str]
     created_at: datetime
-    role: UserRole
+    user_role_id: Private[int]
+    auth_method: AuthMethod
+
+    @strawberry.field
+    async def role(self, info: Info[Context, None]) -> UserRole:
+        role = await info.context.data_loaders.user_roles.load(self.user_role_id)
+        if role is None:
+            raise NotFound(f"User role with id {self.user_role_id} not found")
+        return to_gql_user_role(role)
+
+    @strawberry.field
+    async def api_keys(self, info: Info[Context, None]) -> List[UserApiKey]:
+        async with info.context.db() as session:
+            api_keys = await session.scalars(
+                select(models.ApiKey).where(models.ApiKey.user_id == self.id_attr)
+            )
+        return [to_gql_api_key(api_key) for api_key in api_keys]
+
+
+def to_gql_user(user: models.User, api_keys: Optional[List[models.ApiKey]] = None) -> User:
+    """
+    Converts an ORM user to a GraphQL user.
+    """
+    assert user.auth_method is not None
+    return User(
+        id_attr=user.id,
+        password_needs_reset=user.reset_password,
+        username=user.username,
+        email=user.email,
+        profile_picture_url=user.profile_picture_url,
+        created_at=user.created_at,
+        user_role_id=user.user_role_id,
+        auth_method=AuthMethod(user.auth_method),
+    )

phoenix/server/api/types/UserApiKey.py

@@ -1,11 +1,45 @@
+from typing import TYPE_CHECKING
+
 import strawberry
 from strawberry import Private
-from strawberry.relay.types import Node, NodeID
+from strawberry.relay import Node, NodeID
+from strawberry.types import Info
+from typing_extensions import Annotated
+
+from phoenix.db.models import ApiKey as OrmApiKey
+from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import NotFound
 
 from .ApiKey import ApiKey
 
+if TYPE_CHECKING:
+    from .User import User
+
 
 @strawberry.type
 class UserApiKey(ApiKey, Node):
     id_attr: NodeID[int]
     user_id: Private[int]
+
+    @strawberry.field
+    async def user(self, info: Info[Context, None]) -> Annotated["User", strawberry.lazy(".User")]:
+        user = await info.context.data_loaders.users.load(self.user_id)
+        if user is None:
+            raise NotFound(f"User with id {self.user_id} not found")
+        from .User import to_gql_user
+
+        return to_gql_user(user)
+
+
+def to_gql_api_key(api_key: OrmApiKey) -> UserApiKey:
+    """
+    Converts an ORM API key to a GraphQL UserApiKey type.
+    """
+    return UserApiKey(
+        id_attr=api_key.id,
+        user_id=api_key.user_id,
+        name=api_key.name,
+        description=api_key.description,
+        created_at=api_key.created_at,
+        expires_at=api_key.expires_at,
+    )

phoenix/server/api/types/UserRole.py

@@ -1,8 +1,15 @@
 import strawberry
 from strawberry.relay import Node, NodeID
 
+from phoenix.db import models
+
 
 @strawberry.type
 class UserRole(Node):
     id_attr: NodeID[int]
     name: str
+
+
+def to_gql_user_role(role: models.UserRole) -> UserRole:
+    """Convert an ORM user role to a GraphQL user role."""
+    return UserRole(id_attr=role.id, name=role.name)