arize-phoenix 3.16.1__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/server/api/mutations/user_mutations.py

@@ -0,0 +1,329 @@
+import secrets
+from contextlib import AsyncExitStack
+from datetime import datetime, timezone
+from typing import Literal, Optional
+
+import strawberry
+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 (
+    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, IsLocked, 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.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: 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
+
+
+@strawberry.type
+class UserMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin, IsLocked])  # type: ignore
+    async def create_user(
+        self,
+        info: Info[Context, None],
+        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,
+        )
+        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, IsLocked])  # 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, IsLocked])  # 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, IsLocked])  # 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()
+        )
+        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,
+                    )
+                )
+            )
+            .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 _user_operation_error_message(
+    error: IntegrityError,
+    operation: Literal["create", "modify"] = "create",
+) -> str:
+    """
+    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
+    email_already_exists = "users.email" in original_error_message
+    if username_already_exists:
+        return "Username already exists"
+    elif email_already_exists:
+        return "Email already exists"
+    return f"Failed to {operation} user"

phoenix/server/api/openapi/main.py

@@ -0,0 +1,17 @@
+import json
+from argparse import ArgumentParser
+
+from phoenix.server.api.openapi.schema import get_openapi_schema
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        required=True,
+        help="Path to the output file (e.g., openapi.json)",
+    )
+    args = parser.parse_args()
+    with open(args.output, "w") as f:
+        json.dump(get_openapi_schema(), f, indent=2)

phoenix/server/api/openapi/schema.py

@@ -0,0 +1,16 @@
+from typing import Any
+
+from fastapi.openapi.utils import get_openapi
+
+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,
+        openapi_version="3.1.0",
+        description="Schema for Arize-Phoenix REST API",
+        routes=v1_router.routes,
+    )