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

phoenix/server/api/mutations/trace_annotations_mutations.py

@@ -6,11 +6,11 @@ from strawberry import UNSET
 from strawberry.types import Info
 
 from phoenix.db import models
+from phoenix.server.api.auth import IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.input_types.CreateTraceAnnotationInput import CreateTraceAnnotationInput
 from phoenix.server.api.input_types.DeleteAnnotationsInput import DeleteAnnotationsInput
 from phoenix.server.api.input_types.PatchAnnotationInput import PatchAnnotationInput
-from phoenix.server.api.mutations.auth import IsAuthenticated
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
@@ -25,7 +25,7 @@ class TraceAnnotationMutationPayload:
 
 @strawberry.type
 class TraceAnnotationMutationMixin:
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def create_trace_annotations(
         self, info: Info[Context, None], input: List[CreateTraceAnnotationInput]
     ) -> TraceAnnotationMutationPayload:
@@ -59,7 +59,7 @@ class TraceAnnotationMutationMixin:
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def patch_trace_annotations(
         self, info: Info[Context, None], input: List[PatchAnnotationInput]
     ) -> TraceAnnotationMutationPayload:
@@ -98,7 +98,7 @@ class TraceAnnotationMutationMixin:
                     info.context.event_queue.put(TraceAnnotationInsertEvent((trace_annotation.id,)))
         return TraceAnnotationMutationPayload(trace_annotations=patched_annotations, query=Query())
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def delete_trace_annotations(
         self, info: Info[Context, None], input: DeleteAnnotationsInput
     ) -> TraceAnnotationMutationPayload:

phoenix/server/api/mutations/user_mutations.py

@@ -1,27 +1,79 @@
-import asyncio
-from typing import Optional
+import secrets
+from contextlib import AsyncExitStack
+from datetime import datetime, timezone
+from typing import List, Literal, Optional, Tuple
 
 import strawberry
-from sqlalchemy import insert, select
+from sqlalchemy import Boolean, Select, and_, case, cast, delete, distinct, func, select
+from sqlalchemy.orm import joinedload
 from sqlean.dbapi2 import IntegrityError  # type: ignore[import-untyped]
+from strawberry import UNSET
+from strawberry.relay import GlobalID
 from strawberry.types import Info
 
-from phoenix.auth import compute_password_hash, validate_email_format, validate_password_format
-from phoenix.db import models
+from phoenix.auth import (
+    DEFAULT_ADMIN_EMAIL,
+    DEFAULT_ADMIN_USERNAME,
+    DEFAULT_SECRET_LENGTH,
+    PASSWORD_REQUIREMENTS,
+    PHOENIX_ACCESS_TOKEN_COOKIE_NAME,
+    PHOENIX_REFRESH_TOKEN_COOKIE_NAME,
+    validate_email_format,
+    validate_password_format,
+)
+from phoenix.db import enums, models
+from phoenix.server.api.auth import IsAdmin, IsNotReadOnly
 from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import Conflict, NotFound, Unauthorized
 from phoenix.server.api.input_types.UserRoleInput import UserRoleInput
-from phoenix.server.api.types.User import User
-from phoenix.server.api.types.UserRole import UserRole
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.User import User, to_gql_user
+from phoenix.server.bearer_auth import PhoenixUser
+from phoenix.server.types import AccessTokenId, ApiKeyId, PasswordResetTokenId, RefreshTokenId
 
 
 @strawberry.input
 class CreateUserInput:
     email: str
-    username: Optional[str]
+    username: str
     password: str
     role: UserRoleInput
 
 
+@strawberry.input
+class PatchViewerInput:
+    new_username: Optional[str] = UNSET
+    new_password: Optional[str] = UNSET
+    current_password: Optional[str] = UNSET
+
+    def __post_init__(self) -> None:
+        if not self.new_username and not self.new_password:
+            raise ValueError("At least one field must be set")
+        if self.new_password and not self.current_password:
+            raise ValueError("current_password is required when modifying password")
+        if self.new_password:
+            PASSWORD_REQUIREMENTS.validate(self.new_password)
+
+
+@strawberry.input
+class PatchUserInput:
+    user_id: GlobalID
+    new_role: Optional[UserRoleInput] = UNSET
+    new_username: Optional[str] = UNSET
+    new_password: Optional[str] = UNSET
+
+    def __post_init__(self) -> None:
+        if not self.new_role and not self.new_username and not self.new_password:
+            raise ValueError("At least one field must be set")
+        if self.new_password:
+            PASSWORD_REQUIREMENTS.validate(self.new_password)
+
+
+@strawberry.input
+class DeleteUsersInput:
+    user_ids: List[GlobalID]
+
+
 @strawberry.type
 class UserMutationPayload:
     user: User
