Allow overwriting of request tenant and db from auth (#1498)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- This PR allows the auth system to overwrite a request's tenant and db.
An example use-case is when a client wants to connect to DEFAULT_TENANT
and DEFAULT_DATABASE with an API token in a multi-tenant system, but
have their requests routed to the correct tenant and database based on
the API key.

## Test plan
*How are these changes tested?*

- [x] Tests pass locally with `pytest` for python, `yarn test` for js
- [ ] New unit tests on this functionality

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the [docs
repository](https://github.com/chroma-core/docs)?*

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

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:

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 ["*"],
         )
 
 

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] = ""

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()