codekingpro/portable-devtools

1from typing import Dict
2from fastapi import HTTPException
3from overrides import override
4from chromadb.auth import (
5    AuthzAction,
6    AuthzResource,
7    ServerAuthenticationProvider,
8    ServerAuthorizationProvider,
9    UserIdentity,
10)
11from chromadb.config import System
12
13
14class ExampleAuthenticationProvider(ServerAuthenticationProvider):
15    """In practice the tenant would likely be resolved from some other opaque value (e.g. key/token). Here, it's just passed directly as a header for simplicity."""
16
17    @override
18    def authenticate_or_raise(self, headers: Dict[str, str]) -> UserIdentity:
19        return UserIdentity(
20            user_id="test",
21            tenant=headers.get("x-tenant", None),
22        )
23
24
25class ExampleAuthorizationProvider(ServerAuthorizationProvider):
26    """A simple authz provider that asserts the user's tenant matches the resource's tenant."""
27
28    def __init__(self, system: System) -> None:
29        super().__init__(system)
30        self._settings = system.settings
31
32    @override
33    def authorize_or_raise(
34        self, user: UserIdentity, action: AuthzAction, resource: AuthzResource
35    ) -> None:
36        if user.tenant is None:
37            return
38
39        if action == AuthzAction.RESET:
40            return
41
42        if user.tenant != resource.tenant:
43            raise HTTPException(status_code=403, detail="Unauthorized")
44