[ENH] Auto-set tenant in python CloudClient, add AdminCloudClient

Author: jairad26

chromadb/__init__.py

@@ -1,12 +1,15 @@
 from typing import Dict, Optional, Union
 import logging
 from chromadb.api.client import Client as ClientCreator
-from chromadb.api.client import AdminClient as AdminClientCreator
+from chromadb.api.client import (
+    AdminClient as AdminClientCreator,
+    AdminCloudClient as AdminCloudClientCreator,
+)
 from chromadb.api.async_client import AsyncClient as AsyncClientCreator
 from chromadb.auth.token_authn import TokenTransportHeader
 import chromadb.config
 from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Settings
-from chromadb.api import AdminAPI, AsyncClientAPI, ClientAPI
+from chromadb.api import AdminAPI, AsyncClientAPI, ClientAPI, AdminCloudAPI
 from chromadb.api.models.Collection import Collection
 from chromadb.api.types import (
     CollectionMetadata,
@@ -27,6 +30,7 @@
     UpdateCollectionMetadata,
 )
 from pathlib import Path
+import warnings
 
 # Re-export types from chromadb.types
 __all__ = [
@@ -311,19 +315,39 @@ def CloudClient(
     Creates a client to connect to a tennant and database on the Chroma cloud.
 
     Args:
-        tenant: The tenant to use for this client.
         database: The database to use for this client.
         api_key: The api key to use for this client.
     """
+    if tenant is not None:
+        warnings.warn(
+            "The 'tenant' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+    if cloud_port != 8000:
+        warnings.warn(
+            "The 'cloud_port' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+    if enable_ssl is not True:
+        warnings.warn(
+            "The 'enable_ssl' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     required_args = [
-        CloudClientArg(name="tenant", env_var="CHROMA_TENANT", value=tenant),
         CloudClientArg(name="database", env_var="CHROMA_DATABASE", value=database),
         CloudClientArg(name="api_key", env_var="CHROMA_API_KEY", value=api_key),
     ]
 
     # If any of tenant, database, or api_key is not provided, try to load it from the environment variable
     if not all([arg.value for arg in required_args]):
         import os
+
         for arg in required_args:
             arg.value = arg.value or os.environ.get(arg.env_var)
 
@@ -338,26 +362,45 @@ def CloudClient(
         settings = Settings()
 
     # Make sure paramaters are the correct types -- users can pass anything.
-    tenant = str(tenant)
     database = str(database)
     api_key = str(api_key)
     cloud_host = str(cloud_host)
-    cloud_port = int(cloud_port)
-    enable_ssl = bool(enable_ssl)
 
     settings.chroma_api_impl = "chromadb.api.fastapi.FastAPI"
     settings.chroma_server_host = cloud_host
-    settings.chroma_server_http_port = cloud_port
+    settings.chroma_server_http_port = 443
     # Always use SSL for cloud
-    settings.chroma_server_ssl_enabled = enable_ssl
+    settings.chroma_server_ssl_enabled = True
 
     settings.chroma_client_auth_provider = (
         "chromadb.auth.token_authn.TokenAuthClientProvider"
     )
     settings.chroma_client_auth_credentials = api_key
     settings.chroma_auth_token_transport_header = TokenTransportHeader.X_CHROMA_TOKEN
+    settings.chroma_overwrite_singleton_tenant_database_access_from_auth = True
 
-    return ClientCreator(tenant=tenant, database=database, settings=settings)
+    return ClientCreator(database=database, settings=settings)
+
+
+def AdminCloudClient(
+    api_key: str,
+    settings: Optional[Settings] = None,
+    *,
+    cloud_host: str = "api.trychroma.com",
+) -> AdminCloudAPI:
+    if settings is None:
+        settings = Settings()
+    settings.chroma_api_impl = "chromadb.api.fastapi.FastAPI"
+    settings.chroma_server_host = cloud_host
+    settings.chroma_server_http_port = 443
+    settings.chroma_server_ssl_enabled = True
+    settings.chroma_client_auth_provider = (
+        "chromadb.auth.token_authn.TokenAuthClientProvider"
+    )
+    settings.chroma_client_auth_credentials = api_key
+    settings.chroma_auth_token_transport_header = TokenTransportHeader.X_CHROMA_TOKEN
+    settings.chroma_overwrite_singleton_tenant_database_access_from_auth = True
+    return AdminCloudClientCreator(settings=settings)
 
 
 def Client(

chromadb/api/__init__.py

@@ -562,6 +562,63 @@ def get_tenant(self, name: str) -> Tenant:
         pass
 
 
+class AdminCloudAPI(ABC):
+    @abstractmethod
+    def create_database(self, name: str) -> None:
+        """Create a new database. Raises an error if the database already exists.
+
+        Args:
+            name: The name of the database to create.
+
+        """
+        pass
+
+    @abstractmethod
+    def get_database(self, name: str) -> Database:
+        """Get a database. Raises an error if the database does not exist.
+
+        Args:
+            name: The name of the database to get.
+
+        """
+        pass
+
+    @abstractmethod
+    def delete_database(self, name: str) -> None:
+        """Delete a database. Raises an error if the database does not exist.
+
+        Args:
+            name: The name of the database to delete.
+
+        """
+        pass
+
+    @abstractmethod
+    def list_databases(
+        self,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+    ) -> Sequence[Database]:
+        """List all databases for a tenant. Raises an error if the tenant does not exist.
+
+        Args:
+            limit: The maximum number of databases to return.
+            offset: The offset to start from.
+
+        """
+        pass
+
+    @abstractmethod
+    def get_tenant(self, name: str) -> Tenant:
+        """Get a tenant. Raises an error if the tenant does not exist.
+
+        Args:
+            name: The name of the tenant to get.
+
+        """
+        pass
+
+
 class ServerAPI(BaseAPI, AdminAPI, Component):
     """An API instance that extends the relevant Base API methods by passing
     in a tenant and database. This is the root component of the Chroma System"""

chromadb/api/client.py

@@ -3,7 +3,7 @@
 
 from overrides import override
 import httpx
-from chromadb.api import AdminAPI, ClientAPI, ServerAPI
+from chromadb.api import AdminAPI, ClientAPI, ServerAPI, AdminCloudAPI
 from chromadb.api.collection_configuration import (
     CreateCollectionConfiguration,
     UpdateCollectionConfiguration,
@@ -525,3 +525,59 @@ def from_system(
         SharedSystemClient._populate_data_from_system(system)
         instance = cls(settings=system.settings)
         return instance
+
+
+class AdminCloudClient(SharedSystemClient, AdminCloudAPI):
+    _server: ServerAPI
+    tenant: str
+
+    def __init__(self, settings: Settings = Settings()) -> None:
+        super().__init__(settings)
+        self._server = self._system.instance(ServerAPI)
+
+        user_identity = self._server.get_user_identity()
+
+        maybe_tenant, _ = maybe_set_tenant_and_database(
+            user_identity,
+            overwrite_singleton_tenant_database_access_from_auth=settings.chroma_overwrite_singleton_tenant_database_access_from_auth,
+            user_provided_tenant=None,
+            user_provided_database=None,
+        )
+        if maybe_tenant:
+            self.tenant = maybe_tenant
+        else:
+            raise ValueError("Could not get tenant from user identity")
+
+    @override
+    def create_database(self, name: str) -> None:
+        return self._server.create_database(name=name, tenant=self.tenant)
+
+    @override
+    def get_database(self, name: str) -> Database:
+        return self._server.get_database(name=name, tenant=self.tenant)
+
+    @override
+    def delete_database(self, name: str) -> None:
+        return self._server.delete_database(name=name, tenant=self.tenant)
+
+    @override
+    def list_databases(
+        self,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+    ) -> Sequence[Database]:
+        return self._server.list_databases(limit, offset, tenant=self.tenant)
+
+    @override
+    def get_tenant(self, name: str) -> Tenant:
+        return self._server.get_tenant(name=name)
+
+    @classmethod
+    @override
+    def from_system(
+        cls,
+        system: System,
+    ) -> "AdminCloudClient":
+        SharedSystemClient._populate_data_from_system(system)
+        instance = cls(settings=system.settings)
+        return instance

chromadb/test/client/test_cloud_client.py

@@ -1,101 +1,52 @@
-import multiprocessing
-from typing import Any, Dict, Generator, Optional, Tuple
 import pytest
+from unittest.mock import patch
 from chromadb import CloudClient
-from chromadb.api import ServerAPI
-from chromadb.auth.token_authn import TokenTransportHeader
-from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Settings, System
 from chromadb.errors import ChromaAuthError
+from chromadb.auth import UserIdentity
+from chromadb.types import Tenant, Database
+from uuid import uuid4
+
+
+def test_valid_key() -> None:
+    with patch(
+        "chromadb.api.fastapi.FastAPI.get_user_identity"
+    ) as mock_get_user_identity, patch(
+        "chromadb.api.client.AdminClient.get_tenant"
+    ) as mock_get_tenant, patch(
+        "chromadb.api.client.AdminClient.get_database"
+    ) as mock_get_database, patch(
+        "chromadb.api.fastapi.FastAPI.heartbeat"
+    ) as mock_heartbeat:
+        mock_get_user_identity.return_value = UserIdentity(
+            user_id="test_user", tenant="default_tenant", databases=["testdb"]
+        )
+        mock_get_tenant.return_value = Tenant(name="default_tenant")
+        mock_get_database.return_value = Database(
+            id=uuid4(), name="testdb", tenant="default_tenant"
+        )
+        mock_heartbeat.return_value = 1234567890
 
-from chromadb.test.conftest import _await_server, _run_server, find_free_port
-
-TOKEN_TRANSPORT_HEADER = TokenTransportHeader.X_CHROMA_TOKEN
-TEST_CLOUD_HOST = "localhost"
-
-
-@pytest.fixture(scope="module")
-def valid_token() -> str:
-    return "valid_token"
-
-
-@pytest.fixture(scope="module")
-def mock_cloud_server(valid_token: str) -> Generator[System, None, None]:
-    chroma_server_authn_provider: str = (
-        "chromadb.auth.token_authn.TokenAuthenticationServerProvider"
-    )
-    chroma_server_authn_credentials: str = valid_token
-    chroma_auth_token_transport_header: str = TOKEN_TRANSPORT_HEADER
-
-    port = find_free_port()
-
-    args: Tuple[
-        int,
-        bool,
-        Optional[str],
-        Optional[str],
-        Optional[str],
-        Optional[str],
-        Optional[str],
-        Optional[str],
-        Optional[str],
-        Optional[Dict[str, Any]],
-    ] = (
-        port,
-        False,
-        None,
-        chroma_server_authn_provider,
-        None,
-        chroma_server_authn_credentials,
-        chroma_auth_token_transport_header,
-        None,
-        None,
-        None,
-    )
-    ctx = multiprocessing.get_context("spawn")
-    proc = ctx.Process(target=_run_server, args=args, daemon=True)
-    proc.start()
-
-    settings = Settings(
-        chroma_api_impl="chromadb.api.fastapi.FastAPI",
-        chroma_server_host=TEST_CLOUD_HOST,
-        chroma_server_http_port=port,
-        chroma_client_auth_provider="chromadb.auth.token_authn.TokenAuthClientProvider",
-        chroma_client_auth_credentials=valid_token,
-        chroma_auth_token_transport_header=TOKEN_TRANSPORT_HEADER,
-    )
+        client = CloudClient(database="testdb", api_key="valid_token")
 
-    system = System(settings)
-    api = system.instance(ServerAPI)
-    system.start()
-    _await_server(api)
-    yield system
-    system.stop()
-    proc.kill()
+        assert client.get_user_identity().user_id == "test_user"
+        assert client.get_user_identity().tenant == "default_tenant"
+        assert client.get_user_identity().databases == ["testdb"]
 
+        settings = client.get_settings()
+        assert settings.chroma_client_auth_credentials == "valid_token"
+        assert (
+            settings.chroma_client_auth_provider
+            == "chromadb.auth.token_authn.TokenAuthClientProvider"
+        )
 
-def test_valid_key(mock_cloud_server: System, valid_token: str) -> None:
-    valid_client = CloudClient(
-        tenant=DEFAULT_TENANT,
-        database=DEFAULT_DATABASE,
-        api_key=valid_token,
-        cloud_host=TEST_CLOUD_HOST,
-        cloud_port=mock_cloud_server.settings.chroma_server_http_port or 8000,
-        enable_ssl=False,
-    )
+        assert client.heartbeat() == 1234567890
 
-    assert valid_client.heartbeat()
 
+def test_invalid_key() -> None:
+    with patch(
+        "chromadb.api.fastapi.FastAPI.get_user_identity"
+    ) as mock_get_user_identity:
+        mock_get_user_identity.side_effect = ChromaAuthError("Authentication failed")
 
-def test_invalid_key(mock_cloud_server: System, valid_token: str) -> None:
-    # Try to connect to the default tenant and database with an invalid token
-    invalid_token = valid_token + "_invalid"
-    with pytest.raises(ChromaAuthError):
-        client = CloudClient(
-            tenant=DEFAULT_TENANT,
-            database=DEFAULT_DATABASE,
-            api_key=invalid_token,
-            cloud_host=TEST_CLOUD_HOST,
-            cloud_port=mock_cloud_server.settings.chroma_server_http_port or 8000,
-            enable_ssl=False,
-        )
-        client.heartbeat()
+        with pytest.raises(ChromaAuthError):
+            CloudClient(database="testdb", api_key="invalid_token")