diff --git a/chromadb/auth/__init__.py b/chromadb/auth/__init__.py index 20baf53d95d..7c7223e4ff6 100644 --- a/chromadb/auth/__init__.py +++ b/chromadb/auth/__init__.py @@ -52,6 +52,10 @@ def get_user_id(self) -> str: def get_user_tenant(self) -> Optional[str]: ... + @abstractmethod + def get_user_databases(self) -> Optional[List[str]]: + ... + @abstractmethod def get_user_attributes(self) -> Optional[Dict[str, Any]]: ... @@ -62,11 +66,13 @@ def __init__( self, user_id: str, tenant: Optional[str] = None, + databases: Optional[List[str]] = None, attributes: Optional[Dict[str, Any]] = None, ) -> None: self._user_id = user_id self._tenant = tenant self._attributes = attributes + self._databases = databases @override def get_user_id(self) -> str: @@ -76,6 +82,10 @@ def get_user_id(self) -> str: def get_user_tenant(self) -> Optional[str]: return self._tenant if self._tenant else DEFAULT_TENANT + @override + def get_user_databases(self) -> Optional[List[str]]: + return self._databases + @override def get_user_attributes(self) -> Optional[Dict[str, Any]]: return self._attributes diff --git a/chromadb/auth/fastapi.py b/chromadb/auth/fastapi.py index 749e4677e9f..3e91bca7808 100644 --- a/chromadb/auth/fastapi.py +++ b/chromadb/auth/fastapi.py @@ -116,7 +116,7 @@ def instrument_server(self, app: ASGIApp) -> None: raise NotImplementedError("Not implemented yet") -class FastAPIChromaAuthMiddlewareWrapper(BaseHTTPMiddleware): # type: ignore +class FastAPIChromaAuthMiddlewareWrapper(BaseHTTPMiddleware): def __init__( self, app: ASGIApp, auth_middleware: FastAPIChromaAuthMiddleware ) -> None: @@ -154,6 +154,15 @@ async def dispatch( "authz_provider", default=None ) +# This needs to be module-level config, since it's used in authz_context() where we +# don't have a system (so don't have easy access to the settings). +overwrite_singleton_tenant_database_access_from_auth: bool = False + + +def set_overwrite_singleton_tenant_database_access_from_auth(overwrite: bool = False) -> None: + global overwrite_singleton_tenant_database_access_from_auth + overwrite_singleton_tenant_database_access_from_auth = overwrite + def authz_context( action: Union[str, AuthzResourceActions, List[str], List[AuthzResourceActions]], @@ -204,6 +213,15 @@ def wrapped(*args: Any, **kwargs: Dict[Any, Any]) -> Any: a_authz_responses.append(_provider.authorize(_context)) if not any(a_authz_responses): raise AuthorizationError("Unauthorized") + # In a multi-tenant environment, we may want to allow users to send + # requests without configuring a tenant and DB. If so, they can set + # the request tenant and DB however they like and we simply overwrite it. + if overwrite_singleton_tenant_database_access_from_auth: + if "tenant" in kwargs: + kwargs["tenant"] = request.state.user_identity.get_user_tenant() + databases = request.state.user_identity.get_user_databases() + if len(databases) == 1 and "database" in kwargs: + kwargs["database"] = databases[0] return f(*args, **kwargs) return wrapped @@ -267,7 +285,7 @@ def instrument_server(self, app: ASGIApp) -> None: raise NotImplementedError("Not implemented yet") -class FastAPIChromaAuthzMiddlewareWrapper(BaseHTTPMiddleware): # type: ignore +class FastAPIChromaAuthzMiddlewareWrapper(BaseHTTPMiddleware): def __init__( self, app: ASGIApp, authz_middleware: FastAPIChromaAuthzMiddleware ) -> None: diff --git a/chromadb/auth/token/__init__.py b/chromadb/auth/token/__init__.py index 3280cc2a4e8..4d1998ff5ee 100644 --- a/chromadb/auth/token/__init__.py +++ b/chromadb/auth/token/__init__.py @@ -122,6 +122,7 @@ class User(TypedDict): id: str role: str tenant: Optional[str] + databases: Optional[List[str]] tokens: List[Token] @@ -179,6 +180,7 @@ def get_user_identity( return SimpleUserIdentity( user_id=_user_id, tenant=_user["tenant"] if _user and "tenant" in _user else "*", + databases=_user["databases"] if _user and "databases" in _user else ["*"], ) diff --git a/chromadb/config.py b/chromadb/config.py index 6f59934619b..59bea5ee0e4 100644 --- a/chromadb/config.py +++ b/chromadb/config.py @@ -218,6 +218,9 @@ def chroma_server_authz_config_file_non_empty_file_exists( str ] = "chromadb.auth.authz.LocalUserConfigAuthorizationConfigurationProvider" + # TODO comment + chroma_overwrite_singleton_tenant_database_access_from_auth: bool = False + anonymized_telemetry: bool = True chroma_otel_collection_endpoint: Optional[str] = "" diff --git a/chromadb/server/fastapi/__init__.py b/chromadb/server/fastapi/__init__.py index 1377c64e5ed..ac1b9873bfa 100644 --- a/chromadb/server/fastapi/__init__.py +++ b/chromadb/server/fastapi/__init__.py @@ -22,6 +22,7 @@ FastAPIChromaAuthzMiddleware, FastAPIChromaAuthzMiddlewareWrapper, authz_context, + set_overwrite_singleton_tenant_database_access_from_auth, ) from chromadb.auth.fastapi_utils import ( attr_from_collection_lookup, @@ -159,6 +160,9 @@ def __init__(self, settings: Settings): FastAPIChromaAuthMiddlewareWrapper, auth_middleware=self._api.require(FastAPIChromaAuthMiddleware), ) + set_overwrite_singleton_tenant_database_access_from_auth( + settings.chroma_overwrite_singleton_tenant_database_access_from_auth + ) self.router = ChromaAPIRouter()