@@ -29,7 +81,7 @@ class UserMutationPayload:
 
 @strawberry.type
 class UserMutationMixin:
-    @strawberry.mutation
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin])  # type: ignore
     async def create_user(
         self,
         info: Info[Context, None],
@@ -37,47 +89,235 @@ class UserMutationMixin:
     ) -> UserMutationPayload:
         validate_email_format(email := input.email)
         validate_password_format(password := input.password)
-        role_name = input.role.value
-        user_role_id = (
-            select(models.UserRole.id).where(models.UserRole.name == role_name).scalar_subquery()
+        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,
         )
-        secret = info.context.get_secret()
-        loop = asyncio.get_running_loop()
-        password_hash = await loop.run_in_executor(
-            executor=None,
-            func=lambda: compute_password_hash(password=password, salt=secret),
+        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))
+            if user_role_id is None:
+                raise NotFound(f"Role {input.role.value} not found")
+            stack.enter_context(session.no_autoflush)
+            user.user_role_id = user_role_id
+            session.add(user)
+            try:
+                await session.flush()
+            except IntegrityError as error:
+                raise Conflict(_user_operation_error_message(error))
+        return UserMutationPayload(user=to_gql_user(user))
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin])  # type: ignore
+    async def patch_user(
+        self,
+        info: Info[Context, None],
+        input: PatchUserInput,
+    ) -> UserMutationPayload:
+        assert (request := info.context.request)
+        assert isinstance(request.user, PhoenixUser)
+        assert (requester_id := int(request.user.identity))
+        user_id = from_global_id_with_expected_type(input.user_id, expected_type_name=User.__name__)
+        async with AsyncExitStack() as stack:
+            session = await stack.enter_async_context(info.context.db())
+            requester = await session.scalar(_select_user_by_id(requester_id))
+            assert requester
+            if not (user := await session.scalar(_select_user_by_id(user_id))):
+                raise NotFound("User not found")
+            stack.enter_context(session.no_autoflush)
+            if input.new_role:
+                if user.email == DEFAULT_ADMIN_EMAIL:
+                    raise Unauthorized("Cannot modify role for the default admin user")
+                user_role_id = await session.scalar(_select_role_id_by_name(input.new_role.value))
+                if user_role_id is None:
+                    raise NotFound(f"Role {input.new_role.value} not found")
+                user.user_role_id = user_role_id
+            if password := input.new_password:
+                if user.auth_method != enums.AuthMethod.LOCAL.value:
+                    raise Conflict("Cannot modify password for non-local user")
+                validate_password_format(password)
+                user.password_salt = secrets.token_bytes(DEFAULT_SECRET_LENGTH)
+                user.password_hash = await info.context.hash_password(password, user.password_salt)
+                user.reset_password = True
+            if username := input.new_username:
+                user.username = username
+            assert user in session.dirty
+            try:
+                await session.flush()
+            except IntegrityError as error:
+                raise Conflict(_user_operation_error_message(error, "modify"))
+        assert user
+        if input.new_password:
+            await info.context.log_out(user.id)
+        return UserMutationPayload(user=to_gql_user(user))
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def patch_viewer(
+        self,
+        info: Info[Context, None],
+        input: PatchViewerInput,
+    ) -> UserMutationPayload:
+        assert (request := info.context.request)
+        assert isinstance(user := request.user, PhoenixUser)
+        user_id = int(user.identity)
+        async with AsyncExitStack() as stack:
+            session = await stack.enter_async_context(info.context.db())
+            if not (user := await session.scalar(_select_user_by_id(user_id))):
+                raise NotFound("User not found")
+            stack.enter_context(session.no_autoflush)
+            if password := input.new_password:
+                if user.auth_method != enums.AuthMethod.LOCAL.value:
+                    raise Conflict("Cannot modify password for non-local user")
+                if not (
+                    current_password := input.current_password
+                ) or not await info.context.is_valid_password(current_password, user):
+                    raise Conflict("Valid current password is required to modify password")
+                validate_password_format(password)
+                user.password_salt = secrets.token_bytes(DEFAULT_SECRET_LENGTH)
+                user.password_hash = await info.context.hash_password(password, user.password_salt)
+                user.reset_password = False
+            if username := input.new_username:
+                user.username = username
+            assert user in session.dirty
+            user.updated_at = datetime.now(timezone.utc)
+            try:
+                await session.flush()
+            except IntegrityError as error:
+                raise Conflict(_user_operation_error_message(error, "modify"))
+        assert user
+        if input.new_password:
+            await info.context.log_out(user.id)
+            response = info.context.get_response()
+            response.delete_cookie(PHOENIX_REFRESH_TOKEN_COOKIE_NAME)
+            response.delete_cookie(PHOENIX_ACCESS_TOKEN_COOKIE_NAME)
+        return UserMutationPayload(user=to_gql_user(user))
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin])  # type: ignore
+    async def delete_users(
+        self,
+        info: Info[Context, None],
+        input: DeleteUsersInput,
+    ) -> None:
+        assert (token_store := info.context.token_store) is not None
+        if not input.user_ids:
+            return
+        user_ids = tuple(
+            map(
+                lambda gid: from_global_id_with_expected_type(gid, User.__name__),
+                set(input.user_ids),
+            )
+        )
+        system_user_role_id = (
+            select(models.UserRole.id)
+            .where(models.UserRole.name == enums.UserRole.SYSTEM.value)
+            .scalar_subquery()
+        )
+        admin_user_role_id = (
+            select(models.UserRole.id)
+            .where(models.UserRole.name == enums.UserRole.ADMIN.value)
+            .scalar_subquery()
         )
