diff --git a/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage.py b/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage.py index 5dc612737..f9e55c634 100644 --- a/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage.py +++ b/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage.py @@ -52,34 +52,34 @@ def __init__(self, config: CosmosDBStorageConfig): self._lock: asyncio.Lock = asyncio.Lock() def _create_client(self) -> CosmosClient: - if self._config.url: - if not self._config.credential: - raise ValueError( - storage_errors.InvalidConfiguration.format( - "Credential is required when using a custom service URL" - ) - ) - return CosmosClient( - url=self._config.url, credential=self._config.credential + + if not self._config.cosmos_db_endpoint: + raise ValueError(str(storage_errors.CosmosDbEndpointRequired)) + + if self._config.credential or self._config.auth_key: + cred = ( + self._config.credential + if self._config.credential + else self._config.auth_key ) - connection_policy = self._config.cosmos_client_options.get( - "connection_policy", documents.ConnectionPolicy() - ) + connection_policy = self._config.cosmos_client_options.get( + "connection_policy", documents.ConnectionPolicy() + ) - # kwargs 'connection_verify' is to handle CosmosClient overwriting the - # ConnectionPolicy.DisableSSLVerification value. - return CosmosClient( - self._config.cosmos_db_endpoint, - self._config.auth_key, - consistency_level=self._config.cosmos_client_options.get( - "consistency_level", None - ), - **{ - "connection_policy": connection_policy, - "connection_verify": not connection_policy.DisableSSLVerification, - }, - ) + return CosmosClient( + url=self._config.cosmos_db_endpoint, + credential=cred, + consistency_level=self._config.cosmos_client_options.get( + "consistency_level", None + ), + **{ + "connection_policy": connection_policy, + "connection_verify": not connection_policy.DisableSSLVerification, + }, + ) + + raise ValueError(str(storage_errors.CosmosDbAuthKeyRequired)) def _sanitize(self, key: str) -> str: return sanitize_key( diff --git a/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage_config.py b/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage_config.py index 596a10360..6fd9c0cbd 100644 --- a/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage_config.py +++ b/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage_config.py @@ -1,10 +1,13 @@ import json +import logging from azure.core.credentials_async import AsyncTokenCredential from microsoft_agents.storage.cosmos.errors import storage_errors from .key_ops import sanitize_key +logger = logging.getLogger(__name__) + class CosmosDBStorageConfig: """The class for partitioned CosmosDB configuration for the Azure Bot Framework.""" @@ -19,7 +22,6 @@ def __init__( container_throughput: int | None = None, key_suffix: str = "", compatibility_mode: bool = False, - url: str = "", credential: AsyncTokenCredential | None = None, **kwargs, ): @@ -32,13 +34,12 @@ def __init__( :param cosmos_client_options: The options for the CosmosClient. Currently only supports connection_policy and consistency_level :param container_throughput: The throughput set when creating the Container. Defaults to 400. - :param key_suffix: The suffix to be added to every key. The keySuffix must contain only valid ComosDb + :param key_suffix: The suffix to be added to every key. The keySuffix must contain only valid CosmosDb key characters. (e.g. not: '\\', '?', '/', '#', '*') :param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb max key length of 255. - :param url: The URL to the CosmosDB resource. :param credential: The TokenCredential to use for authentication. - :return CosmosDBConfig: + :return CosmosDBStorageConfig: """ config_file: str = kwargs.get("filename", "") if config_file: @@ -47,6 +48,14 @@ def __init__( self.cosmos_db_endpoint: str = cosmos_db_endpoint or kwargs.get( "cosmos_db_endpoint", "" ) + + if "url" in kwargs: + logger.warning( + "The 'url' parameter is deprecated. Please use 'cosmos_db_endpoint' instead." + ) + if not self.cosmos_db_endpoint: + self.cosmos_db_endpoint = kwargs.get("url", "") + self.auth_key: str = auth_key or kwargs.get("auth_key", "") self.database_id: str = database_id or kwargs.get("database_id", "") self.container_id: str = container_id or kwargs.get("container_id", "") @@ -60,7 +69,6 @@ def __init__( self.compatibility_mode: bool = compatibility_mode or kwargs.get( "compatibility_mode", False ) - self.url = url or kwargs.get("url", "") self.credential: AsyncTokenCredential | None = credential @staticmethod diff --git a/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/errors/error_resources.py b/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/errors/error_resources.py index ddec9e90a..332c9bbeb 100644 --- a/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/errors/error_resources.py +++ b/libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/errors/error_resources.py @@ -28,7 +28,7 @@ class StorageErrorResources: ) CosmosDbAuthKeyRequired = ErrorMessage( - "CosmosDBStorage: auth_key is required.", + "CosmosDBStorage: auth_key or credential is required.", -61002, ) diff --git a/tests/storage_cosmos/test_cosmos_db_storage.py b/tests/storage_cosmos/test_cosmos_db_storage.py index d71da51b2..e42df9406 100644 --- a/tests/storage_cosmos/test_cosmos_db_storage.py +++ b/tests/storage_cosmos/test_cosmos_db_storage.py @@ -194,7 +194,6 @@ async def test_cosmos_db_storage_flow_existing_container_and_persistence( @pytest.mark.cosmos class TestCosmosDBStorage(QuickCRUDStorageTests): - def get_compat_mode(self): return False @@ -253,9 +252,9 @@ async def test_cosmos_db_from_azure_cred(self): load_dotenv() cred = DefaultAzureCredential() - url = os.environ.get("TEST_COSMOS_DB_ENDPOINT") + cosmos_db_endpoint = os.environ.get("TEST_COSMOS_DB_ENDPOINT", "") config = CosmosDBStorageConfig( - url=url, + cosmos_db_endpoint=cosmos_db_endpoint, credential=cred, database_id="test-db", container_id="bot-storage", @@ -284,7 +283,6 @@ async def test_cosmos_db_from_azure_cred(self): # @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") @pytest.mark.cosmos class TestCosmosDBStorageInit: - def test_raises_error_when_suffix_provided_but_compat(self, config): config.auth_key = None config.compatibility_mode = True @@ -302,6 +300,17 @@ def test_raises_error_when_no_container_id_provided(self, config): with pytest.raises(ValueError): CosmosDBStorage(config) + def test_raises_error_when_no_endpoint_provided(self, config): + config.cosmos_db_endpoint = None + with pytest.raises(ValueError, match="cosmos_db_endpoint is required"): + CosmosDBStorage(config) + + def test_raises_error_when_no_auth_key_or_cred_provided(self, config): + config.auth_key = None + config.credential = None + with pytest.raises(ValueError, match="auth_key or credential is required"): + CosmosDBStorage(config) + @pytest.mark.asyncio @pytest.mark.parametrize("compat_mode", [True, False]) async def test_raises_error_different_partition_key(self, compat_mode):