arize-phoenix 11.3.0__py3-none-any.whl → 11.5.0__py3-none-any.whl

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

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from fastapi import APIRouter, HTTPException, Path, Query
+from fastapi import APIRouter, Depends, HTTPException, Path, Query
 from pydantic import Field
 from sqlalchemy import select
 from starlette.requests import Request
@@ -23,6 +23,7 @@ from phoenix.server.api.routers.v1.utils import (
     add_errors_to_responses,
 )
 from phoenix.server.api.types.Project import Project as ProjectNodeType
+from phoenix.server.authorization import is_not_locked
 
 router = APIRouter(tags=["projects"])
 
@@ -173,6 +174,7 @@ async def get_project(
 
 @router.post(
     "/projects",
+    dependencies=[Depends(is_not_locked)],
     operation_id="createProject",
     summary="Create a new project",  # noqa: E501
     description="Create a new project with the specified configuration.",  # noqa: E501
@@ -213,6 +215,7 @@ async def create_project(
 
 @router.put(
     "/projects/{project_identifier}",
+    dependencies=[Depends(is_not_locked)],
     operation_id="updateProject",
     summary="Update a project by ID or name",  # noqa: E501
     description="Update an existing project with new configuration. Project names cannot be changed. The project identifier is either project ID or project name. Note: When using a project name as the identifier, it cannot contain slash (/), question mark (?), or pound sign (#) characters.",  # noqa: E501

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

@@ -1,7 +1,7 @@
 import logging
 from typing import Any, Optional, Union
 
-from fastapi import APIRouter, HTTPException, Path, Query
+from fastapi import APIRouter, Depends, HTTPException, Path, Query
 from pydantic import ValidationError, model_validator
 from sqlalchemy import select
 from sqlalchemy.sql import Select
@@ -33,6 +33,7 @@ from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Prompt import Prompt as PromptNodeType
 from phoenix.server.api.types.PromptVersion import PromptVersion as PromptVersionNodeType
 from phoenix.server.api.types.PromptVersionTag import PromptVersionTag as PromptVersionTagNodeType
+from phoenix.server.authorization import is_not_locked
 from phoenix.server.bearer_auth import PhoenixUser
 
 logger = logging.getLogger(__name__)
@@ -393,6 +394,7 @@ async def get_prompt_version_by_latest(
 
 @router.post(
     "/prompts",
+    dependencies=[Depends(is_not_locked)],
     operation_id="postPromptVersion",
     summary="Create a new prompt",
     description="Create a new prompt and its initial version. A prompt can have multiple versions.",
@@ -602,6 +604,7 @@ async def list_prompt_version_tags(
 
 @router.post(
     "/prompt_versions/{prompt_version_id}/tags",
+    dependencies=[Depends(is_not_locked)],
     operation_id="createPromptVersionTag",
     summary="Add tag to prompt version",
     description="Add a new tag to a specific prompt version. Tags help identify and categorize "

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

@@ -8,7 +8,7 @@ from secrets import token_urlsafe
 from typing import Annotated, Any, Literal, Optional, Union
 
 import pandas as pd
-from fastapi import APIRouter, Header, HTTPException, Path, Query
+from fastapi import APIRouter, Depends, Header, HTTPException, Path, Query
 from pydantic import BaseModel, Field
 from sqlalchemy import select
 from starlette.requests import Request
@@ -28,6 +28,7 @@ from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
 from phoenix.db.insertion.types import Precursors
 from phoenix.server.api.routers.utils import df_to_bytes
+from phoenix.server.authorization import is_not_locked
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import SpanAnnotationInsertEvent
 from phoenix.trace.attributes import flatten
@@ -907,6 +908,7 @@ class AnnotateSpansResponseBody(ResponseBody[list[InsertedSpanAnnotation]]):
 
 @router.post(
     "/span_annotations",
+    dependencies=[Depends(is_not_locked)],
     operation_id="annotateSpans",
     summary="Create span annotations",
     responses=add_errors_to_responses(
@@ -990,6 +992,7 @@ class CreateSpansResponseBody(V1RoutesBaseModel):
 
 @router.post(
     "/projects/{project_identifier}/spans",
+    dependencies=[Depends(is_not_locked)],
     operation_id="createSpans",
     summary="Create spans",
     description=(

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

@@ -2,7 +2,7 @@ import gzip
 import zlib
 from typing import Any, Literal, Optional
 
-from fastapi import APIRouter, BackgroundTasks, Header, HTTPException, Query
+from fastapi import APIRouter, BackgroundTasks, Depends, Header, HTTPException, Query
 from google.protobuf.message import DecodeError
 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
     ExportTraceServiceRequest,
@@ -24,6 +24,7 @@ from strawberry.relay import GlobalID
 from phoenix.db import models
 from phoenix.db.insertion.helpers import as_kv
 from phoenix.db.insertion.types import Precursors
+from phoenix.server.authorization import is_not_locked
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import TraceAnnotationInsertEvent
 from phoenix.trace.otel import decode_otlp_span
@@ -37,6 +38,7 @@ router = APIRouter(tags=["traces"])
 
 @router.post(
     "/traces",
+    dependencies=[Depends(is_not_locked)],
     operation_id="addTraces",
     summary="Send traces",
     responses=add_errors_to_responses(
@@ -160,6 +162,7 @@ class AnnotateTracesResponseBody(ResponseBody[list[InsertedTraceAnnotation]]):
 
 @router.post(
     "/trace_annotations",
+    dependencies=[Depends(is_not_locked)],
     operation_id="annotateTraces",
     summary="Create trace annotations",
     responses=add_errors_to_responses(

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

@@ -43,7 +43,7 @@ from phoenix.server.api.routers.v1.utils import (
     add_errors_to_responses,
 )
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.authorization import require_admin
+from phoenix.server.authorization import is_not_locked, require_admin
 
 logger = logging.getLogger(__name__)
 
@@ -194,7 +194,7 @@ async def list_users(
             HTTP_422_UNPROCESSABLE_ENTITY,
         ]
     ),
-    dependencies=[Depends(require_admin)],
+    dependencies=[Depends(require_admin), Depends(is_not_locked)],
     response_model_by_alias=True,
     response_model_exclude_unset=True,
     response_model_exclude_defaults=True,

phoenix/server/app.py

@@ -16,17 +16,20 @@ from typing import (
     Any,
     NamedTuple,
     Optional,
+    Protocol,
     TypedDict,
     Union,
     cast,
 )
 from urllib.parse import urlparse
 
+import grpc
 import strawberry
 from fastapi import APIRouter, Depends, FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.utils import is_body_allowed_for_status_code
 from grpc.aio import ServerInterceptor
+from grpc_interceptor import AsyncServerInterceptor
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
 from starlette.datastructures import URL, Secret
@@ -44,7 +47,7 @@ from starlette.types import Scope, StatefulLifespan
 from strawberry.extensions import SchemaExtension
 from strawberry.fastapi import GraphQLRouter
 from strawberry.subscriptions import GRAPHQL_TRANSPORT_WS_PROTOCOL
-from typing_extensions import TypeAlias
+from typing_extensions import TypeAlias, override
 
 import phoenix.trace.v1 as pb
 from phoenix.config import (
@@ -134,6 +137,7 @@ from phoenix.server.api.routers import (
 from phoenix.server.api.routers.v1 import REST_API_VERSION
 from phoenix.server.api.schema import build_graphql_schema
 from phoenix.server.bearer_auth import BearerTokenAuthBackend, is_authenticated
+from phoenix.server.daemons.db_disk_usage_monitor import DbDiskUsageMonitor
 from phoenix.server.daemons.generative_model_store import GenerativeModelStore
 from phoenix.server.daemons.span_cost_calculator import SpanCostCalculator
 from phoenix.server.dml_event import DmlEvent
@@ -523,6 +527,7 @@ def _lifespan(
     trace_data_sweeper: Optional[TraceDataSweeper],
     span_cost_calculator: SpanCostCalculator,
     generative_model_store: GenerativeModelStore,
+    db_disk_usage_monitor: DbDiskUsageMonitor,
     token_store: Optional[TokenStore] = None,
     tracer_provider: Optional["TracerProvider"] = None,
     enable_prometheus: bool = False,
@@ -530,6 +535,7 @@ def _lifespan(
     shutdown_callbacks: Iterable[_Callback] = (),
     read_only: bool = False,
     scaffolder_config: Optional[ScaffolderConfig] = None,
+    grpc_interceptors: Iterable[AsyncServerInterceptor] = (),
 ) -> StatefulLifespan[FastAPI]:
     @contextlib.asynccontextmanager
     async def lifespan(_: FastAPI) -> AsyncIterator[dict[str, Any]]:
@@ -551,7 +557,7 @@ def _lifespan(
                 tracer_provider=tracer_provider,
                 enable_prometheus=enable_prometheus,
                 token_store=token_store,
-                interceptors=user_grpc_interceptors(),
+                interceptors=user_grpc_interceptors() + list(grpc_interceptors),
             )
             await stack.enter_async_context(grpc_server)
             await stack.enter_async_context(dml_event_handler)
@@ -559,6 +565,7 @@ def _lifespan(
                 await stack.enter_async_context(trace_data_sweeper)
             await stack.enter_async_context(span_cost_calculator)
             await stack.enter_async_context(generative_model_store)
+            await stack.enter_async_context(db_disk_usage_monitor)
             if scaffolder_config:
                 scaffolder = Scaffolder(
                     config=scaffolder_config,
@@ -826,6 +833,34 @@ async def plain_text_http_exception_handler(request: Request, exc: HTTPException
     return PlainTextResponse(str(exc.detail), status_code=exc.status_code, headers=headers)
 
 
+class _HasDbStatus(Protocol):
+    @property
+    def should_not_insert_or_update(self) -> bool: ...
+
+
+class DbDiskUsageInterceptor(AsyncServerInterceptor):
+    def __init__(self, db: _HasDbStatus) -> None:
+        self._db = db
+
+    @override
+    async def intercept(
+        self,
+        method: Callable[[Any, grpc.aio.ServicerContext], Awaitable[Any]],
+        request_or_iterator: Any,
+        context: grpc.aio.ServicerContext,
+        method_name: str,
+    ) -> Any:
+        if (
+            method_name.endswith("trace.v1.TraceService/Export")
+            and self._db.should_not_insert_or_update
+        ):
+            await context.abort(
+                grpc.StatusCode.RESOURCE_EXHAUSTED,
+                "Database disk usage threshold exceeded",
+            )
+        return await method(request_or_iterator, context)
+
+
 def create_app(
     db: DbSessionFactory,
     export_path: Path,
@@ -971,6 +1006,8 @@ def create_app(
         from phoenix.server.prometheus import PrometheusMiddleware
 
         middlewares.append(Middleware(PrometheusMiddleware))
+    grpc_interceptors: list[AsyncServerInterceptor] = []
+    grpc_interceptors.append(DbDiskUsageInterceptor(db))
     app = FastAPI(
         title="Arize-Phoenix REST API",
         version=REST_API_VERSION,
@@ -982,6 +1019,8 @@ def create_app(
             trace_data_sweeper=trace_data_sweeper,
             span_cost_calculator=span_cost_calculator,
             generative_model_store=generative_model_store,
+            db_disk_usage_monitor=DbDiskUsageMonitor(db, email_sender),
+            grpc_interceptors=grpc_interceptors,
             token_store=token_store,
             tracer_provider=tracer_provider,
             enable_prometheus=enable_prometheus,

phoenix/server/authorization.py

@@ -51,3 +51,12 @@ def require_admin(request: Request) -> None:
             status_code=fastapi_status.HTTP_403_FORBIDDEN,
             detail="Only admin or system users can perform this action.",
         )
+
+
+def is_not_locked(request: Request) -> None:
+    if request.app.state.db.should_not_insert_or_update:
+        raise HTTPException(
+            status_code=fastapi_status.HTTP_507_INSUFFICIENT_STORAGE,
+            detail="Operations that insert or update database "
+            "records are currently not allowed.",
+        )

phoenix/server/bearer_auth.py

@@ -7,10 +7,10 @@ from typing import Any, Optional, cast
 import grpc
 from fastapi import HTTPException, Request, WebSocket, WebSocketException
 from grpc_interceptor import AsyncServerInterceptor
-from grpc_interceptor.exceptions import Unauthenticated
 from starlette.authentication import AuthCredentials, AuthenticationBackend, BaseUser
 from starlette.requests import HTTPConnection
 from starlette.status import HTTP_401_UNAUTHORIZED
+from typing_extensions import override
 
 from phoenix import config
 from phoenix.auth import (
@@ -100,16 +100,19 @@ class PhoenixSystemUser(PhoenixUser):
 
 
 class ApiKeyInterceptor(HasTokenStore, AsyncServerInterceptor):
+    @override
     async def intercept(
         self,
-        method: Callable[[Any, grpc.ServicerContext], Awaitable[Any]],
+        method: Callable[[Any, grpc.aio.ServicerContext], Awaitable[Any]],
         request_or_iterator: Any,
-        context: grpc.ServicerContext,
+        context: grpc.aio.ServicerContext,
         method_name: str,
     ) -> Any:
-        for datum in context.invocation_metadata():
-            if datum.key.lower() == "authorization":
-                scheme, _, token = datum.value.partition(" ")
+        for key, value in context.invocation_metadata() or ():
+            if key.lower() == "authorization":
+                if isinstance(value, bytes):
+                    value = value.decode("utf-8")
+                scheme, _, token = value.partition(" ")
                 if scheme.lower() != "bearer" or not token:
                     break
                 if (
@@ -119,16 +122,16 @@ class ApiKeyInterceptor(HasTokenStore, AsyncServerInterceptor):
                 ):
                     return await method(request_or_iterator, context)
                 claims = await self._token_store.read(Token(token))
-                if not (isinstance(claims, UserClaimSet) and isinstance(claims.subject, UserId)):
+                if (
+                    not (
+                        isinstance(claims, (ApiKeyClaims, AccessTokenClaims))
+                        and isinstance(claims.subject, UserId)
+                    )
+                    or claims.status is not ClaimSetStatus.VALID
+                ):
                     break
-                if not isinstance(claims, (ApiKeyClaims, AccessTokenClaims)):
-                    raise Unauthenticated(details="Invalid token")
-                if claims.status is ClaimSetStatus.EXPIRED:
-                    raise Unauthenticated(details="Expired token")
-                if claims.status is ClaimSetStatus.VALID:
-                    return await method(request_or_iterator, context)
-                raise Unauthenticated()
-        raise Unauthenticated()
+                return await method(request_or_iterator, context)
+        await context.abort(grpc.StatusCode.UNAUTHENTICATED)
 
 
 async def is_authenticated(