arize-phoenix 4.35.2__py3-none-any.whl → 5.0.0__py3-none-any.whl

phoenix/server/jwt_store.py

@@ -0,0 +1,504 @@
+import logging
+from abc import ABC, abstractmethod
+from asyncio import create_task, gather, sleep
+from copy import deepcopy
+from dataclasses import replace
+from datetime import datetime, timezone
+from functools import cached_property, singledispatchmethod
+from typing import Any, Callable, Coroutine, Dict, Generic, List, Optional, Tuple, Type, TypeVar
+
+from authlib.jose import jwt
+from authlib.jose.errors import JoseError
+from sqlalchemy import Select, delete, select
+
+from phoenix.auth import (
+    JWT_ALGORITHM,
+    ClaimSet,
+    Token,
+)
+from phoenix.config import get_env_enable_prometheus
+from phoenix.db import models
+from phoenix.db.enums import UserRole
+from phoenix.server.types import (
+    AccessToken,
+    AccessTokenAttributes,
+    AccessTokenClaims,
+    AccessTokenId,
+    ApiKey,
+    ApiKeyAttributes,
+    ApiKeyClaims,
+    ApiKeyId,
+    DaemonTask,
+    DbSessionFactory,
+    PasswordResetToken,
+    PasswordResetTokenAttributes,
+    PasswordResetTokenClaims,
+    PasswordResetTokenId,
+    RefreshToken,
+    RefreshTokenAttributes,
+    RefreshTokenClaims,
+    RefreshTokenId,
+    TokenId,
+    UserId,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class JwtStore:
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        secret: str,
+        algorithm: str = JWT_ALGORITHM,
+        sleep_seconds: int = 10,
+        **kwargs: Any,
+    ) -> None:
+        assert secret
+        super().__init__(**kwargs)
+        self._db = db
+        self._secret = secret
+        args = (db, secret, algorithm, sleep_seconds)
+        self._password_reset_token_store = _PasswordResetTokenStore(*args, **kwargs)
+        self._access_token_store = _AccessTokenStore(*args, **kwargs)
+        self._refresh_token_store = _RefreshTokenStore(*args, **kwargs)
+        self._api_key_store = _ApiKeyStore(*args, **kwargs)
+
+    @cached_property
+    def _stores(self) -> Tuple[DaemonTask, ...]:
+        return tuple(dt for dt in self.__dict__.values() if isinstance(dt, _Store))
+
+    async def __aenter__(self) -> None:
+        await gather(*(s.__aenter__() for s in self._stores))
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        await gather(*(s.__aexit__(*args, **kwargs) for s in self._stores))
+
+    async def read(self, token: Token) -> Optional[ClaimSet]:
+        try:
+            payload = jwt.decode(
+                s=token,
+                key=self._secret,
+            )
+        except JoseError:
+            return None
+        if (jti := payload.get("jti")) is None:
+            return None
+        if (token_id := TokenId.parse(jti)) is None:
+            return None
+        return await self._get(token_id)
+
+    @singledispatchmethod
+    async def _get(self, _: TokenId) -> Optional[ClaimSet]:
+        return None
+
+    @_get.register
+    async def _(self, token_id: PasswordResetTokenId) -> Optional[ClaimSet]:
+        return await self._password_reset_token_store.get(token_id)
+
+    @_get.register
+    async def _(self, token_id: AccessTokenId) -> Optional[ClaimSet]:
+        return await self._access_token_store.get(token_id)
+
+    @_get.register
+    async def _(self, token_id: RefreshTokenId) -> Optional[ClaimSet]:
+        return await self._refresh_token_store.get(token_id)
+
+    @_get.register
+    async def _(self, token_id: ApiKeyId) -> Optional[ClaimSet]:
+        return await self._api_key_store.get(token_id)
+
+    @singledispatchmethod
+    async def _evict(self, _: TokenId) -> Optional[ClaimSet]:
+        return None
+
+    @_evict.register
+    async def _(self, token_id: PasswordResetTokenId) -> Optional[ClaimSet]:
+        return await self._password_reset_token_store.evict(token_id)
+
+    @_evict.register
+    async def _(self, token_id: AccessTokenId) -> Optional[ClaimSet]:
+        return await self._access_token_store.evict(token_id)
+
+    @_evict.register
+    async def _(self, token_id: RefreshTokenId) -> Optional[ClaimSet]:
+        return await self._refresh_token_store.evict(token_id)
+
+    @_evict.register
+    async def _(self, token_id: ApiKeyId) -> Optional[ClaimSet]:
+        return await self._api_key_store.evict(token_id)
+
+    async def create_password_reset_token(
+        self,
+        claim: PasswordResetTokenClaims,
+    ) -> Tuple[PasswordResetToken, PasswordResetTokenId]:
+        return await self._password_reset_token_store.create(claim)
+
+    async def create_access_token(
+        self,
+        claim: AccessTokenClaims,
+    ) -> Tuple[AccessToken, AccessTokenId]:
+        return await self._access_token_store.create(claim)
+
+    async def create_refresh_token(
+        self,
+        claim: RefreshTokenClaims,
+    ) -> Tuple[RefreshToken, RefreshTokenId]:
+        return await self._refresh_token_store.create(claim)
+
+    async def create_api_key(
+        self,
+        claim: ApiKeyClaims,
+    ) -> Tuple[ApiKey, ApiKeyId]:
+        return await self._api_key_store.create(claim)
+
+    async def revoke(self, *token_ids: TokenId) -> None:
+        if not token_ids:
+            return
+        password_reset_token_ids: List[PasswordResetTokenId] = []
+        access_token_ids: List[AccessTokenId] = []
+        refresh_token_ids: List[RefreshTokenId] = []
+        api_key_ids: List[ApiKeyId] = []
+        for token_id in token_ids:
+            if isinstance(token_id, PasswordResetTokenId):
+                password_reset_token_ids.append(token_id)
+            if isinstance(token_id, AccessTokenId):
+                access_token_ids.append(token_id)
+            elif isinstance(token_id, RefreshTokenId):
+                refresh_token_ids.append(token_id)
+            elif isinstance(token_id, ApiKeyId):
+                api_key_ids.append(token_id)
+        coroutines: List[Coroutine[None, None, None]] = []
+        if password_reset_token_ids:
+            coroutines.append(self._password_reset_token_store.revoke(*password_reset_token_ids))
+        if access_token_ids:
+            coroutines.append(self._access_token_store.revoke(*access_token_ids))
+        if refresh_token_ids:
+            coroutines.append(self._refresh_token_store.revoke(*refresh_token_ids))
+        if api_key_ids:
+            coroutines.append(self._api_key_store.revoke(*api_key_ids))
+        await gather(*coroutines)
+
+    async def log_out(self, user_id: UserId) -> None:
+        for cls in (AccessTokenId, RefreshTokenId):
+            table = cls.table
+            stmt = delete(table).where(table.user_id == int(user_id)).returning(table.id)
+            async with self._db() as session:
+                async for id_ in await session.stream_scalars(stmt):
+                    await self._evict(cls(id_))
+
+
+_TokenT = TypeVar("_TokenT", bound=Token)
+_TokenIdT = TypeVar("_TokenIdT", bound=TokenId)
+_ClaimSetT = TypeVar("_ClaimSetT", bound=ClaimSet)
+_RecordT = TypeVar(
+    "_RecordT",
+    models.PasswordResetToken,
+    models.AccessToken,
+    models.RefreshToken,
+    models.ApiKey,
+)
+
+
+class _Claims(Generic[_TokenIdT, _ClaimSetT]):
+    def __init__(self) -> None:
+        self._cache: Dict[_TokenIdT, _ClaimSetT] = {}
+
+    def __getitem__(self, token_id: _TokenIdT) -> Optional[_ClaimSetT]:
+        claim = self._cache.get(token_id)
+        return deepcopy(claim) if claim else None
+
+    def __setitem__(self, token_id: _TokenIdT, claim: _ClaimSetT) -> None:
+        self._cache[token_id] = deepcopy(claim)
+
+    def get(self, token_id: _TokenIdT) -> Optional[_ClaimSetT]:
+        claim = self._cache.get(token_id)
+        return deepcopy(claim) if claim else None
+
+    def pop(
+        self, token_id: _TokenIdT, default: Optional[_ClaimSetT] = None
+    ) -> Optional[_ClaimSetT]:
+        claim = self._cache.pop(token_id, default)
+        return deepcopy(claim) if claim else None
+
+
+class _Store(DaemonTask, Generic[_ClaimSetT, _TokenT, _TokenIdT, _RecordT], ABC):
+    _table: Type[_RecordT]
+    _token_id: Callable[[int], _TokenIdT]
+    _token: Callable[[str], _TokenT]
+
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        secret: str,
+        algorithm: str = JWT_ALGORITHM,
+        sleep_seconds: int = 10,
+        **kwargs: Any,
+    ) -> None:
+        assert secret
+        super().__init__(**kwargs)
+        self._db = db
+        self._seconds = sleep_seconds
+        self._claims: _Claims[_TokenIdT, _ClaimSetT] = _Claims()
+        self._secret = secret
+        self._algorithm = algorithm
+
+    def _encode(self, claim: ClaimSet) -> str:
+        payload: Dict[str, Any] = dict(jti=claim.token_id)
+        header = {"alg": self._algorithm}
+        jwt_bytes: bytes = jwt.encode(header=header, payload=payload, key=self._secret)
+        return jwt_bytes.decode()
+
+    async def get(self, token_id: _TokenIdT) -> Optional[_ClaimSetT]:
+        if claims := self._claims.get(token_id):
+            return claims
+        stmt = self._update_stmt.where(self._table.id == int(token_id))
+        async with self._db() as session:
+            record = (await session.execute(stmt)).first()
+        if not record:
+            return None
+        token, role = record
+        _, claims = self._from_db(token, UserRole(role))
+        self._claims[token_id] = claims
+        return claims
+
+    async def evict(self, token_id: _TokenIdT) -> Optional[_ClaimSetT]:
+        return self._claims.pop(token_id, None)
+
+    async def revoke(self, *token_ids: _TokenIdT) -> None:
+        if not token_ids:
+            return
+        for token_id in token_ids:
+            await self.evict(token_id)
+        stmt = delete(self._table).where(self._table.id.in_(map(int, token_ids)))
+        async with self._db() as session:
+            await session.execute(stmt)
+
+    @abstractmethod
+    def _from_db(self, record: _RecordT, role: UserRole) -> Tuple[_TokenIdT, _ClaimSetT]: ...
+
+    @abstractmethod
+    def _to_db(self, claims: _ClaimSetT) -> _RecordT: ...
+
+    async def create(self, claim: _ClaimSetT) -> Tuple[_TokenT, _TokenIdT]:
+        record = self._to_db(claim)
+        async with self._db() as session:
+            session.add(record)
+            await session.flush()
+        token_id = self._token_id(record.id)
+        claim = replace(claim, token_id=token_id)
+        self._claims[token_id] = claim
+        token = self._token(self._encode(claim))
+        return token, token_id
+
+    async def _update(self) -> None:
+        claims: _Claims[_TokenIdT, _ClaimSetT] = _Claims()
+        async with self._db() as session:
+            async with session.begin_nested():
+                await self._delete_expired_tokens(session)
+            async with session.begin_nested():
+                async for record, role in await session.stream(self._update_stmt):
+                    token_id, claim_set = self._from_db(record, UserRole(role))
+                    claims[token_id] = claim_set
+        self._claims = claims
+
+    @cached_property
+    def _update_stmt(self) -> Select[Tuple[_RecordT, str]]:
+        return (
+            select(self._table, models.UserRole.name)
+            .join_from(self._table, models.User)
+            .join_from(models.User, models.UserRole)
+        )
+
+    async def _delete_expired_tokens(self, session: Any) -> None:
+        now = datetime.now(timezone.utc)
+        await session.execute(delete(self._table).where(self._table.expires_at < now))
+
+    async def _run(self) -> None:
+        while self._running:
+            self._tasks.append(create_task(self._update()))
+            await self._tasks[-1]
+            self._tasks.pop()
+            self._tasks.append(create_task(sleep(self._seconds)))
+            await self._tasks[-1]
+            self._tasks.pop()
+
+
+class _PasswordResetTokenStore(
+    _Store[
+        PasswordResetTokenClaims,
+        PasswordResetToken,
+        PasswordResetTokenId,
+        models.PasswordResetToken,
+    ]
+):
+    _table = models.PasswordResetToken
+    _token_id = PasswordResetTokenId
+    _token = PasswordResetToken
+
+    def _from_db(
+        self,
+        record: models.PasswordResetToken,
+        user_role: UserRole,
+    ) -> Tuple[PasswordResetTokenId, PasswordResetTokenClaims]:
+        token_id = PasswordResetTokenId(record.id)
+        return token_id, PasswordResetTokenClaims(
+            token_id=token_id,
+            subject=UserId(record.user_id),
+            issued_at=record.created_at,
+            expiration_time=record.expires_at,
+            attributes=PasswordResetTokenAttributes(
+                user_role=user_role,
+            ),
+        )
+
+    def _to_db(self, claim: PasswordResetTokenClaims) -> models.PasswordResetToken:
+        assert claim.expiration_time
+        assert claim.subject
+        user_id = int(claim.subject)
+        return models.PasswordResetToken(
+            user_id=user_id,
+            created_at=claim.issued_at,
+            expires_at=claim.expiration_time,
+        )
+
+
+class _AccessTokenStore(
+    _Store[
+        AccessTokenClaims,
+        AccessToken,
+        AccessTokenId,
+        models.AccessToken,
+    ]
+):
+    _table = models.AccessToken
+    _token_id = AccessTokenId
+    _token = AccessToken
+
+    def _from_db(
+        self,
+        record: models.AccessToken,
+        user_role: UserRole,
+    ) -> Tuple[AccessTokenId, AccessTokenClaims]:
+        token_id = AccessTokenId(record.id)
+        refresh_token_id = RefreshTokenId(record.refresh_token_id)
+        return token_id, AccessTokenClaims(
+            token_id=token_id,
+            subject=UserId(record.user_id),
+            issued_at=record.created_at,
+            expiration_time=record.expires_at,
+            attributes=AccessTokenAttributes(
+                user_role=user_role,
+                refresh_token_id=refresh_token_id,
+            ),
+        )
+
+    def _to_db(self, claim: AccessTokenClaims) -> models.AccessToken:
+        assert claim.expiration_time
+        assert claim.subject
+        user_id = int(claim.subject)
+        assert claim.attributes
+        refresh_token_id = int(claim.attributes.refresh_token_id)
+        return models.AccessToken(
+            user_id=user_id,
+            created_at=claim.issued_at,
+            expires_at=claim.expiration_time,
+            refresh_token_id=refresh_token_id,
+        )
+
+
+class _RefreshTokenStore(
+    _Store[
+        RefreshTokenClaims,
+        RefreshToken,
+        RefreshTokenId,
+        models.RefreshToken,
+    ]
+):
+    _table = models.RefreshToken
+    _token_id = RefreshTokenId
+    _token = RefreshToken
+
+    def _from_db(
+        self,
+        record: models.RefreshToken,
+        user_role: UserRole,
+    ) -> Tuple[RefreshTokenId, RefreshTokenClaims]:
+        token_id = RefreshTokenId(record.id)
+        return token_id, RefreshTokenClaims(
+            token_id=token_id,
+            subject=UserId(record.user_id),
+            issued_at=record.created_at,
+            expiration_time=record.expires_at,
+            attributes=RefreshTokenAttributes(
+                user_role=user_role,
+            ),
+        )
+
+    def _to_db(self, claims: RefreshTokenClaims) -> models.RefreshToken:
+        assert claims.expiration_time
+        assert claims.subject
+        user_id = int(claims.subject)
+        return models.RefreshToken(
+            user_id=user_id,
+            created_at=claims.issued_at,
+            expires_at=claims.expiration_time,
+        )
+
+    async def _update(self) -> None:
+        await super()._update()
+        if get_env_enable_prometheus():
+            from phoenix.server.prometheus import JWT_STORE_TOKENS_ACTIVE
+
+            JWT_STORE_TOKENS_ACTIVE.set(len(self._claims._cache))
+
+
+class _ApiKeyStore(
+    _Store[
+        ApiKeyClaims,
+        ApiKey,
+        ApiKeyId,
+        models.ApiKey,
+    ]
+):
+    _table = models.ApiKey
+    _token_id = ApiKeyId
+    _token = ApiKey
+
+    def _from_db(
+        self,
+        record: models.ApiKey,
+        user_role: UserRole,
+    ) -> Tuple[ApiKeyId, ApiKeyClaims]:
+        token_id = ApiKeyId(record.id)
+        return token_id, ApiKeyClaims(
+            token_id=token_id,
+            subject=UserId(record.user_id),
+            issued_at=record.created_at,
+            expiration_time=record.expires_at,
+            attributes=ApiKeyAttributes(
+                user_role=user_role,
+                name=record.name,
+                description=record.description,
+            ),
+        )
+
+    def _to_db(self, claims: ApiKeyClaims) -> models.ApiKey:
+        assert claims.attributes
+        assert claims.attributes.name
+        assert claims.subject
+        user_id = int(claims.subject)
+        return models.ApiKey(
+            user_id=user_id,
+            name=claims.attributes.name,
+            description=claims.attributes.description or None,
+            created_at=claims.issued_at,
+            expires_at=claims.expiration_time or None,
+        )
+
+    async def _update(self) -> None:
+        await super()._update()
+        if get_env_enable_prometheus():
+            from phoenix.server.prometheus import JWT_STORE_API_KEYS_ACTIVE
+
+            JWT_STORE_API_KEYS_ACTIVE.set(len(self._claims._cache))