arize-phoenix 11.23.1__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/server/api/types/User.py

@@ -3,7 +3,6 @@ from typing import Optional
 
 import strawberry
 from sqlalchemy import select
-from strawberry import Private
 from strawberry.relay import Node, NodeID
 from strawberry.types import Info
 
@@ -12,59 +11,155 @@ 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 phoenix.server.api.types.UserApiKey import UserApiKey
+from phoenix.server.ldap import is_null_email_marker
 
 from .UserRole import UserRole, to_gql_user_role
 
 
 @strawberry.type
 class User(Node):
-    id_attr: NodeID[int]
-    password_needs_reset: bool
-    email: str
-    username: str
-    profile_picture_url: Optional[str]
-    created_at: datetime
-    user_role_id: Private[int]
-    auth_method: AuthMethod
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.User]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("User ID mismatch")
+
+    @strawberry.field
+    async def password_needs_reset(
+        self,
+        info: Info[Context, None],
+    ) -> bool:
+        if self.db_record:
+            val = self.db_record.reset_password
+        else:
+            val = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.reset_password),
+            )
+        return val
+
+    @strawberry.field
+    async def email(
+        self,
+        info: Info[Context, None],
+    ) -> str | None:
+        if self.db_record:
+            val = self.db_record.email
+        else:
+            val = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.email),
+            )
+        if is_null_email_marker(val):
+            return None
+        return val
+
+    @strawberry.field
+    async def username(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.username
+        else:
+            val = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.username),
+            )
+        return val
+
+    @strawberry.field
+    async def profile_picture_url(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.profile_picture_url
+        else:
+            val = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.profile_picture_url),
+            )
+        return val
+
+    @strawberry.field
+    async def created_at(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.created_at
+        else:
+            val = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.created_at),
+            )
+        return val
+
+    @strawberry.field
+    async def auth_method(
+        self,
+        info: Info[Context, None],
+    ) -> AuthMethod:
+        """Return semantic auth method (translated from database for LDAP users).
+
+        For Approach 1 (zero-migration), LDAP users are stored with auth_method='OAUTH2'
+        but have a special Unicode marker in oauth2_client_id. This resolver translates
+        that storage convention to the correct semantic AuthMethod.LDAP for the frontend.
+        """
+        if self.db_record:
+            auth_method_val = self.db_record.auth_method
+            oauth2_client_id = self.db_record.oauth2_client_id
+        else:
+            (
+                auth_method_val,
+                oauth2_client_id,
+            ) = await info.context.data_loaders.user_fields.load_many(
+                (
+                    (self.id, models.User.auth_method),
+                    (self.id, models.User.oauth2_client_id),
+                )
+            )
+
+        # Translate LDAP users from database storage to semantic type
+        if auth_method_val == "OAUTH2":
+            # Import here to avoid circular dependency
+            from phoenix.server.ldap import is_ldap_user
+
+            if is_ldap_user(oauth2_client_id):
+                return AuthMethod.LDAP
+
+        return AuthMethod(auth_method_val)
 
     @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 self.db_record:
+            user_role_id = self.db_record.user_role_id
+        else:
+            user_role_id = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.user_role_id),
+            )
+        role = await info.context.data_loaders.user_roles.load(user_role_id)
         if role is None:
-            raise NotFound(f"User role with id {self.user_role_id} not found")
+            raise NotFound(f"User role with id {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)
+                select(models.ApiKey).where(models.ApiKey.user_id == self.id)
             )
-        return [to_gql_api_key(api_key) for api_key in api_keys]
+        return [UserApiKey(id=api_key.id, db_record=api_key) for api_key in api_keys]
 
     @strawberry.field
-    async def is_management_user(self) -> bool:
+    async def is_management_user(self, info: Info[Context, None]) -> bool:
         initial_admins = get_env_admins()
         # this field is only visible to initial admins as they are the ones likely to have access to
         # a management interface / the phoenix environment.
-        if self.email in initial_admins or self.email == "admin@localhost":
+        if self.db_record:
+            email = self.db_record.email
+        else:
+            email = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.email),
+            )
+        if email in initial_admins or email == "admin@localhost":
             return True
         return False
-
-
-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,14 +1,13 @@
-from typing import TYPE_CHECKING
+from datetime import datetime
+from typing import TYPE_CHECKING, Optional
 
 import strawberry
