From 7c33c0d5dcfd9d019bb94cd288e4715663fa10ed Mon Sep 17 00:00:00 2001 From: Tracy Boehrer Date: Mon, 9 Sep 2024 15:59:34 -0500 Subject: [PATCH] Updated azure-cosmos to 4.7.0, requiring dropped support for obsolete CosmosDBStorage class. --- .../botbuilder/azure/__init__.py | 4 +- .../azure/cosmosdb_partitioned_storage.py | 179 +++++---- .../botbuilder/azure/cosmosdb_storage.py | 378 ------------------ libraries/botbuilder-azure/setup.py | 2 +- .../tests/test_cosmos_partitioned_storage.py | 13 +- .../tests/test_cosmos_storage.py | 300 -------------- .../botbuilder/testing/storage_base_tests.py | 4 +- 7 files changed, 105 insertions(+), 775 deletions(-) delete mode 100644 libraries/botbuilder-azure/botbuilder/azure/cosmosdb_storage.py delete mode 100644 libraries/botbuilder-azure/tests/test_cosmos_storage.py diff --git a/libraries/botbuilder-azure/botbuilder/azure/__init__.py b/libraries/botbuilder-azure/botbuilder/azure/__init__.py index e625500a3..e6c70e7fc 100644 --- a/libraries/botbuilder-azure/botbuilder/azure/__init__.py +++ b/libraries/botbuilder-azure/botbuilder/azure/__init__.py @@ -7,10 +7,10 @@ from .about import __version__ from .azure_queue_storage import AzureQueueStorage -from .cosmosdb_storage import CosmosDbStorage, CosmosDbConfig, CosmosDbKeyEscape from .cosmosdb_partitioned_storage import ( CosmosDbPartitionedStorage, CosmosDbPartitionedConfig, + CosmosDbKeyEscape, ) from .blob_storage import BlobStorage, BlobStorageSettings @@ -18,8 +18,6 @@ "AzureQueueStorage", "BlobStorage", "BlobStorageSettings", - "CosmosDbStorage", - "CosmosDbConfig", "CosmosDbKeyEscape", "CosmosDbPartitionedStorage", "CosmosDbPartitionedConfig", diff --git a/libraries/botbuilder-azure/botbuilder/azure/cosmosdb_partitioned_storage.py b/libraries/botbuilder-azure/botbuilder/azure/cosmosdb_partitioned_storage.py index 982ac5974..cfe66f8d8 100644 --- a/libraries/botbuilder-azure/botbuilder/azure/cosmosdb_partitioned_storage.py +++ b/libraries/botbuilder-azure/botbuilder/azure/cosmosdb_partitioned_storage.py @@ -6,14 +6,14 @@ from typing import Dict, List from threading import Lock import json - +from hashlib import sha256 +from azure.core import MatchConditions from azure.cosmos import documents, http_constants from jsonpickle.pickler import Pickler from jsonpickle.unpickler import Unpickler import azure.cosmos.cosmos_client as cosmos_client # pylint: disable=no-name-in-module,import-error -import azure.cosmos.errors as cosmos_errors # pylint: disable=no-name-in-module,import-error +import azure.cosmos.exceptions as cosmos_exceptions from botbuilder.core.storage import Storage -from botbuilder.azure import CosmosDbKeyEscape class CosmosDbPartitionedConfig: @@ -63,6 +63,49 @@ def __init__( self.compatibility_mode = compatibility_mode or kwargs.get("compatibility_mode") +class CosmosDbKeyEscape: + @staticmethod + def sanitize_key( + key: str, key_suffix: str = "", compatibility_mode: bool = True + ) -> str: + """Return the sanitized key. + + Replace characters that are not allowed in keys in Cosmos. + + :param key: The provided key to be escaped. + :param key_suffix: The string to add a the end of all RowKeys. + :param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb + max key length of 255. This behavior can be overridden by setting + cosmosdb_partitioned_config.compatibility_mode to False. + :return str: + """ + # forbidden characters + bad_chars = ["\\", "?", "/", "#", "\t", "\n", "\r", "*"] + # replace those with with '*' and the + # Unicode code point of the character and return the new string + key = "".join(map(lambda x: "*" + str(ord(x)) if x in bad_chars else x, key)) + + if key_suffix is None: + key_suffix = "" + + return CosmosDbKeyEscape.truncate_key(f"{key}{key_suffix}", compatibility_mode) + + @staticmethod + def truncate_key(key: str, compatibility_mode: bool = True) -> str: + max_key_len = 255 + + if not compatibility_mode: + return key + + if len(key) > max_key_len: + aux_hash = sha256(key.encode("utf-8")) + aux_hex = aux_hash.hexdigest() + + key = key[0 : max_key_len - len(aux_hex)] + aux_hex + + return key + + class CosmosDbPartitionedStorage(Storage): """A CosmosDB based storage provider using partitioning for a bot.""" @@ -99,7 +142,8 @@ async def read(self, keys: List[str]) -> Dict[str, object]: :return dict: """ if not keys: - raise Exception("Keys are required when reading") + # No keys passed in, no result to return. Back-compat with original CosmosDBStorage. + return {} await self.initialize() @@ -111,8 +155,8 @@ async def read(self, keys: List[str]) -> Dict[str, object]: key, self.config.key_suffix, self.config.compatibility_mode ) - read_item_response = self.client.ReadItem( - self.__item_link(escaped_key), self.__get_partition_key(escaped_key) + read_item_response = self.container.read_item( + escaped_key, self.__get_partition_key(escaped_key) ) document_store_item = read_item_response if document_store_item: @@ -122,13 +166,8 @@ async def read(self, keys: List[str]) -> Dict[str, object]: # When an item is not found a CosmosException is thrown, but we want to # return an empty collection so in this instance we catch and do not rethrow. # Throw for any other exception. - except cosmos_errors.HTTPFailure as err: - if ( - err.status_code - == cosmos_errors.http_constants.StatusCodes.NOT_FOUND - ): - continue - raise err + except cosmos_exceptions.CosmosResourceNotFoundError: + continue except Exception as err: raise err return store_items @@ -162,20 +201,16 @@ async def write(self, changes: Dict[str, object]): if e_tag == "": raise Exception("cosmosdb_storage.write(): etag missing") - access_condition = { - "accessCondition": {"type": "IfMatch", "condition": e_tag} - } - options = ( - access_condition if e_tag != "*" and e_tag and e_tag != "" else None - ) + access_condition = e_tag != "*" and e_tag and e_tag != "" + try: - self.client.UpsertItem( - database_or_Container_link=self.__container_link, - document=doc, - options=options, + self.container.upsert_item( + body=doc, + etag=e_tag if access_condition else None, + match_condition=( + MatchConditions.IfNotModified if access_condition else None + ), ) - except cosmos_errors.HTTPFailure as err: - raise err except Exception as err: raise err @@ -192,69 +227,66 @@ async def delete(self, keys: List[str]): key, self.config.key_suffix, self.config.compatibility_mode ) try: - self.client.DeleteItem( - document_link=self.__item_link(escaped_key), - options=self.__get_partition_key(escaped_key), + self.container.delete_item( + escaped_key, + self.__get_partition_key(escaped_key), ) - except cosmos_errors.HTTPFailure as err: - if ( - err.status_code - == cosmos_errors.http_constants.StatusCodes.NOT_FOUND - ): - continue - raise err + except cosmos_exceptions.CosmosResourceNotFoundError: + continue except Exception as err: raise err async def initialize(self): if not self.container: if not self.client: + connection_policy = self.config.cosmos_client_options.get( + "connection_policy", documents.ConnectionPolicy() + ) + + # kwargs 'connection_verify' is to handle CosmosClient overwriting the + # ConnectionPolicy.DisableSSLVerification value. self.client = cosmos_client.CosmosClient( self.config.cosmos_db_endpoint, - {"masterKey": self.config.auth_key}, - self.config.cosmos_client_options.get("connection_policy", None), + self.config.auth_key, self.config.cosmos_client_options.get("consistency_level", None), + **{ + "connection_policy": connection_policy, + "connection_verify": not connection_policy.DisableSSLVerification, + }, ) if not self.database: with self.__lock: - try: - if not self.database: - self.database = self.client.CreateDatabase( - {"id": self.config.database_id} - ) - except cosmos_errors.HTTPFailure: - self.database = self.client.ReadDatabase( - "dbs/" + self.config.database_id + if not self.database: + self.database = self.client.create_database_if_not_exists( + self.config.database_id ) self.__get_or_create_container() def __get_or_create_container(self): with self.__lock: - container_def = { - "id": self.config.container_id, - "partitionKey": { - "paths": ["/id"], - "kind": documents.PartitionKind.Hash, - }, + partition_key = { + "paths": ["/id"], + "kind": documents.PartitionKind.Hash, } try: if not self.container: - self.container = self.client.CreateContainer( - "dbs/" + self.database["id"], - container_def, - {"offerThroughput": self.config.container_throughput}, + self.container = self.database.create_container( + self.config.container_id, + partition_key, + offer_throughput=self.config.container_throughput, ) - except cosmos_errors.HTTPFailure as err: + except cosmos_exceptions.CosmosHttpResponseError as err: if err.status_code == http_constants.StatusCodes.CONFLICT: - self.container = self.client.ReadContainer( - "dbs/" + self.database["id"] + "/colls/" + container_def["id"] + self.container = self.database.get_container_client( + self.config.container_id ) - if "partitionKey" not in self.container: + properties = self.container.read() + if "partitionKey" not in properties: self.compatability_mode_partition_key = True else: - paths = self.container["partitionKey"]["paths"] + paths = properties["partitionKey"]["paths"] if "/partitionKey" in paths: self.compatability_mode_partition_key = True elif "/id" not in paths: @@ -267,7 +299,7 @@ def __get_or_create_container(self): raise err def __get_partition_key(self, key: str) -> str: - return None if self.compatability_mode_partition_key else {"partitionKey": key} + return None if self.compatability_mode_partition_key else key @staticmethod def __create_si(result) -> object: @@ -303,28 +335,3 @@ def __create_dict(store_item: object) -> Dict: # loop through attributes and write and return a dict return json_dict - - def __item_link(self, identifier) -> str: - """Return the item link of a item in the container. - - :param identifier: - :return str: - """ - return self.__container_link + "/docs/" + identifier - - @property - def __container_link(self) -> str: - """Return the container link in the database. - - :param: - :return str: - """ - return self.__database_link + "/colls/" + self.config.container_id - - @property - def __database_link(self) -> str: - """Return the database link. - - :return str: - """ - return "dbs/" + self.config.database_id diff --git a/libraries/botbuilder-azure/botbuilder/azure/cosmosdb_storage.py b/libraries/botbuilder-azure/botbuilder/azure/cosmosdb_storage.py deleted file mode 100644 index 2e383666f..000000000 --- a/libraries/botbuilder-azure/botbuilder/azure/cosmosdb_storage.py +++ /dev/null @@ -1,378 +0,0 @@ -"""Implements a CosmosDB based storage provider. -""" - -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -from hashlib import sha256 -from typing import Dict, List -from threading import Semaphore -import json -import warnings -from jsonpickle.pickler import Pickler -from jsonpickle.unpickler import Unpickler -import azure.cosmos.cosmos_client as cosmos_client # pylint: disable=no-name-in-module,import-error -import azure.cosmos.errors as cosmos_errors # pylint: disable=no-name-in-module,import-error -from botbuilder.core.storage import Storage - - -class CosmosDbConfig: - """The class for CosmosDB configuration for the Azure Bot Framework.""" - - def __init__( - self, - endpoint: str = None, - masterkey: str = None, - database: str = None, - container: str = None, - partition_key: str = None, - database_creation_options: dict = None, - container_creation_options: dict = None, - **kwargs, - ): - """Create the Config object. - - :param endpoint: - :param masterkey: - :param database: - :param container: - :param filename: - :return CosmosDbConfig: - """ - self.__config_file = kwargs.get("filename") - if self.__config_file: - kwargs = json.load(open(self.__config_file)) - self.endpoint = endpoint or kwargs.get("endpoint") - self.masterkey = masterkey or kwargs.get("masterkey") - self.database = database or kwargs.get("database", "bot_db") - self.container = container or kwargs.get("container", "bot_container") - self.partition_key = partition_key or kwargs.get("partition_key") - self.database_creation_options = database_creation_options or kwargs.get( - "database_creation_options" - ) - self.container_creation_options = container_creation_options or kwargs.get( - "container_creation_options" - ) - - -class CosmosDbKeyEscape: - @staticmethod - def sanitize_key( - key: str, key_suffix: str = "", compatibility_mode: bool = True - ) -> str: - """Return the sanitized key. - - Replace characters that are not allowed in keys in Cosmos. - - :param key: The provided key to be escaped. - :param key_suffix: The string to add a the end of all RowKeys. - :param compatibility_mode: True if keys should be truncated in order to support previous CosmosDb - max key length of 255. This behavior can be overridden by setting - cosmosdb_partitioned_config.compatibility_mode to False. - :return str: - """ - # forbidden characters - bad_chars = ["\\", "?", "/", "#", "\t", "\n", "\r", "*"] - # replace those with with '*' and the - # Unicode code point of the character and return the new string - key = "".join(map(lambda x: "*" + str(ord(x)) if x in bad_chars else x, key)) - - if key_suffix is None: - key_suffix = "" - - return CosmosDbKeyEscape.truncate_key(f"{key}{key_suffix}", compatibility_mode) - - @staticmethod - def truncate_key(key: str, compatibility_mode: bool = True) -> str: - max_key_len = 255 - - if not compatibility_mode: - return key - - if len(key) > max_key_len: - aux_hash = sha256(key.encode("utf-8")) - aux_hex = aux_hash.hexdigest() - - key = key[0 : max_key_len - len(aux_hex)] + aux_hex - - return key - - -class CosmosDbStorage(Storage): - """A CosmosDB based storage provider for a bot.""" - - def __init__( - self, config: CosmosDbConfig, client: cosmos_client.CosmosClient = None - ): - """Create the storage object. - - :param config: - """ - super(CosmosDbStorage, self).__init__() - warnings.warn( - "CosmosDbStorage is obsolete. Use CosmosDbPartitionedStorage instead." - ) - self.config = config - self.client = client or cosmos_client.CosmosClient( - self.config.endpoint, {"masterKey": self.config.masterkey} - ) - # these are set by the functions that check - # the presence of the database and container or creates them - self.database = None - self.container = None - self._database_creation_options = config.database_creation_options - self._container_creation_options = config.container_creation_options - self.__semaphore = Semaphore() - - async def read(self, keys: List[str]) -> Dict[str, object]: - """Read storeitems from storage. - - :param keys: - :return dict: - """ - try: - # check if the database and container exists and if not create - if not self.__container_exists: - self.__create_db_and_container() - if keys: - # create the parameters object - parameters = [ - { - "name": f"@id{i}", - "value": f"{CosmosDbKeyEscape.sanitize_key(key)}", - } - for i, key in enumerate(keys) - ] - # get the names of the params - parameter_sequence = ",".join(param.get("name") for param in parameters) - # create the query - query = { - "query": f"SELECT c.id, c.realId, c.document, c._etag FROM c WHERE c.id in ({parameter_sequence})", - "parameters": parameters, - } - - if self.config.partition_key: - options = {"partitionKey": self.config.partition_key} - else: - options = {"enableCrossPartitionQuery": True} - - # run the query and store the results as a list - results = list( - self.client.QueryItems(self.__container_link, query, options) - ) - # return a dict with a key and an object - return {r.get("realId"): self.__create_si(r) for r in results} - - # No keys passed in, no result to return. - return {} - except TypeError as error: - raise error - - async def write(self, changes: Dict[str, object]): - """Save storeitems to storage. - - :param changes: - :return: - """ - if changes is None: - raise Exception("Changes are required when writing") - if not changes: - return - try: - # check if the database and container exists and if not create - if not self.__container_exists: - self.__create_db_and_container() - # iterate over the changes - for key, change in changes.items(): - # store the e_tag - e_tag = None - if isinstance(change, dict): - e_tag = change.get("e_tag", None) - elif hasattr(change, "e_tag"): - e_tag = change.e_tag - # create the new document - doc = { - "id": CosmosDbKeyEscape.sanitize_key(key), - "realId": key, - "document": self.__create_dict(change), - } - if e_tag == "": - raise Exception("cosmosdb_storage.write(): etag missing") - # the e_tag will be * for new docs so do an insert - if e_tag == "*" or not e_tag: - self.client.UpsertItem( - database_or_Container_link=self.__container_link, - document=doc, - options={"disableAutomaticIdGeneration": True}, - ) - # if we have an etag, do opt. concurrency replace - elif e_tag: - access_condition = {"type": "IfMatch", "condition": e_tag} - self.client.ReplaceItem( - document_link=self.__item_link( - CosmosDbKeyEscape.sanitize_key(key) - ), - new_document=doc, - options={"accessCondition": access_condition}, - ) - except Exception as error: - raise error - - async def delete(self, keys: List[str]): - """Remove storeitems from storage. - - :param keys: - :return: - """ - try: - # check if the database and container exists and if not create - if not self.__container_exists: - self.__create_db_and_container() - - options = {} - if self.config.partition_key: - options["partitionKey"] = self.config.partition_key - - # call the function for each key - for key in keys: - self.client.DeleteItem( - document_link=self.__item_link(CosmosDbKeyEscape.sanitize_key(key)), - options=options, - ) - # print(res) - except cosmos_errors.HTTPFailure as http_failure: - # print(h.status_code) - if http_failure.status_code != 404: - raise http_failure - except TypeError as error: - raise error - - def __create_si(self, result) -> object: - """Create an object from a result out of CosmosDB. - - :param result: - :return object: - """ - # get the document item from the result and turn into a dict - doc = result.get("document") - # read the e_tag from Cosmos - if result.get("_etag"): - doc["e_tag"] = result["_etag"] - - result_obj = Unpickler().restore(doc) - - # create and return the object - return result_obj - - def __create_dict(self, store_item: object) -> Dict: - """Return the dict of an object. - - This eliminates non_magic attributes and the e_tag. - - :param store_item: - :return dict: - """ - # read the content - json_dict = Pickler().flatten(store_item) - if "e_tag" in json_dict: - del json_dict["e_tag"] - - # loop through attributes and write and return a dict - return json_dict - - def __item_link(self, identifier) -> str: - """Return the item link of a item in the container. - - :param identifier: - :return str: - """ - return self.__container_link + "/docs/" + identifier - - @property - def __container_link(self) -> str: - """Return the container link in the database. - - :param: - :return str: - """ - return self.__database_link + "/colls/" + self.container - - @property - def __database_link(self) -> str: - """Return the database link. - - :return str: - """ - return "dbs/" + self.database - - @property - def __container_exists(self) -> bool: - """Return whether the database and container have been created. - - :return bool: - """ - return self.database and self.container - - def __create_db_and_container(self): - """Call the get or create methods.""" - with self.__semaphore: - db_id = self.config.database - container_name = self.config.container - self.database = self._get_or_create_database(self.client, db_id) - self.container = self._get_or_create_container(self.client, container_name) - - def _get_or_create_database( # pylint: disable=invalid-name - self, doc_client, id - ) -> str: - """Return the database link. - - Check if the database exists or create the database. - - :param doc_client: - :param id: - :return str: - """ - # query CosmosDB for a database with that name/id - dbs = list( - doc_client.QueryDatabases( - { - "query": "SELECT * FROM r WHERE r.id=@id", - "parameters": [{"name": "@id", "value": id}], - } - ) - ) - # if there are results, return the first (database names are unique) - if dbs: - return dbs[0]["id"] - - # create the database if it didn't exist - res = doc_client.CreateDatabase({"id": id}, self._database_creation_options) - return res["id"] - - def _get_or_create_container(self, doc_client, container) -> str: - """Return the container link. - - Check if the container exists or create the container. - - :param doc_client: - :param container: - :return str: - """ - # query CosmosDB for a container in the database with that name - containers = list( - doc_client.QueryContainers( - self.__database_link, - { - "query": "SELECT * FROM r WHERE r.id=@id", - "parameters": [{"name": "@id", "value": container}], - }, - ) - ) - # if there are results, return the first (container names are unique) - if containers: - return containers[0]["id"] - - # Create a container if it didn't exist - res = doc_client.CreateContainer( - self.__database_link, {"id": container}, self._container_creation_options - ) - return res["id"] diff --git a/libraries/botbuilder-azure/setup.py b/libraries/botbuilder-azure/setup.py index 04fd479cb..9c40b3ab5 100644 --- a/libraries/botbuilder-azure/setup.py +++ b/libraries/botbuilder-azure/setup.py @@ -5,7 +5,7 @@ from setuptools import setup REQUIRES = [ - "azure-cosmos==3.2.0", + "azure-cosmos==4.7.0", "azure-storage-blob==12.7.0", "azure-storage-queue==12.4.0", "botbuilder-schema==4.17.0", diff --git a/libraries/botbuilder-azure/tests/test_cosmos_partitioned_storage.py b/libraries/botbuilder-azure/tests/test_cosmos_partitioned_storage.py index cb6dd0822..d52733fd9 100644 --- a/libraries/botbuilder-azure/tests/test_cosmos_partitioned_storage.py +++ b/libraries/botbuilder-azure/tests/test_cosmos_partitioned_storage.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. -import azure.cosmos.errors as cosmos_errors +import azure.cosmos.exceptions as cosmos_exceptions from azure.cosmos import documents import pytest from botbuilder.azure import CosmosDbPartitionedStorage, CosmosDbPartitionedConfig @@ -27,8 +27,8 @@ async def reset(): storage = CosmosDbPartitionedStorage(get_settings()) await storage.initialize() try: - storage.client.DeleteDatabase(database_link="dbs/" + get_settings().database_id) - except cosmos_errors.HTTPFailure: + storage.client.delete_database(get_settings().database_id) + except cosmos_exceptions.HttpResponseError: pass @@ -99,9 +99,12 @@ async def test_passes_cosmos_client_options(self): client = CosmosDbPartitionedStorage(settings_with_options) await client.initialize() - assert client.client.connection_policy.DisableSSLVerification is True assert ( - client.client.default_headers["x-ms-consistency-level"] + client.client.client_connection.connection_policy.DisableSSLVerification + is True + ) + assert ( + client.client.client_connection.default_headers["x-ms-consistency-level"] == documents.ConsistencyLevel.Eventual ) diff --git a/libraries/botbuilder-azure/tests/test_cosmos_storage.py b/libraries/botbuilder-azure/tests/test_cosmos_storage.py deleted file mode 100644 index c66660857..000000000 --- a/libraries/botbuilder-azure/tests/test_cosmos_storage.py +++ /dev/null @@ -1,300 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. - -from unittest.mock import Mock -import azure.cosmos.errors as cosmos_errors -from azure.cosmos.cosmos_client import CosmosClient -import pytest -from botbuilder.core import StoreItem -from botbuilder.azure import CosmosDbStorage, CosmosDbConfig -from botbuilder.testing import StorageBaseTests - -# local cosmosdb emulator instance cosmos_db_config -COSMOS_DB_CONFIG = CosmosDbConfig( - endpoint="https://localhost:8081", - masterkey="C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==", - database="test-db", - container="bot-storage", -) -EMULATOR_RUNNING = False - - -def get_storage(): - return CosmosDbStorage(COSMOS_DB_CONFIG) - - -async def reset(): - storage = CosmosDbStorage(COSMOS_DB_CONFIG) - try: - storage.client.DeleteDatabase(database_link="dbs/" + COSMOS_DB_CONFIG.database) - except cosmos_errors.HTTPFailure: - pass - - -def get_mock_client(identifier: str = "1"): - # pylint: disable=attribute-defined-outside-init, invalid-name - mock = MockClient() - - mock.QueryDatabases = Mock(return_value=[]) - mock.QueryContainers = Mock(return_value=[]) - mock.CreateDatabase = Mock(return_value={"id": identifier}) - mock.CreateContainer = Mock(return_value={"id": identifier}) - - return mock - - -class MockClient(CosmosClient): - def __init__(self): # pylint: disable=super-init-not-called - pass - - -class SimpleStoreItem(StoreItem): - def __init__(self, counter=1, e_tag="*"): - super(SimpleStoreItem, self).__init__() - self.counter = counter - self.e_tag = e_tag - - -class TestCosmosDbStorageConstructor: - @pytest.mark.asyncio - async def test_cosmos_storage_init_should_error_without_cosmos_db_config(self): - try: - CosmosDbStorage(CosmosDbConfig()) - except Exception as error: - assert error - - @pytest.mark.asyncio - async def test_creation_request_options_are_being_called(self): - # pylint: disable=protected-access - test_config = CosmosDbConfig( - endpoint="https://localhost:8081", - masterkey="C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==", - database="test-db", - container="bot-storage", - database_creation_options={"OfferThroughput": 1000}, - container_creation_options={"OfferThroughput": 500}, - ) - - test_id = "1" - client = get_mock_client(identifier=test_id) - storage = CosmosDbStorage(test_config, client) - storage.database = test_id - - assert storage._get_or_create_database(doc_client=client, id=test_id), test_id - client.CreateDatabase.assert_called_with( - {"id": test_id}, test_config.database_creation_options - ) - assert storage._get_or_create_container( - doc_client=client, container=test_id - ), test_id - client.CreateContainer.assert_called_with( - "dbs/" + test_id, {"id": test_id}, test_config.container_creation_options - ) - - -class TestCosmosDbStorageBaseStorageTests: - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_return_empty_object_when_reading_unknown_key(self): - await reset() - - test_ran = await StorageBaseTests.return_empty_object_when_reading_unknown_key( - get_storage() - ) - - assert test_ran - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_handle_null_keys_when_reading(self): - await reset() - - test_ran = await StorageBaseTests.handle_null_keys_when_reading(get_storage()) - - assert test_ran - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_handle_null_keys_when_writing(self): - await reset() - - test_ran = await StorageBaseTests.handle_null_keys_when_writing(get_storage()) - - assert test_ran - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_does_not_raise_when_writing_no_items(self): - await reset() - - test_ran = await StorageBaseTests.does_not_raise_when_writing_no_items( - get_storage() - ) - - assert test_ran - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_create_object(self): - await reset() - - test_ran = await StorageBaseTests.create_object(get_storage()) - - assert test_ran - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_handle_crazy_keys(self): - await reset() - - test_ran = await StorageBaseTests.handle_crazy_keys(get_storage()) - - assert test_ran - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_update_object(self): - await reset() - - test_ran = await StorageBaseTests.update_object(get_storage()) - - assert test_ran - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_delete_object(self): - await reset() - - test_ran = await StorageBaseTests.delete_object(get_storage()) - - assert test_ran - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_perform_batch_operations(self): - await reset() - - test_ran = await StorageBaseTests.perform_batch_operations(get_storage()) - - assert test_ran - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_proceeds_through_waterfall(self): - await reset() - - test_ran = await StorageBaseTests.proceeds_through_waterfall(get_storage()) - - assert test_ran - - -class TestCosmosDbStorage: - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_cosmos_storage_init_should_work_with_just_endpoint_and_key(self): - storage = CosmosDbStorage( - CosmosDbConfig( - endpoint=COSMOS_DB_CONFIG.endpoint, masterkey=COSMOS_DB_CONFIG.masterkey - ) - ) - await storage.write({"user": SimpleStoreItem()}) - data = await storage.read(["user"]) - assert "user" in data - assert data["user"].counter == 1 - assert len(data.keys()) == 1 - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_cosmos_storage_read_update_should_return_new_etag(self): - await reset() - storage = CosmosDbStorage(COSMOS_DB_CONFIG) - await storage.write({"test": SimpleStoreItem(counter=1)}) - data_result = await storage.read(["test"]) - data_result["test"].counter = 2 - await storage.write(data_result) - data_updated = await storage.read(["test"]) - assert data_updated["test"].counter == 2 - assert data_updated["test"].e_tag != data_result["test"].e_tag - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_cosmos_storage_read_with_invalid_key_should_return_empty_dict(self): - await reset() - storage = CosmosDbStorage(COSMOS_DB_CONFIG) - data = await storage.read(["test"]) - - assert isinstance(data, dict) - assert not data.keys() - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_cosmos_storage_write_should_overwrite_when_new_e_tag_is_an_asterisk( - self, - ): - await reset() - storage = CosmosDbStorage(COSMOS_DB_CONFIG) - await storage.write({"user": SimpleStoreItem()}) - - await storage.write({"user": SimpleStoreItem(counter=10, e_tag="*")}) - data = await storage.read(["user"]) - assert data["user"].counter == 10 - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_cosmos_storage_delete_should_delete_multiple_values_when_given_multiple_valid_keys( - self, - ): - await reset() - storage = CosmosDbStorage(COSMOS_DB_CONFIG) - await storage.write({"test": SimpleStoreItem(), "test2": SimpleStoreItem(2)}) - - await storage.delete(["test", "test2"]) - data = await storage.read(["test", "test2"]) - assert not data.keys() - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_cosmos_storage_delete_should_delete_values_when_given_multiple_valid_keys_and_ignore_other_data( - self, - ): - await reset() - storage = CosmosDbStorage(COSMOS_DB_CONFIG) - await storage.write( - { - "test": SimpleStoreItem(), - "test2": SimpleStoreItem(counter=2), - "test3": SimpleStoreItem(counter=3), - } - ) - - await storage.delete(["test", "test2"]) - data = await storage.read(["test", "test2", "test3"]) - assert len(data.keys()) == 1 - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_cosmos_storage_delete_invalid_key_should_do_nothing_and_not_affect_cached_data( - self, - ): - await reset() - storage = CosmosDbStorage(COSMOS_DB_CONFIG) - await storage.write({"test": SimpleStoreItem()}) - - await storage.delete(["foo"]) - data = await storage.read(["test"]) - assert len(data.keys()) == 1 - data = await storage.read(["foo"]) - assert not data.keys() - - @pytest.mark.skipif(not EMULATOR_RUNNING, reason="Needs the emulator to run.") - @pytest.mark.asyncio - async def test_cosmos_storage_delete_invalid_keys_should_do_nothing_and_not_affect_cached_data( - self, - ): - await reset() - storage = CosmosDbStorage(COSMOS_DB_CONFIG) - await storage.write({"test": SimpleStoreItem()}) - - await storage.delete(["foo", "bar"]) - data = await storage.read(["test"]) - assert len(data.keys()) == 1 diff --git a/libraries/botbuilder-testing/botbuilder/testing/storage_base_tests.py b/libraries/botbuilder-testing/botbuilder/testing/storage_base_tests.py index e196099a0..e374a3401 100644 --- a/libraries/botbuilder-testing/botbuilder/testing/storage_base_tests.py +++ b/libraries/botbuilder-testing/botbuilder/testing/storage_base_tests.py @@ -24,7 +24,7 @@ async def test_handle_null_keys_when_reading(self): assert test_ran """ import pytest -from botbuilder.azure import CosmosDbStorage +from botbuilder.azure import CosmosDbPartitionedStorage from botbuilder.core import ( ConversationState, TurnContext, @@ -57,7 +57,7 @@ async def return_empty_object_when_reading_unknown_key(storage) -> bool: @staticmethod async def handle_null_keys_when_reading(storage) -> bool: - if isinstance(storage, (CosmosDbStorage, MemoryStorage)): + if isinstance(storage, (CosmosDbPartitionedStorage, MemoryStorage)): result = await storage.read(None) assert len(result.keys()) == 0 # Catch-all