-        try:
-            async with info.context.db() as session:
-                user = await session.scalar(
-                    insert(models.User)
-                    .values(
-                        user_role_id=user_role_id,
-                        username=input.username,
-                        email=email,
-                        auth_method="LOCAL",
-                        password_hash=password_hash,
-                        reset_password=True,
+        default_admin_user_id = (
+            select(models.User.id)
+            .where(
+                (
+                    and_(
+                        models.User.user_role_id == admin_user_role_id,
+                        models.User.username == DEFAULT_ADMIN_USERNAME,
+                        models.User.email == DEFAULT_ADMIN_EMAIL,
                     )
-                    .returning(models.User)
                 )
-            assert user is not None
-        except IntegrityError as error:
-            raise ValueError(_get_user_create_error_message(error))
-        return UserMutationPayload(
-            user=User(
-                id_attr=user.id,
-                email=user.email,
-                username=user.username,
-                created_at=user.created_at,
-                role=UserRole(id_attr=user.user_role_id, name=role_name),
             )
+            .scalar_subquery()
         )
+        async with info.context.db() as session:
+            [
+                (
+                    deletes_default_admin,
+                    num_resolved_user_ids,
+                )
+            ] = (
+                await session.execute(
+                    select(
+                        cast(
+                            func.coalesce(
+                                func.max(
+                                    case((models.User.id == default_admin_user_id, 1), else_=0)
+                                ),
+                                0,
+                            ),
+                            Boolean,
+                        ).label("deletes_default_admin"),
+                        func.count(distinct(models.User.id)).label("num_resolved_user_ids"),
+                    )
+                    .select_from(models.User)
+                    .where(
+                        and_(
+                            models.User.id.in_(user_ids),
+                            models.User.user_role_id != system_user_role_id,
+                        )
+                    )
+                )
+            ).all()
+            if deletes_default_admin:
+                raise Conflict("Cannot delete the default admin user")
+            if num_resolved_user_ids < len(user_ids):
+                raise NotFound("Some user IDs could not be found")
+            password_reset_token_ids = [
+                PasswordResetTokenId(id_)
+                async for id_ in await session.stream_scalars(
+                    select(models.PasswordResetToken.id).where(
+                        models.PasswordResetToken.user_id.in_(user_ids)
+                    )
+                )
+            ]
+            access_token_ids = [
+                AccessTokenId(id_)
+                async for id_ in await session.stream_scalars(
+                    select(models.AccessToken.id).where(models.AccessToken.user_id.in_(user_ids))
+                )
+            ]
+            refresh_token_ids = [
+                RefreshTokenId(id_)
+                async for id_ in await session.stream_scalars(
+                    select(models.RefreshToken.id).where(models.RefreshToken.user_id.in_(user_ids))
+                )
+            ]
+            api_key_ids = [
+                ApiKeyId(id_)
+                async for id_ in await session.stream_scalars(
+                    select(models.ApiKey.id).where(models.ApiKey.user_id.in_(user_ids))
+                )
+            ]
+            await session.execute(delete(models.User).where(models.User.id.in_(user_ids)))
+        await token_store.revoke(
+            *password_reset_token_ids,
+            *access_token_ids,
+            *refresh_token_ids,
+            *api_key_ids,
+        )
+
+
+def _select_role_id_by_name(role_name: str) -> Select[Tuple[int]]:
+    return select(models.UserRole.id).where(models.UserRole.name == role_name)
+
+
+def _select_user_by_id(user_id: int) -> Select[Tuple[models.User]]:
+    return (
+        select(models.User).where(models.User.id == user_id).options(joinedload(models.User.role))
+    )
 
 
-def _get_user_create_error_message(error: IntegrityError) -> str:
+def _user_operation_error_message(
+    error: IntegrityError,
+    operation: Literal["create", "modify"] = "create",
+) -> str:
     """
-    Gets a user-facing error message to explain why user creation failed.
+    User-facing error message to explain why user creation/modification failed.
     """
     original_error_message = str(error)
     username_already_exists = "users.username" in original_error_message
@@ -86,4 +326,4 @@ def _get_user_create_error_message(error: IntegrityError) -> str:
         return "Username already exists"
     elif email_already_exists:
         return "Email already exists"
-    return "Failed to create user"
+    return f"Failed to {operation} user"

phoenix/server/api/openapi/schema.py

@@ -2,11 +2,11 @@ from typing import Any, Dict
 
 from fastapi.openapi.utils import get_openapi
 
-from phoenix.server.api.routers.v1 import REST_API_VERSION
-from phoenix.server.api.routers.v1 import router as v1_router
+from phoenix.server.api.routers.v1 import REST_API_VERSION, create_v1_router
 
 
 def get_openapi_schema() -> Dict[str, Any]:
+    v1_router = create_v1_router(authentication_enabled=False)
     return get_openapi(
         title="Arize-Phoenix REST API",
         version=REST_API_VERSION,

phoenix/server/api/queries.py

@@ -7,12 +7,13 @@ import numpy.typing as npt
 import strawberry
 from sqlalchemy import and_, distinct, func, select
 from sqlalchemy.orm import joinedload
+from starlette.authentication import UnauthenticatedUser
 from strawberry import ID, UNSET
 from strawberry.relay import Connection, GlobalID, Node
 from strawberry.types import Info
 from typing_extensions import Annotated, TypeAlias
 
-from phoenix.db import models
+from phoenix.db import enums, models
 from phoenix.db.models import (
     DatasetExample as OrmExample,
 )
@@ -32,8 +33,9 @@ from phoenix.db.models import (
     Trace as OrmTrace,
 )
 from phoenix.pointcloud.clustering import Hdbscan
+from phoenix.server.api.auth import MSG_ADMIN_ONLY, IsAdmin
 from phoenix.server.api.context import Context
-from phoenix.server.api.exceptions import NotFound
+from phoenix.server.api.exceptions import NotFound, Unauthorized
 from phoenix.server.api.helpers import ensure_list
 from phoenix.server.api.input_types.ClusterInput import ClusterInput
 from phoenix.server.api.input_types.Coordinates import (
@@ -69,14 +71,14 @@ from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span, to_gql_span
 from phoenix.server.api.types.SystemApiKey import SystemApiKey
 from phoenix.server.api.types.Trace import Trace
-from phoenix.server.api.types.User import User
-from phoenix.server.api.types.UserApiKey import UserApiKey
+from phoenix.server.api.types.User import User, to_gql_user
+from phoenix.server.api.types.UserApiKey import UserApiKey, to_gql_api_key
 from phoenix.server.api.types.UserRole import UserRole
 
 
 @strawberry.type
 class Query:
-    @strawberry.field
+    @strawberry.field(permission_classes=[IsAdmin])  # type: ignore
     async def users(
         self,
         info: Info[Context, None],
@@ -94,25 +96,13 @@ class Query:
         stmt = (
             select(models.User)
             .join(models.UserRole)
-            .where(models.UserRole.name != "SYSTEM")
+            .where(models.UserRole.name != enums.UserRole.SYSTEM.value)
             .order_by(models.User.email)
             .options(joinedload(models.User.role))
         )
         async with info.context.db() as session:
             users = await session.stream_scalars(stmt)
-            data = [
-                User(
-                    id_attr=user.id,
-                    email=user.email,
-                    username=user.username,
-                    created_at=user.created_at,
-                    role=UserRole(
-                        id_attr=user.role.id,
-                        name=user.role.name,
-                    ),
-                )
-                async for user in users
-            ]
+            data = [to_gql_user(user) async for user in users]
         return connection_from_list(data=data, args=args)
 
     @strawberry.field
@@ -122,7 +112,7 @@ class Query:
     ) -> List[UserRole]:
         async with info.context.db() as session:
             roles = await session.scalars(
-                select(models.UserRole).where(models.UserRole.name != "SYSTEM")
+                select(models.UserRole).where(models.UserRole.name != enums.UserRole.SYSTEM.value)
             )
         return [
             UserRole(
@@ -132,37 +122,25 @@ class Query:
             for role in roles
         ]
 
-    @strawberry.field
+    @strawberry.field(permission_classes=[IsAdmin])  # type: ignore
     async def user_api_keys(self, info: Info[Context, None]) -> List[UserApiKey]:
-        # TODO(auth): add access control
         stmt = (
-            select(models.APIKey)
+            select(models.ApiKey)
             .join(models.User)
             .join(models.UserRole)
-            .where(models.UserRole.name != "SYSTEM")
+            .where(models.UserRole.name != enums.UserRole.SYSTEM.value)
         )
         async with info.context.db() as session:
             api_keys = await session.scalars(stmt)
-        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,
-            )
-            for api_key in api_keys
-        ]
+        return [to_gql_api_key(api_key) for api_key in api_keys]
 
-    @strawberry.field
+    @strawberry.field(permission_classes=[IsAdmin])  # type: ignore
     async def system_api_keys(self, info: Info[Context, None]) -> List[SystemApiKey]:
-        # TODO(auth): add access control
         stmt = (
-            select(models.APIKey)
+            select(models.ApiKey)
             .join(models.User)
             .join(models.UserRole)
-            .where(models.UserRole.name == "SYSTEM")
+            .where(models.UserRole.name == enums.UserRole.SYSTEM.value)
         )
         async with info.context.db() as session:
             api_keys = await session.scalars(stmt)
@@ -488,8 +466,39 @@ class Query:
                 ):
                     raise NotFound(f"Unknown experiment run: {id}")
             return to_gql_experiment_run(run)
+        elif type_name == User.__name__:
+            if int((user := info.context.user).identity) != node_id and not user.is_admin:
+                raise Unauthorized(MSG_ADMIN_ONLY)
+            async with info.context.db() as session:
+                if not (
+                    user := await session.scalar(
+                        select(models.User).where(models.User.id == node_id)
+                    )
+                ):
+                    raise NotFound(f"Unknown user: {id}")
+            return to_gql_user(user)
         raise NotFound(f"Unknown node type: {type_name}")
 
+    @strawberry.field
+    async def viewer(self, info: Info[Context, None]) -> Optional[User]:
+        request = info.context.get_request()
+        try:
+            user = request.user
+        except AssertionError:
+            return None
+        if isinstance(user, UnauthenticatedUser):
+            return None
+        async with info.context.db() as session:
+            if (
+                user := await session.scalar(
+                    select(models.User)
+                    .where(models.User.id == int(user.identity))
+                    .options(joinedload(models.User.role))
+                )
+            ) is None:
+                return None
+        return to_gql_user(user)
+
     @strawberry.field
     def clusters(
         self,

phoenix/server/api/routers/__init__.py

@@ -0,0 +1,11 @@
+from .auth import router as auth_router
+from .embeddings import create_embeddings_router
+from .oauth2 import router as oauth2_router
+from .v1 import create_v1_router
+
+__all__ = [
+    "auth_router",
+    "create_embeddings_router",
+    "create_v1_router",
+    "oauth2_router",
+]