-from strawberry import Private
 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
 
@@ -18,28 +17,76 @@ if TYPE_CHECKING:
 
 @strawberry.type
 class UserApiKey(ApiKey, Node):
-    id_attr: NodeID[int]
-    user_id: Private[int]
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[OrmApiKey]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("UserApiKey ID mismatch")
+
+    @strawberry.field(description="Name of the API key.")  # type: ignore
+    async def name(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.name
+        else:
+            val = await info.context.data_loaders.user_api_key_fields.load(
+                (self.id, OrmApiKey.name),
+            )
+        return val
+
+    @strawberry.field(description="Description of the API key.")  # type: ignore
+    async def description(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.description
+        else:
+            val = await info.context.data_loaders.user_api_key_fields.load(
+                (self.id, OrmApiKey.description),
+            )
+        return val
+
+    @strawberry.field(description="The date and time the API key was created.")  # type: ignore
+    async def created_at(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.created_at
+        else:
+            val = await info.context.data_loaders.user_api_key_fields.load(
+                (self.id, OrmApiKey.created_at),
+            )
+        return val
+
+    @strawberry.field(description="The date and time the API key will expire.")  # type: ignore
+    async def expires_at(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[datetime]:
+        if self.db_record:
+            val = self.db_record.expires_at
+        else:
+            val = await info.context.data_loaders.user_api_key_fields.load(
+                (self.id, OrmApiKey.expires_at),
+            )
+        return val
 
     @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,
-    )
+    async def user(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["User", strawberry.lazy(".User")]:
+        if self.db_record:
+            user_id = self.db_record.user_id
+        else:
+            user_id = await info.context.data_loaders.user_api_key_fields.load(
+                (self.id, OrmApiKey.user_id),
+            )
+        from .User import User
+
+        return User(id=user_id)

phoenix/server/api/types/node.py

@@ -1,5 +1,15 @@
+import re
+from base64 import b64decode
+
 from strawberry.relay import GlobalID
 
+_GLOBAL_ID_PATTERN = re.compile(r"[a-zA-Z]+:[0-9]+")
+
+
+def is_global_id(node_id: str) -> bool:
+    decoded_node_id = b64decode(node_id).decode()
+    return _GLOBAL_ID_PATTERN.match(decoded_node_id) is not None
+
 
 def from_global_id(global_id: GlobalID) -> tuple[str, int]:
     """

phoenix/server/api/types/pagination.py

@@ -9,7 +9,7 @@ from strawberry.relay.types import Connection, Edge, NodeType, PageInfo
 from typing_extensions import TypeAlias, assert_never
 
 ID: TypeAlias = int
-CursorSortColumnValue: TypeAlias = Union[str, int, float, datetime]
+CursorSortColumnValue: TypeAlias = Union[str, int, float, datetime, None]
 
 # A type alias for the connection cursor implementation
 CursorString = str
@@ -19,6 +19,7 @@ CURSOR_PREFIX = "connection:"
 
 
 class CursorSortColumnDataType(Enum):
+    NULL = auto()
     STRING = auto()
     INT = auto()
     FLOAT = auto()
@@ -30,6 +31,10 @@ class CursorSortColumn:
     type: CursorSortColumnDataType
     value: CursorSortColumnValue
 
+    def __post_init__(self) -> None:
+        if self.value is None:
+            self.type = CursorSortColumnDataType.NULL
+
     def __str__(self) -> str:
         if isinstance(self.value, str):
             return self.value
@@ -37,7 +42,9 @@ class CursorSortColumn:
             return str(self.value)
         if isinstance(self.value, datetime):
             return self.value.isoformat()
-        assert_never(self.type)
+        if self.value is None:
+            return ""
+        assert_never(self.value)
 
     @classmethod
     def from_string(cls, type: CursorSortColumnDataType, cursor_string: str) -> "CursorSortColumn":
@@ -50,6 +57,8 @@ class CursorSortColumn:
             value = float(cursor_string)
         elif type is CursorSortColumnDataType.DATETIME:
             value = datetime.fromisoformat(cursor_string)
+        elif type is CursorSortColumnDataType.NULL:
+            value = None
         else:
             assert_never(type)
         return cls(type=type, value=value)