From 15dcceb042c2dabf04004a6f9af27ef83a89a7c7 Mon Sep 17 00:00:00 2001 From: annatisch Date: Fri, 20 Aug 2021 07:01:18 -0700 Subject: [PATCH] Client Constructor (#20310) * Removed some stuff * Looking at constructors * Updated request * Added client close * working client creation Co-authored-by: simorenoh --- .../azure/cosmos/aio/_asynchronous_request.py | 4 +- .../aio/_cosmos_client_connection_async.py | 2466 +---------------- .../aio/_global_endpoint_manager_async.py | 24 +- .../azure/cosmos/aio/_retry_utility.py | 196 ++ .../azure/cosmos/aio/container_async.py | 802 ------ .../azure/cosmos/aio/cosmos_client.py | 163 ++ .../azure/cosmos/aio/cosmos_client_async.py | 456 --- .../azure/cosmos/aio/database_async.py | 768 ----- .../azure/cosmos/cosmos_client.py | 3 + .../azure-cosmos/samples/simon_testfile.py | 31 +- 10 files changed, 395 insertions(+), 4518 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/aio/_retry_utility.py delete mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/aio/container_async.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/aio/cosmos_client.py delete mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/aio/cosmos_client_async.py delete mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/aio/database_async.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py index 986cb130b9ae..a1afc3f39e93 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py @@ -31,7 +31,7 @@ from .. import exceptions from .. import http_constants -from .. import _retry_utility +from . import _retry_utility from .._synchronized_request import _request_body_from_data @@ -174,7 +174,7 @@ async def AsynchronousRequest( request.headers[http_constants.HttpHeaders.ContentLength] = 0 # Pass _Request function with it's parameters to retry_utility's Execute method that wraps the call with retries - return await _retry_utility.Execute( + return await _retry_utility.ExecuteAsync( client, global_endpoint_manager, _Request, diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index c23ac5bde52e..645cb0b47e53 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -58,16 +58,15 @@ from .. import _global_endpoint_manager as global_endpoint_manager from . import _global_endpoint_manager_async as global_endpoint_manager_async from .._routing import routing_map_provider -from .._retry_utility import ConnectionRetryPolicy +from ._retry_utility import ConnectionRetryPolicy from .. import _session from .. import _utils from ..partition_key import _Undefined, _Empty -from .._cosmos_client_connection import CosmosClientConnection as BaseCosmosConnection # pylint: disable=protected-access -class CosmosClientConnection(BaseCosmosConnection, object): # pylint: disable=too-many-public-methods,too-many-instance-attributes +class CosmosClientConnection(object): # pylint: disable=too-many-public-methods,too-many-instance-attributes """Represents a document client. Provides a client-side logical representation of the Azure Cosmos @@ -114,12 +113,6 @@ def __init__( The default consistency policy for client operations. """ - # super(CosmosClientConnection, self).__init__( - # url_connection=url_connection, - # auth=auth, - # connection_policy=connection_policy, - # consistency_level=consistency_level, - # **kwargs) self.url_connection = url_connection self.master_key = None @@ -204,12 +197,9 @@ def __init__( HttpLoggingPolicy(**kwargs), ] - # print(asyncio.get_event_loop()) - database_account = self._global_endpoint_manager._GetDatabaseAccount(**kwargs) - self._global_endpoint_manager.force_refresh(database_account) - transport = kwargs.pop("transport", None) self.pipeline_client = AsyncPipelineClient(base_url=url_connection, transport=transport, policies=policies) + self._setup_kwargs = kwargs # Query compatibility mode. # Allows to specify compatibility mode used by client when making query requests. Should be removed when @@ -219,2449 +209,7 @@ def __init__( # Routing map provider self._routing_map_provider = routing_map_provider.SmartRoutingMapProvider(self) - @property - def Session(self): - """Gets the session object from the client. """ - return self.session - - @Session.setter - def Session(self, session): - """Sets a session object on the document client. - - This will override the existing session - """ - self.session = session - - @property - def WriteEndpoint(self): - """Gets the curent write endpoint for a geo-replicated database account. - """ - return self._global_endpoint_manager.get_write_endpoint() - - @property - def ReadEndpoint(self): - """Gets the curent read endpoint for a geo-replicated database account. - """ - return self._global_endpoint_manager.get_read_endpoint() - - def RegisterPartitionResolver(self, database_link, partition_resolver): - """Registers the partition resolver associated with the database link - - :param str database_link: - Database Self Link or ID based link. - :param object partition_resolver: - An instance of PartitionResolver. - - """ - if not database_link: - raise ValueError("database_link is None or empty.") - - if partition_resolver is None: - raise ValueError("partition_resolver is None.") - - self.partition_resolvers = {base.TrimBeginningAndEndingSlashes(database_link): partition_resolver} - - def GetPartitionResolver(self, database_link): - """Gets the partition resolver associated with the database link - - :param str database_link: - Database self link or ID based link. - - :return: - An instance of PartitionResolver. - :rtype: object - - """ - if not database_link: - raise ValueError("database_link is None or empty.") - - return self.partition_resolvers.get(base.TrimBeginningAndEndingSlashes(database_link)) - - async def CreateDatabase(self, database, options=None, **kwargs): - """Creates a database. - - :param dict database: - The Azure Cosmos database to create. - :param dict options: - The request options for the request. - - :return: - The Database that was created. - :rtype: dict - - """ - if options is None: - options = {} - - CosmosClientConnection.__ValidateResource(database) - path = "/dbs" - return await self.Create(database, path, "dbs", None, None, options, **kwargs) - - def ReadDatabase(self, database_link, options=None, **kwargs): - """Reads a database. - - :param str database_link: - The link to the database. - :param dict options: - The request options for the request. - - :return: - The Database that was read. - :rtype: dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(database_link) - database_id = base.GetResourceIdOrFullNameFromLink(database_link) - return self.Read(path, "dbs", database_id, None, options, **kwargs) - - def ReadDatabases(self, options=None, **kwargs): - """Reads all databases. - - :param dict options: - The request options for the request. - - :return: - Query Iterable of Databases. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - return self.QueryDatabases(None, options, **kwargs) - - def QueryDatabases(self, query, options=None, **kwargs): - """Queries databases. - - :param (str or dict) query: - :param dict options: - The request options for the request. - - :return: Query Iterable of Databases. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - def fetch_fn(options): - return ( - self.__QueryFeed( - "/dbs", "dbs", "", lambda r: r["Databases"], - lambda _, b: b, query, options, **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable - ) - - def ReadContainers(self, database_link, options=None, **kwargs): - """Reads all collections in a database. - - :param str database_link: - The link to the database. - :param dict options: - The request options for the request. - - :return: Query Iterable of Collections. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - return self.QueryContainers(database_link, None, options, **kwargs) - - def QueryContainers(self, database_link, query, options=None, **kwargs): - """Queries collections in a database. - - :param str database_link: - The link to the database. - :param (str or dict) query: - :param dict options: - The request options for the request. - - :return: Query Iterable of Collections. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(database_link, "colls") - database_id = base.GetResourceIdOrFullNameFromLink(database_link) - - def fetch_fn(options): - return ( - self.__QueryFeed( - path, "colls", database_id, lambda r: r["DocumentCollections"], - lambda _, body: body, query, options, **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable - ) - - def CreateContainer(self, database_link, collection, options=None, **kwargs): - """Creates a collection in a database. - - :param str database_link: - The link to the database. - :param dict collection: - The Azure Cosmos collection to create. - :param dict options: - The request options for the request. - - :return: The Collection that was created. - :rtype: dict - - """ - if options is None: - options = {} - - CosmosClientConnection.__ValidateResource(collection) - path = base.GetPathFromLink(database_link, "colls") - database_id = base.GetResourceIdOrFullNameFromLink(database_link) - return self.Create(collection, path, "colls", database_id, None, options, **kwargs) - - def ReplaceContainer(self, collection_link, collection, options=None, **kwargs): - """Replaces a collection and return it. - - :param str collection_link: - The link to the collection entity. - :param dict collection: - The collection to be used. - :param dict options: - The request options for the request. - - :return: - The new Collection. - :rtype: - dict - - """ - if options is None: - options = {} - - CosmosClientConnection.__ValidateResource(collection) - path = base.GetPathFromLink(collection_link) - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return self.Replace(collection, path, "colls", collection_id, None, options, **kwargs) - - def ReadContainer(self, collection_link, options=None, **kwargs): - """Reads a collection. - - :param str collection_link: - The link to the document collection. - :param dict options: - The request options for the request. - - :return: - The read Collection. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(collection_link) - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return self.Read(path, "colls", collection_id, None, options, **kwargs) - - def CreateUser(self, database_link, user, options=None, **kwargs): - """Creates a user. - - :param str database_link: - The link to the database. - :param dict user: - The Azure Cosmos user to create. - :param dict options: - The request options for the request. - - :return: - The created User. - :rtype: - dict - - """ - if options is None: - options = {} - - database_id, path = self._GetDatabaseIdWithPathForUser(database_link, user) - return self.Create(user, path, "users", database_id, None, options, **kwargs) - - def UpsertUser(self, database_link, user, options=None, **kwargs): - """Upserts a user. - - :param str database_link: - The link to the database. - :param dict user: - The Azure Cosmos user to upsert. - :param dict options: - The request options for the request. - - :return: - The upserted User. - :rtype: dict - """ - if options is None: - options = {} - - database_id, path = self._GetDatabaseIdWithPathForUser(database_link, user) - return self.Upsert(user, path, "users", database_id, None, options, **kwargs) - - def _GetDatabaseIdWithPathForUser(self, database_link, user): # pylint: disable=no-self-use - CosmosClientConnection.__ValidateResource(user) - path = base.GetPathFromLink(database_link, "users") - database_id = base.GetResourceIdOrFullNameFromLink(database_link) - return database_id, path - - def ReadUser(self, user_link, options=None, **kwargs): - """Reads a user. - - :param str user_link: - The link to the user entity. - :param dict options: - The request options for the request. - - :return: - The read User. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(user_link) - user_id = base.GetResourceIdOrFullNameFromLink(user_link) - return self.Read(path, "users", user_id, None, options, **kwargs) - - def ReadUsers(self, database_link, options=None, **kwargs): - """Reads all users in a database. - - :params str database_link: - The link to the database. - :params dict options: - The request options for the request. - :return: - Query iterable of Users. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - return self.QueryUsers(database_link, None, options, **kwargs) - - def QueryUsers(self, database_link, query, options=None, **kwargs): - """Queries users in a database. - - :param str database_link: - The link to the database. - :param (str or dict) query: - :param dict options: - The request options for the request. - - :return: - Query Iterable of Users. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(database_link, "users") - database_id = base.GetResourceIdOrFullNameFromLink(database_link) - - def fetch_fn(options): - return ( - self.__QueryFeed( - path, "users", database_id, lambda r: r["Users"], - lambda _, b: b, query, options, **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable - ) - - def DeleteDatabase(self, database_link, options=None, **kwargs): - """Deletes a database. - - :param str database_link: - The link to the database. - :param dict options: - The request options for the request. - - :return: - The deleted Database. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(database_link) - database_id = base.GetResourceIdOrFullNameFromLink(database_link) - return self.DeleteResource(path, "dbs", database_id, None, options, **kwargs) - - def CreatePermission(self, user_link, permission, options=None, **kwargs): - """Creates a permission for a user. - - :param str user_link: - The link to the user entity. - :param dict permission: - The Azure Cosmos user permission to create. - :param dict options: - The request options for the request. - - :return: - The created Permission. - :rtype: - dict - - """ - if options is None: - options = {} - - path, user_id = self._GetUserIdWithPathForPermission(permission, user_link) - return self.Create(permission, path, "permissions", user_id, None, options, **kwargs) - - def UpsertPermission(self, user_link, permission, options=None, **kwargs): - """Upserts a permission for a user. - - :param str user_link: - The link to the user entity. - :param dict permission: - The Azure Cosmos user permission to upsert. - :param dict options: - The request options for the request. - - :return: - The upserted permission. - :rtype: - dict - - """ - if options is None: - options = {} - - path, user_id = self._GetUserIdWithPathForPermission(permission, user_link) - return self.Upsert(permission, path, "permissions", user_id, None, options, **kwargs) - - def _GetUserIdWithPathForPermission(self, permission, user_link): # pylint: disable=no-self-use - CosmosClientConnection.__ValidateResource(permission) - path = base.GetPathFromLink(user_link, "permissions") - user_id = base.GetResourceIdOrFullNameFromLink(user_link) - return path, user_id - - def ReadPermission(self, permission_link, options=None, **kwargs): - """Reads a permission. - - :param str permission_link: - The link to the permission. - :param dict options: - The request options for the request. - - :return: - The read permission. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(permission_link) - permission_id = base.GetResourceIdOrFullNameFromLink(permission_link) - return self.Read(path, "permissions", permission_id, None, options, **kwargs) - - def ReadPermissions(self, user_link, options=None, **kwargs): - """Reads all permissions for a user. - - :param str user_link: - The link to the user entity. - :param dict options: - The request options for the request. - - :return: - Query Iterable of Permissions. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - return self.QueryPermissions(user_link, None, options, **kwargs) - - def QueryPermissions(self, user_link, query, options=None, **kwargs): - """Queries permissions for a user. - - :param str user_link: - The link to the user entity. - :param (str or dict) query: - :param dict options: - The request options for the request. - - :return: - Query Iterable of Permissions. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(user_link, "permissions") - user_id = base.GetResourceIdOrFullNameFromLink(user_link) - - def fetch_fn(options): - return ( - self.__QueryFeed( - path, "permissions", user_id, lambda r: r["Permissions"], lambda _, b: b, query, options, **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable - ) - - def ReplaceUser(self, user_link, user, options=None, **kwargs): - """Replaces a user and return it. - - :param str user_link: - The link to the user entity. - :param dict user: - :param dict options: - The request options for the request. - - :return: - The new User. - :rtype: - dict - - """ - if options is None: - options = {} - - CosmosClientConnection.__ValidateResource(user) - path = base.GetPathFromLink(user_link) - user_id = base.GetResourceIdOrFullNameFromLink(user_link) - return self.Replace(user, path, "users", user_id, None, options, **kwargs) - - def DeleteUser(self, user_link, options=None, **kwargs): - """Deletes a user. - - :param str user_link: - The link to the user entity. - :param dict options: - The request options for the request. - - :return: - The deleted user. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(user_link) - user_id = base.GetResourceIdOrFullNameFromLink(user_link) - return self.DeleteResource(path, "users", user_id, None, options, **kwargs) - - def ReplacePermission(self, permission_link, permission, options=None, **kwargs): - """Replaces a permission and return it. - - :param str permission_link: - The link to the permission. - :param dict permission: - :param dict options: - The request options for the request. - - :return: - The new Permission. - :rtype: - dict - - """ - if options is None: - options = {} - - CosmosClientConnection.__ValidateResource(permission) - path = base.GetPathFromLink(permission_link) - permission_id = base.GetResourceIdOrFullNameFromLink(permission_link) - return self.Replace(permission, path, "permissions", permission_id, None, options, **kwargs) - - def DeletePermission(self, permission_link, options=None, **kwargs): - """Deletes a permission. - - :param str permission_link: - The link to the permission. - :param dict options: - The request options for the request. - - :return: - The deleted Permission. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(permission_link) - permission_id = base.GetResourceIdOrFullNameFromLink(permission_link) - return self.DeleteResource(path, "permissions", permission_id, None, options, **kwargs) - - def ReadItems(self, collection_link, feed_options=None, response_hook=None, **kwargs): - """Reads all documents in a collection. - - :param str collection_link: - The link to the document collection. - :param dict feed_options: - - :return: - Query Iterable of Documents. - :rtype: - query_iterable.QueryIterable - - """ - if feed_options is None: - feed_options = {} - - return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, **kwargs) - - def QueryItems( - self, - database_or_container_link, - query, - options=None, - partition_key=None, - response_hook=None, - **kwargs - ): - """Queries documents in a collection. - - :param str database_or_container_link: - The link to the database when using partitioning, otherwise link to the document collection. - :param (str or dict) query: - :param dict options: - The request options for the request. - :param str partition_key: - Partition key for the query(default value None) - :param response_hook: - A callable invoked with the response metadata - - :return: - Query Iterable of Documents. - :rtype: - query_iterable.QueryIterable - - """ - database_or_container_link = base.TrimBeginningAndEndingSlashes(database_or_container_link) - - if options is None: - options = {} - - if base.IsDatabaseLink(database_or_container_link): - return AsyncItemPaged( - self, - query, - options, - database_link=database_or_container_link, - partition_key=partition_key, - page_iterator_class=query_iterable.QueryIterable - ) - - path = base.GetPathFromLink(database_or_container_link, "docs") - collection_id = base.GetResourceIdOrFullNameFromLink(database_or_container_link) - - def fetch_fn(options): - return ( - self.__QueryFeed( - path, - "docs", - collection_id, - lambda r: r["Documents"], - lambda _, b: b, - query, - options, - response_hook=response_hook, - **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, - query, - options, - fetch_function=fetch_fn, - collection_link=database_or_container_link, - page_iterator_class=query_iterable.QueryIterable - ) - - def QueryItemsChangeFeed(self, collection_link, options=None, response_hook=None, **kwargs): - """Queries documents change feed in a collection. - - :param str collection_link: - The link to the document collection. - :param dict options: - The request options for the request. - options may also specify partition key range id. - :param response_hook: - A callable invoked with the response metadata - - :return: - Query Iterable of Documents. - :rtype: - query_iterable.QueryIterable - - """ - - partition_key_range_id = None - if options is not None and "partitionKeyRangeId" in options: - partition_key_range_id = options["partitionKeyRangeId"] - - return self._QueryChangeFeed( - collection_link, "Documents", options, partition_key_range_id, response_hook=response_hook, **kwargs - ) - - def _QueryChangeFeed( - self, collection_link, resource_type, options=None, partition_key_range_id=None, response_hook=None, **kwargs - ): - """Queries change feed of a resource in a collection. - - :param str collection_link: - The link to the document collection. - :param str resource_type: - The type of the resource. - :param dict options: - The request options for the request. - :param str partition_key_range_id: - Specifies partition key range id. - :param response_hook: - A callable invoked with the response metadata - - :return: - Query Iterable of Documents. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - options["changeFeed"] = True - - resource_key_map = {"Documents": "docs"} - - # For now, change feed only supports Documents and Partition Key Range resouce type - if resource_type not in resource_key_map: - raise NotImplementedError(resource_type + " change feed query is not supported.") - - resource_key = resource_key_map[resource_type] - path = base.GetPathFromLink(collection_link, resource_key) - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - - def fetch_fn(options): - return ( - self.__QueryFeed( - path, - resource_key, - collection_id, - lambda r: r[resource_type], - lambda _, b: b, - None, - options, - partition_key_range_id, - response_hook=response_hook, - **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, - None, - options, - fetch_function=fetch_fn, - collection_link=collection_link, - page_iterator_class=query_iterable.QueryIterable - ) - - def _ReadPartitionKeyRanges(self, collection_link, feed_options=None, **kwargs): - """Reads Partition Key Ranges. - - :param str collection_link: - The link to the document collection. - :param dict feed_options: - - :return: - Query Iterable of PartitionKeyRanges. - :rtype: - query_iterable.QueryIterable - - """ - if feed_options is None: - feed_options = {} - - return self._QueryPartitionKeyRanges(collection_link, None, feed_options, **kwargs) - - def _QueryPartitionKeyRanges(self, collection_link, query, options=None, **kwargs): - """Queries Partition Key Ranges in a collection. - - :param str collection_link: - The link to the document collection. - :param (str or dict) query: - :param dict options: - The request options for the request. - - :return: - Query Iterable of PartitionKeyRanges. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(collection_link, "pkranges") - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - - def fetch_fn(options): - return ( - self.__QueryFeed( - path, "pkranges", collection_id, lambda r: r["PartitionKeyRanges"], - lambda _, b: b, query, options, **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable - ) - - def CreateItem(self, database_or_container_link, document, options=None, **kwargs): - """Creates a document in a collection. - - :param str database_or_container_link: - The link to the database when using partitioning, otherwise link to the document collection. - :param dict document: - The Azure Cosmos document to create. - :param dict options: - The request options for the request. - :param bool options['disableAutomaticIdGeneration']: - Disables the automatic id generation. If id is missing in the body and this - option is true, an error will be returned. - - :return: - The created Document. - :rtype: - dict - - """ - # Python's default arguments are evaluated once when the function is defined, - # not each time the function is called (like it is in say, Ruby). This means - # that if you use a mutable default argument and mutate it, you will and have - # mutated that object for all future calls to the function as well. So, using - # a non-mutable default in this case(None) and assigning an empty dict(mutable) - # inside the method For more details on this gotcha, please refer - # http://docs.python-guide.org/en/latest/writing/gotchas/ - if options is None: - options = {} - - # We check the link to be document collection link since it can be database - # link in case of client side partitioning - collection_id, document, path = self._GetContainerIdWithPathForItem( - database_or_container_link, document, options - ) - - if base.IsItemContainerLink(database_or_container_link): - options = self._AddPartitionKey(database_or_container_link, document, options) - - return self.Create(document, path, "docs", collection_id, None, options, **kwargs) - - def UpsertItem(self, database_or_container_link, document, options=None, **kwargs): - """Upserts a document in a collection. - - :param str database_or_container_link: - The link to the database when using partitioning, otherwise link to the document collection. - :param dict document: - The Azure Cosmos document to upsert. - :param dict options: - The request options for the request. - :param bool options['disableAutomaticIdGeneration']: - Disables the automatic id generation. If id is missing in the body and this - option is true, an error will be returned. - - :return: - The upserted Document. - :rtype: - dict - - """ - # Python's default arguments are evaluated once when the function is defined, - # not each time the function is called (like it is in say, Ruby). This means - # that if you use a mutable default argument and mutate it, you will and have - # mutated that object for all future calls to the function as well. So, using - # a non-mutable deafult in this case(None) and assigning an empty dict(mutable) - # inside the method For more details on this gotcha, please refer - # http://docs.python-guide.org/en/latest/writing/gotchas/ - if options is None: - options = {} - - # We check the link to be document collection link since it can be database - # link in case of client side partitioning - if base.IsItemContainerLink(database_or_container_link): - options = self._AddPartitionKey(database_or_container_link, document, options) - - collection_id, document, path = self._GetContainerIdWithPathForItem( - database_or_container_link, document, options - ) - return self.Upsert(document, path, "docs", collection_id, None, options, **kwargs) - - PartitionResolverErrorMessage = ( - "Couldn't find any partition resolvers for the database link provided. " - + "Ensure that the link you used when registering the partition resolvers " - + "matches the link provided or you need to register both types of database " - + "link(self link as well as ID based link)." - ) - - # Gets the collection id and path for the document - def _GetContainerIdWithPathForItem(self, database_or_container_link, document, options): - - if not database_or_container_link: - raise ValueError("database_or_container_link is None or empty.") - - if document is None: - raise ValueError("document is None.") - - CosmosClientConnection.__ValidateResource(document) - document = document.copy() - if not document.get("id") and not options.get("disableAutomaticIdGeneration"): - document["id"] = base.GenerateGuidId() - - collection_link = database_or_container_link - - if base.IsDatabaseLink(database_or_container_link): - partition_resolver = self.GetPartitionResolver(database_or_container_link) - - if partition_resolver is not None: - collection_link = partition_resolver.ResolveForCreate(document) - else: - raise ValueError(CosmosClientConnection.PartitionResolverErrorMessage) - - path = base.GetPathFromLink(collection_link, "docs") - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return collection_id, document, path - - def ReadItem(self, document_link, options=None, **kwargs): - """Reads a document. - - :param str document_link: - The link to the document. - :param dict options: - The request options for the request. - - :return: - The read Document. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(document_link) - document_id = base.GetResourceIdOrFullNameFromLink(document_link) - return self.Read(path, "docs", document_id, None, options, **kwargs) - - async def ReadItemAsync(self, document_link, options=None, **kwargs): - """Reads a document. - - :param str document_link: - The link to the document. - :param dict options: - The request options for the request. - - :return: - The read Document. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(document_link) - document_id = base.GetResourceIdOrFullNameFromLink(document_link) - return await self.ReadAsync(path, "docs", document_id, None, options, **kwargs) - - def ReadTriggers(self, collection_link, options=None, **kwargs): - """Reads all triggers in a collection. - - :param str collection_link: - The link to the document collection. - :param dict options: - The request options for the request. - - :return: - Query Iterable of Triggers. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - return self.QueryTriggers(collection_link, None, options, **kwargs) - - def QueryTriggers(self, collection_link, query, options=None, **kwargs): - """Queries triggers in a collection. - - :param str collection_link: - The link to the document collection. - :param (str or dict) query: - :param dict options: - The request options for the request. - - :return: - Query Iterable of Triggers. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(collection_link, "triggers") - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - - def fetch_fn(options): - return ( - self.__QueryFeed( - path, "triggers", collection_id, lambda r: r["Triggers"], lambda _, b: b, query, options, **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable - ) - - def CreateTrigger(self, collection_link, trigger, options=None, **kwargs): - """Creates a trigger in a collection. - - :param str collection_link: - The link to the document collection. - :param dict trigger: - :param dict options: - The request options for the request. - - :return: - The created Trigger. - :rtype: - dict - - """ - if options is None: - options = {} - - collection_id, path, trigger = self._GetContainerIdWithPathForTrigger(collection_link, trigger) - return self.Create(trigger, path, "triggers", collection_id, None, options, **kwargs) - - def UpsertTrigger(self, collection_link, trigger, options=None, **kwargs): - """Upserts a trigger in a collection. - - :param str collection_link: - The link to the document collection. - :param dict trigger: - :param dict options: - The request options for the request. - - :return: - The upserted Trigger. - :rtype: - dict - - """ - if options is None: - options = {} - - collection_id, path, trigger = self._GetContainerIdWithPathForTrigger(collection_link, trigger) - return self.Upsert(trigger, path, "triggers", collection_id, None, options, **kwargs) - - def _GetContainerIdWithPathForTrigger(self, collection_link, trigger): # pylint: disable=no-self-use - CosmosClientConnection.__ValidateResource(trigger) - trigger = trigger.copy() - if trigger.get("serverScript"): - trigger["body"] = str(trigger.pop("serverScript", "")) - elif trigger.get("body"): - trigger["body"] = str(trigger["body"]) - - path = base.GetPathFromLink(collection_link, "triggers") - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return collection_id, path, trigger - - def ReadTrigger(self, trigger_link, options=None, **kwargs): - """Reads a trigger. - - :param str trigger_link: - The link to the trigger. - :param dict options: - The request options for the request. - - :return: - The read Trigger. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(trigger_link) - trigger_id = base.GetResourceIdOrFullNameFromLink(trigger_link) - return self.Read(path, "triggers", trigger_id, None, options, **kwargs) - - def ReadUserDefinedFunctions(self, collection_link, options=None, **kwargs): - """Reads all user-defined functions in a collection. - - :param str collection_link: - The link to the document collection. - :param dict options: - The request options for the request. - - :return: - Query Iterable of UDFs. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - return self.QueryUserDefinedFunctions(collection_link, None, options, **kwargs) - - def QueryUserDefinedFunctions(self, collection_link, query, options=None, **kwargs): - """Queries user-defined functions in a collection. - - :param str collection_link: - The link to the collection. - :param (str or dict) query: - :param dict options: - The request options for the request. - - :return: - Query Iterable of UDFs. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(collection_link, "udfs") - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - - def fetch_fn(options): - return ( - self.__QueryFeed( - path, "udfs", collection_id, lambda r: r["UserDefinedFunctions"], - lambda _, b: b, query, options, **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable - ) - - def CreateUserDefinedFunction(self, collection_link, udf, options=None, **kwargs): - """Creates a user-defined function in a collection. - - :param str collection_link: - The link to the collection. - :param str udf: - :param dict options: - The request options for the request. - - :return: - The created UDF. - :rtype: - dict - - """ - if options is None: - options = {} - - collection_id, path, udf = self._GetContainerIdWithPathForUDF(collection_link, udf) - return self.Create(udf, path, "udfs", collection_id, None, options, **kwargs) - - def UpsertUserDefinedFunction(self, collection_link, udf, options=None, **kwargs): - """Upserts a user-defined function in a collection. - - :param str collection_link: - The link to the collection. - :param str udf: - :param dict options: - The request options for the request. - - :return: - The upserted UDF. - :rtype: - dict - - """ - if options is None: - options = {} - - collection_id, path, udf = self._GetContainerIdWithPathForUDF(collection_link, udf) - return self.Upsert(udf, path, "udfs", collection_id, None, options, **kwargs) - - def _GetContainerIdWithPathForUDF(self, collection_link, udf): # pylint: disable=no-self-use - CosmosClientConnection.__ValidateResource(udf) - udf = udf.copy() - if udf.get("serverScript"): - udf["body"] = str(udf.pop("serverScript", "")) - elif udf.get("body"): - udf["body"] = str(udf["body"]) - - path = base.GetPathFromLink(collection_link, "udfs") - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return collection_id, path, udf - - def ReadUserDefinedFunction(self, udf_link, options=None, **kwargs): - """Reads a user-defined function. - - :param str udf_link: - The link to the user-defined function. - :param dict options: - The request options for the request. - - :return: - The read UDF. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(udf_link) - udf_id = base.GetResourceIdOrFullNameFromLink(udf_link) - return self.Read(path, "udfs", udf_id, None, options, **kwargs) - - def ReadStoredProcedures(self, collection_link, options=None, **kwargs): - """Reads all store procedures in a collection. - - :param str collection_link: - The link to the document collection. - :param dict options: - The request options for the request. - - :return: - Query Iterable of Stored Procedures. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - return self.QueryStoredProcedures(collection_link, None, options, **kwargs) - - def QueryStoredProcedures(self, collection_link, query, options=None, **kwargs): - """Queries stored procedures in a collection. - - :param str collection_link: - The link to the document collection. - :param (str or dict) query: - :param dict options: - The request options for the request. - - :return: - Query Iterable of Stored Procedures. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(collection_link, "sprocs") - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - - def fetch_fn(options): - return ( - self.__QueryFeed( - path, "sprocs", collection_id, lambda r: r["StoredProcedures"], - lambda _, b: b, query, options, **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable - ) - - def CreateStoredProcedure(self, collection_link, sproc, options=None, **kwargs): - """Creates a stored procedure in a collection. - - :param str collection_link: - The link to the document collection. - :param str sproc: - :param dict options: - The request options for the request. - - :return: - The created Stored Procedure. - :rtype: - dict - - """ - if options is None: - options = {} - - collection_id, path, sproc = self._GetContainerIdWithPathForSproc(collection_link, sproc) - return self.Create(sproc, path, "sprocs", collection_id, None, options, **kwargs) - - def UpsertStoredProcedure(self, collection_link, sproc, options=None, **kwargs): - """Upserts a stored procedure in a collection. - - :param str collection_link: - The link to the document collection. - :param str sproc: - :param dict options: - The request options for the request. - - :return: - The upserted Stored Procedure. - :rtype: - dict - - """ - if options is None: - options = {} - - collection_id, path, sproc = self._GetContainerIdWithPathForSproc(collection_link, sproc) - return self.Upsert(sproc, path, "sprocs", collection_id, None, options, **kwargs) - - def _GetContainerIdWithPathForSproc(self, collection_link, sproc): # pylint: disable=no-self-use - CosmosClientConnection.__ValidateResource(sproc) - sproc = sproc.copy() - if sproc.get("serverScript"): - sproc["body"] = str(sproc.pop("serverScript", "")) - elif sproc.get("body"): - sproc["body"] = str(sproc["body"]) - path = base.GetPathFromLink(collection_link, "sprocs") - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return collection_id, path, sproc - - def ReadStoredProcedure(self, sproc_link, options=None, **kwargs): - """Reads a stored procedure. - - :param str sproc_link: - The link to the stored procedure. - :param dict options: - The request options for the request. - - :return: - The read Stored Procedure. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(sproc_link) - sproc_id = base.GetResourceIdOrFullNameFromLink(sproc_link) - return self.Read(path, "sprocs", sproc_id, None, options, **kwargs) - - def ReadConflicts(self, collection_link, feed_options=None, **kwargs): - """Reads conflicts. - - :param str collection_link: - The link to the document collection. - :param dict feed_options: - - :return: - Query Iterable of Conflicts. - :rtype: - query_iterable.QueryIterable - - """ - if feed_options is None: - feed_options = {} - - return self.QueryConflicts(collection_link, None, feed_options, **kwargs) - - def QueryConflicts(self, collection_link, query, options=None, **kwargs): - """Queries conflicts in a collection. - - :param str collection_link: - The link to the document collection. - :param (str or dict) query: - :param dict options: - The request options for the request. - - :return: - Query Iterable of Conflicts. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(collection_link, "conflicts") - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - - def fetch_fn(options): - return ( - self.__QueryFeed( - path, "conflicts", collection_id, lambda r: r["Conflicts"], - lambda _, b: b, query, options, **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable - ) - - def ReadConflict(self, conflict_link, options=None, **kwargs): - """Reads a conflict. - - :param str conflict_link: - The link to the conflict. - :param dict options: - - :return: - The read Conflict. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(conflict_link) - conflict_id = base.GetResourceIdOrFullNameFromLink(conflict_link) - return self.Read(path, "conflicts", conflict_id, None, options, **kwargs) - - def DeleteContainer(self, collection_link, options=None, **kwargs): - """Deletes a collection. - - :param str collection_link: - The link to the document collection. - :param dict options: - The request options for the request. - - :return: - The deleted Collection. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(collection_link) - collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return self.DeleteResource(path, "colls", collection_id, None, options, **kwargs) - - def ReplaceItem(self, document_link, new_document, options=None, **kwargs): - """Replaces a document and returns it. - - :param str document_link: - The link to the document. - :param dict new_document: - :param dict options: - The request options for the request. - - :return: - The new Document. - :rtype: - dict - - """ - CosmosClientConnection.__ValidateResource(new_document) - path = base.GetPathFromLink(document_link) - document_id = base.GetResourceIdOrFullNameFromLink(document_link) - - # Python's default arguments are evaluated once when the function is defined, - # not each time the function is called (like it is in say, Ruby). This means - # that if you use a mutable default argument and mutate it, you will and have - # mutated that object for all future calls to the function as well. So, using - # a non-mutable deafult in this case(None) and assigning an empty dict(mutable) - # inside the function so that it remains local For more details on this gotcha, - # please refer http://docs.python-guide.org/en/latest/writing/gotchas/ - if options is None: - options = {} - - # Extract the document collection link and add the partition key to options - collection_link = base.GetItemContainerLink(document_link) - options = self._AddPartitionKey(collection_link, new_document, options) - - return self.Replace(new_document, path, "docs", document_id, None, options, **kwargs) - - def DeleteItem(self, document_link, options=None, **kwargs): - """Deletes a document. - - :param str document_link: - The link to the document. - :param dict options: - The request options for the request. - - :return: - The deleted Document. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(document_link) - document_id = base.GetResourceIdOrFullNameFromLink(document_link) - return self.DeleteResource(path, "docs", document_id, None, options, **kwargs) - - def ReplaceTrigger(self, trigger_link, trigger, options=None, **kwargs): - """Replaces a trigger and returns it. - - :param str trigger_link: - The link to the trigger. - :param dict trigger: - :param dict options: - The request options for the request. - - :return: - The replaced Trigger. - :rtype: - dict - - """ - if options is None: - options = {} - - CosmosClientConnection.__ValidateResource(trigger) - trigger = trigger.copy() - if trigger.get("serverScript"): - trigger["body"] = str(trigger["serverScript"]) - elif trigger.get("body"): - trigger["body"] = str(trigger["body"]) - - path = base.GetPathFromLink(trigger_link) - trigger_id = base.GetResourceIdOrFullNameFromLink(trigger_link) - return self.Replace(trigger, path, "triggers", trigger_id, None, options, **kwargs) - - def DeleteTrigger(self, trigger_link, options=None, **kwargs): - """Deletes a trigger. - - :param str trigger_link: - The link to the trigger. - :param dict options: - The request options for the request. - - :return: - The deleted Trigger. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(trigger_link) - trigger_id = base.GetResourceIdOrFullNameFromLink(trigger_link) - return self.DeleteResource(path, "triggers", trigger_id, None, options, **kwargs) - - def ReplaceUserDefinedFunction(self, udf_link, udf, options=None, **kwargs): - """Replaces a user-defined function and returns it. - - :param str udf_link: - The link to the user-defined function. - :param dict udf: - :param dict options: - The request options for the request. - - :return: - The new UDF. - :rtype: - dict - - """ - if options is None: - options = {} - - CosmosClientConnection.__ValidateResource(udf) - udf = udf.copy() - if udf.get("serverScript"): - udf["body"] = str(udf["serverScript"]) - elif udf.get("body"): - udf["body"] = str(udf["body"]) - - path = base.GetPathFromLink(udf_link) - udf_id = base.GetResourceIdOrFullNameFromLink(udf_link) - return self.Replace(udf, path, "udfs", udf_id, None, options, **kwargs) - - def DeleteUserDefinedFunction(self, udf_link, options=None, **kwargs): - """Deletes a user-defined function. - - :param str udf_link: - The link to the user-defined function. - :param dict options: - The request options for the request. - - :return: - The deleted UDF. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(udf_link) - udf_id = base.GetResourceIdOrFullNameFromLink(udf_link) - return self.DeleteResource(path, "udfs", udf_id, None, options, **kwargs) - - def ExecuteStoredProcedure(self, sproc_link, params, options=None, **kwargs): - """Executes a store procedure. - - :param str sproc_link: - The link to the stored procedure. - :param dict params: - List or None - :param dict options: - The request options for the request. - - :return: - The Stored Procedure response. - :rtype: - dict - - """ - if options is None: - options = {} - - initial_headers = dict(self.default_headers) - initial_headers.update({http_constants.HttpHeaders.Accept: (runtime_constants.MediaTypes.Json)}) - - if params and not isinstance(params, list): - params = [params] - - path = base.GetPathFromLink(sproc_link) - sproc_id = base.GetResourceIdOrFullNameFromLink(sproc_link) - headers = base.GetHeaders(self, initial_headers, "post", path, sproc_id, "sprocs", options) - - # ExecuteStoredProcedure will use WriteEndpoint since it uses POST operation - request_params = _request_object.RequestObject("sprocs", documents._OperationType.ExecuteJavaScript) - result, self.last_response_headers = self.__Post(path, request_params, params, headers, **kwargs) - return result - - def ReplaceStoredProcedure(self, sproc_link, sproc, options=None, **kwargs): - """Replaces a stored procedure and returns it. - - :param str sproc_link: - The link to the stored procedure. - :param dict sproc: - :param dict options: - The request options for the request. - - :return: - The replaced Stored Procedure. - :rtype: - dict - - """ - if options is None: - options = {} - - CosmosClientConnection.__ValidateResource(sproc) - sproc = sproc.copy() - if sproc.get("serverScript"): - sproc["body"] = str(sproc["serverScript"]) - elif sproc.get("body"): - sproc["body"] = str(sproc["body"]) - - path = base.GetPathFromLink(sproc_link) - sproc_id = base.GetResourceIdOrFullNameFromLink(sproc_link) - return self.Replace(sproc, path, "sprocs", sproc_id, None, options, **kwargs) - - def DeleteStoredProcedure(self, sproc_link, options=None, **kwargs): - """Deletes a stored procedure. - - :param str sproc_link: - The link to the stored procedure. - :param dict options: - The request options for the request. - - :return: - The deleted Stored Procedure. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(sproc_link) - sproc_id = base.GetResourceIdOrFullNameFromLink(sproc_link) - return self.DeleteResource(path, "sprocs", sproc_id, None, options, **kwargs) - - def DeleteConflict(self, conflict_link, options=None, **kwargs): - """Deletes a conflict. - - :param str conflict_link: - The link to the conflict. - :param dict options: - The request options for the request. - - :return: - The deleted Conflict. - :rtype: - dict - - """ - if options is None: - options = {} - - path = base.GetPathFromLink(conflict_link) - conflict_id = base.GetResourceIdOrFullNameFromLink(conflict_link) - return self.DeleteResource(path, "conflicts", conflict_id, None, options, **kwargs) - - def ReplaceOffer(self, offer_link, offer, **kwargs): - """Replaces an offer and returns it. - - :param str offer_link: - The link to the offer. - :param dict offer: - - :return: - The replaced Offer. - :rtype: - dict - - """ - CosmosClientConnection.__ValidateResource(offer) - path = base.GetPathFromLink(offer_link) - offer_id = base.GetResourceIdOrFullNameFromLink(offer_link) - return self.Replace(offer, path, "offers", offer_id, None, None, **kwargs) - - def ReadOffer(self, offer_link, **kwargs): - """Reads an offer. - - :param str offer_link: - The link to the offer. - - :return: - The read Offer. - :rtype: - dict - - """ - path = base.GetPathFromLink(offer_link) - offer_id = base.GetResourceIdOrFullNameFromLink(offer_link) - return self.Read(path, "offers", offer_id, None, {}, **kwargs) - - def ReadOffers(self, options=None, **kwargs): - """Reads all offers. - - :param dict options: - The request options for the request - - :return: - Query Iterable of Offers. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - return self.QueryOffers(None, options, **kwargs) - - def QueryOffers(self, query, options=None, **kwargs): - """Query for all offers. - - :param (str or dict) query: - :param dict options: - The request options for the request - - :return: - Query Iterable of Offers. - :rtype: - query_iterable.QueryIterable - - """ - if options is None: - options = {} - - def fetch_fn(options): - return ( - self.__QueryFeed( - "/offers", "offers", "", lambda r: r["Offers"], lambda _, b: b, query, options, **kwargs - ), - self.last_response_headers, - ) - - return AsyncItemPaged( - self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable - ) - - async def GetDatabaseAccount(self, url_connection=None, **kwargs): - """Gets database account info. - - :return: - The Database Account. - :rtype: - documents.DatabaseAccount - - """ - if url_connection is None: - url_connection = self.url_connection - - initial_headers = dict(self.default_headers) - headers = base.GetHeaders(self, initial_headers, "get", "", "", "", {}) # path # id # type - - request_params = _request_object.RequestObject("databaseaccount", documents._OperationType.Read, url_connection) - result, self.last_response_headers = await self.__Get("", request_params, headers, **kwargs) - database_account = documents.DatabaseAccount() - database_account.DatabasesLink = "/dbs/" - database_account.MediaLink = "/media/" - if http_constants.HttpHeaders.MaxMediaStorageUsageInMB in self.last_response_headers: - database_account.MaxMediaStorageUsageInMB = self.last_response_headers[ - http_constants.HttpHeaders.MaxMediaStorageUsageInMB - ] - if http_constants.HttpHeaders.CurrentMediaStorageUsageInMB in self.last_response_headers: - database_account.CurrentMediaStorageUsageInMB = self.last_response_headers[ - http_constants.HttpHeaders.CurrentMediaStorageUsageInMB - ] - database_account.ConsistencyPolicy = result.get(constants._Constants.UserConsistencyPolicy) - - # WritableLocations and ReadableLocations fields will be available only for geo-replicated database accounts - if constants._Constants.WritableLocations in result: - database_account._WritableLocations = result[constants._Constants.WritableLocations] - if constants._Constants.ReadableLocations in result: - database_account._ReadableLocations = result[constants._Constants.ReadableLocations] - if constants._Constants.EnableMultipleWritableLocations in result: - database_account._EnableMultipleWritableLocations = result[ - constants._Constants.EnableMultipleWritableLocations - ] - - self._useMultipleWriteLocations = ( - self.connection_policy.UseMultipleWriteLocations and database_account._EnableMultipleWritableLocations - ) - return database_account - - async def Create(self, body, path, typ, id, initial_headers, options=None, **kwargs): # pylint: disable=redefined-builtin - """Creates a Azure Cosmos resource and returns it. - - :param dict body: - :param str path: - :param str typ: - :param str id: - :param dict initial_headers: - :param dict options: - The request options for the request. - - :return: - The created Azure Cosmos resource. - :rtype: - dict - - """ - if options is None: - options = {} - - initial_headers = initial_headers or self.default_headers - headers = base.GetHeaders(self, initial_headers, "post", path, id, typ, options) - # Create will use WriteEndpoint since it uses POST operation - - request_params = _request_object.RequestObject(typ, documents._OperationType.Create) - result, self.last_response_headers = await self.__Post(path, request_params, body, headers, **kwargs) - - # update session for write request - self._UpdateSessionIfRequired(headers, result, self.last_response_headers) - return result - - def Upsert(self, body, path, typ, id, initial_headers, options=None, **kwargs): # pylint: disable=redefined-builtin - """Upserts a Azure Cosmos resource and returns it. - - :param dict body: - :param str path: - :param str typ: - :param str id: - :param dict initial_headers: - :param dict options: - The request options for the request. - - :return: - The upserted Azure Cosmos resource. - :rtype: - dict - - """ - if options is None: - options = {} - - initial_headers = initial_headers or self.default_headers - headers = base.GetHeaders(self, initial_headers, "post", path, id, typ, options) - - headers[http_constants.HttpHeaders.IsUpsert] = True - - # Upsert will use WriteEndpoint since it uses POST operation - request_params = _request_object.RequestObject(typ, documents._OperationType.Upsert) - result, self.last_response_headers = self.__Post(path, request_params, body, headers, **kwargs) - # update session for write request - self._UpdateSessionIfRequired(headers, result, self.last_response_headers) - return result - - def Replace(self, resource, path, typ, id, initial_headers, options=None, **kwargs): # pylint: disable=redefined-builtin - """Replaces a Azure Cosmos resource and returns it. - - :param dict resource: - :param str path: - :param str typ: - :param str id: - :param dict initial_headers: - :param dict options: - The request options for the request. - - :return: - The new Azure Cosmos resource. - :rtype: - dict - - """ - if options is None: - options = {} - - initial_headers = initial_headers or self.default_headers - headers = base.GetHeaders(self, initial_headers, "put", path, id, typ, options) - # Replace will use WriteEndpoint since it uses PUT operation - request_params = _request_object.RequestObject(typ, documents._OperationType.Replace) - result, self.last_response_headers = self.__Put(path, request_params, resource, headers, **kwargs) - - # update session for request mutates data on server side - self._UpdateSessionIfRequired(headers, result, self.last_response_headers) - return result - - def Read(self, path, typ, id, initial_headers, options=None, **kwargs): # pylint: disable=redefined-builtin - """Reads a Azure Cosmos resource and returns it. - - :param str path: - :param str typ: - :param str id: - :param dict initial_headers: - :param dict options: - The request options for the request. - - :return: - The upserted Azure Cosmos resource. - :rtype: - dict - - """ - if options is None: - options = {} - - initial_headers = initial_headers or self.default_headers - headers = base.GetHeaders(self, initial_headers, "get", path, id, typ, options) - # Read will use ReadEndpoint since it uses GET operation - request_params = _request_object.RequestObject(typ, documents._OperationType.Read) - result, self.last_response_headers = self.__Get(path, request_params, headers, **kwargs) - return result - - async def ReadAsync(self, path, typ, id, initial_headers, options=None, **kwargs): # pylint: disable=redefined-builtin - """Reads a Azure Cosmos resource and returns it. - - :param str path: - :param str typ: - :param str id: - :param dict initial_headers: - :param dict options: - The request options for the request. - - :return: - The upserted Azure Cosmos resource. - :rtype: - dict - - """ - if options is None: - options = {} - - initial_headers = initial_headers or self.default_headers - headers = base.GetHeaders(self, initial_headers, "get", path, id, typ, options) - # Read will use ReadEndpoint since it uses GET operation - request_params = _request_object.RequestObject(typ, documents._OperationType.Read) - result, self.last_response_headers = await self.__Get(path, request_params, headers, **kwargs) - return result - - def DeleteResource(self, path, typ, id, initial_headers, options=None, **kwargs): # pylint: disable=redefined-builtin - """Deletes a Azure Cosmos resource and returns it. - - :param str path: - :param str typ: - :param str id: - :param dict initial_headers: - :param dict options: - The request options for the request. - - :return: - The deleted Azure Cosmos resource. - :rtype: - dict - - """ - if options is None: - options = {} - - initial_headers = initial_headers or self.default_headers - headers = base.GetHeaders(self, initial_headers, "delete", path, id, typ, options) - # Delete will use WriteEndpoint since it uses DELETE operation - request_params = _request_object.RequestObject(typ, documents._OperationType.Delete) - result, self.last_response_headers = self.__Delete(path, request_params, headers, **kwargs) - - # update session for request mutates data on server side - self._UpdateSessionIfRequired(headers, result, self.last_response_headers) - - return result - - async def __Get(self, path, request_params, req_headers, **kwargs): - """Azure Cosmos 'GET' http request. - - :params str url: - :params str path: - :params dict req_headers: - - :return: - Tuple of (result, headers). - :rtype: - tuple of (dict, dict) - - """ - request = self.pipeline_client.get(url=path, headers=req_headers) - return await asynchronous_request.AsynchronousRequest( - client=self, - request_params=request_params, - global_endpoint_manager=self._global_endpoint_manager, - connection_policy=self.connection_policy, - pipeline_client=self.pipeline_client, - request=request, - request_data=None, - **kwargs - ) - - def __GetSync(self, path, request_params, req_headers, **kwargs): - """Azure Cosmos 'GET' http request. - - :params str url: - :params str path: - :params dict req_headers: - - :return: - Tuple of (result, headers). - :rtype: - tuple of (dict, dict) - - """ - request = self.pipeline_client.get(url=path, headers=req_headers) - return synchronized_request.SynchronizedRequest( - client=self, - request_params=request_params, - global_endpoint_manager=self._global_endpoint_manager, - connection_policy=self.connection_policy, - pipeline_client=self.pipeline_client, - request=request, - request_data=None, - **kwargs - ) - - def __Post(self, path, request_params, body, req_headers, **kwargs): - """Azure Cosmos 'POST' http request. - - :params str url: - :params str path: - :params (str, unicode, dict) body: - :params dict req_headers: - - :return: - Tuple of (result, headers). - :rtype: - tuple of (dict, dict) - - """ - request = self.pipeline_client.post(url=path, headers=req_headers) - return synchronized_request.SynchronizedRequest( - client=self, - request_params=request_params, - global_endpoint_manager=self._global_endpoint_manager, - connection_policy=self.connection_policy, - pipeline_client=self.pipeline_client, - request=request, - request_data=body, - **kwargs - ) - - def __Put(self, path, request_params, body, req_headers, **kwargs): - """Azure Cosmos 'PUT' http request. - - :params str url: - :params str path: - :params (str, unicode, dict) body: - :params dict req_headers: - - :return: - Tuple of (result, headers). - :rtype: - tuple of (dict, dict) - - """ - request = self.pipeline_client.put(url=path, headers=req_headers) - return synchronized_request.SynchronizedRequest( - client=self, - request_params=request_params, - global_endpoint_manager=self._global_endpoint_manager, - connection_policy=self.connection_policy, - pipeline_client=self.pipeline_client, - request=request, - request_data=body, - **kwargs - ) - - def __Delete(self, path, request_params, req_headers, **kwargs): - """Azure Cosmos 'DELETE' http request. - - :params str url: - :params str path: - :params dict req_headers: - - :return: - Tuple of (result, headers). - :rtype: - tuple of (dict, dict) - - """ - request = self.pipeline_client.delete(url=path, headers=req_headers) - return synchronized_request.SynchronizedRequest( - client=self, - request_params=request_params, - global_endpoint_manager=self._global_endpoint_manager, - connection_policy=self.connection_policy, - pipeline_client=self.pipeline_client, - request=request, - request_data=None, - **kwargs - ) - - def QueryFeed(self, path, collection_id, query, options, partition_key_range_id=None, **kwargs): - """Query Feed for Document Collection resource. - - :param str path: - Path to the document collection. - :param str collection_id: - Id of the document collection. - :param (str or dict) query: - :param dict options: - The request options for the request. - :param str partition_key_range_id: - Partition key range id. - :rtype: - tuple - - """ - return ( - self.__QueryFeed( - path, - "docs", - collection_id, - lambda r: r["Documents"], - lambda _, b: b, - query, - options, - partition_key_range_id, - **kwargs - ), - self.last_response_headers, - ) - - def __QueryFeed( - self, - path, - typ, - id_, - result_fn, - create_fn, - query, - options=None, - partition_key_range_id=None, - response_hook=None, - is_query_plan=False, - **kwargs - ): - """Query for more than one Azure Cosmos resources. - - :param str path: - :param str typ: - :param str id_: - :param function result_fn: - :param function create_fn: - :param (str or dict) query: - :param dict options: - The request options for the request. - :param str partition_key_range_id: - Specifies partition key range id. - :param function response_hook: - :param bool is_query_plan: - Specififes if the call is to fetch query plan - - :rtype: - list - - :raises SystemError: If the query compatibility mode is undefined. - - """ - if options is None: - options = {} - - if query: - __GetBodiesFromQueryResult = result_fn - else: - - def __GetBodiesFromQueryResult(result): - if result is not None: - return [create_fn(self, body) for body in result_fn(result)] - # If there is no change feed, the result data is empty and result is None. - # This case should be interpreted as an empty array. - return [] - - initial_headers = self.default_headers.copy() - # Copy to make sure that default_headers won't be changed. - if query is None: - # Query operations will use ReadEndpoint even though it uses GET(for feed requests) - request_params = _request_object.RequestObject(typ, - documents._OperationType.QueryPlan if is_query_plan else documents._OperationType.ReadFeed) - headers = base.GetHeaders(self, initial_headers, "get", path, id_, typ, options, partition_key_range_id) - result, self.last_response_headers = self.__Get(path, request_params, headers, **kwargs) - if response_hook: - response_hook(self.last_response_headers, result) - return __GetBodiesFromQueryResult(result) - - query = self.__CheckAndUnifyQueryFormat(query) - - initial_headers[http_constants.HttpHeaders.IsQuery] = "true" - if not is_query_plan: - initial_headers[http_constants.HttpHeaders.IsQuery] = "true" - - if ( - self._query_compatibility_mode == CosmosClientConnection._QueryCompatibilityMode.Default - or self._query_compatibility_mode == CosmosClientConnection._QueryCompatibilityMode.Query - ): - initial_headers[http_constants.HttpHeaders.ContentType] = runtime_constants.MediaTypes.QueryJson - elif self._query_compatibility_mode == CosmosClientConnection._QueryCompatibilityMode.SqlQuery: - initial_headers[http_constants.HttpHeaders.ContentType] = runtime_constants.MediaTypes.SQL - else: - raise SystemError("Unexpected query compatibility mode.") - - # Query operations will use ReadEndpoint even though it uses POST(for regular query operations) - request_params = _request_object.RequestObject(typ, documents._OperationType.SqlQuery) - req_headers = base.GetHeaders(self, initial_headers, "post", path, id_, typ, options, partition_key_range_id) - result, self.last_response_headers = self.__Post(path, request_params, query, req_headers, **kwargs) - - if response_hook: - response_hook(self.last_response_headers, result) - - return __GetBodiesFromQueryResult(result) - - def _GetQueryPlanThroughGateway(self, query, resource_link, **kwargs): - supported_query_features = (documents._QueryFeature.Aggregate + "," + - documents._QueryFeature.CompositeAggregate + "," + - documents._QueryFeature.Distinct + "," + - documents._QueryFeature.MultipleOrderBy + "," + - documents._QueryFeature.OffsetAndLimit + "," + - documents._QueryFeature.OrderBy + "," + - documents._QueryFeature.Top) - - options = { - "contentType": runtime_constants.MediaTypes.Json, - "isQueryPlanRequest": True, - "supportedQueryFeatures": supported_query_features, - "queryVersion": http_constants.Versions.QueryVersion - } - - resource_link = base.TrimBeginningAndEndingSlashes(resource_link) - path = base.GetPathFromLink(resource_link, "docs") - resource_id = base.GetResourceIdOrFullNameFromLink(resource_link) - - return self.__QueryFeed(path, - "docs", - resource_id, - lambda r: r, - None, - query, - options, - is_query_plan=True, - **kwargs) - - def __CheckAndUnifyQueryFormat(self, query_body): - """Checks and unifies the format of the query body. - - :raises TypeError: If query_body is not of expected type (depending on the query compatibility mode). - :raises ValueError: If query_body is a dict but doesn\'t have valid query text. - :raises SystemError: If the query compatibility mode is undefined. - - :param (str or dict) query_body: - - :return: - The formatted query body. - :rtype: - dict or string - """ - if ( - self._query_compatibility_mode == CosmosClientConnection._QueryCompatibilityMode.Default - or self._query_compatibility_mode == CosmosClientConnection._QueryCompatibilityMode.Query - ): - if not isinstance(query_body, dict) and not isinstance(query_body, six.string_types): - raise TypeError("query body must be a dict or string.") - if isinstance(query_body, dict) and not query_body.get("query"): - raise ValueError('query body must have valid query text with key "query".') - if isinstance(query_body, six.string_types): - return {"query": query_body} - elif ( - self._query_compatibility_mode == CosmosClientConnection._QueryCompatibilityMode.SqlQuery - and not isinstance(query_body, six.string_types) - ): - raise TypeError("query body must be a string.") - else: - raise SystemError("Unexpected query compatibility mode.") - - return query_body - - @staticmethod - def __ValidateResource(resource): - id_ = resource.get("id") - if id_: - try: - if id_.find("/") != -1 or id_.find("\\") != -1 or id_.find("?") != -1 or id_.find("#") != -1: - raise ValueError("Id contains illegal chars.") - - if id_[-1] == " ": - raise ValueError("Id ends with a space.") - except AttributeError: - raise_with_traceback(TypeError, message="Id type must be a string.") - - # Adds the partition key to options - def _AddPartitionKey(self, collection_link, document, options): - collection_link = base.TrimBeginningAndEndingSlashes(collection_link) - - # TODO: Refresh the cache if partition is extracted automatically and we get a 400.1001 - - # If the document collection link is present in the cache, then use the cached partitionkey definition - if collection_link in self.partition_key_definition_cache: - partitionKeyDefinition = self.partition_key_definition_cache.get(collection_link) - # Else read the collection from backend and add it to the cache - else: - collection = self.ReadContainer(collection_link) - partitionKeyDefinition = collection.get("partitionKey") - self.partition_key_definition_cache[collection_link] = partitionKeyDefinition - - # If the collection doesn't have a partition key definition, skip it as it's a legacy collection - if partitionKeyDefinition: - # If the user has passed in the partitionKey in options use that elase extract it from the document - if "partitionKey" not in options: - partitionKeyValue = self._ExtractPartitionKey(partitionKeyDefinition, document) - options["partitionKey"] = partitionKeyValue - - return options - - # Extracts the partition key from the document using the partitionKey definition - def _ExtractPartitionKey(self, partitionKeyDefinition, document): - - # Parses the paths into a list of token each representing a property - partition_key_parts = base.ParsePaths(partitionKeyDefinition.get("paths")) - # Check if the partitionKey is system generated or not - is_system_key = partitionKeyDefinition["systemKey"] if "systemKey" in partitionKeyDefinition else False - - # Navigates the document to retrieve the partitionKey specified in the paths - return self._retrieve_partition_key(partition_key_parts, document, is_system_key) - - # Navigates the document to retrieve the partitionKey specified in the partition key parts - def _retrieve_partition_key(self, partition_key_parts, document, is_system_key): - expected_matchCount = len(partition_key_parts) - matchCount = 0 - partitionKey = document - - for part in partition_key_parts: - # At any point if we don't find the value of a sub-property in the document, we return as Undefined - if part not in partitionKey: - return self._return_undefined_or_empty_partition_key(is_system_key) - - partitionKey = partitionKey.get(part) - matchCount += 1 - # Once we reach the "leaf" value(not a dict), we break from loop - if not isinstance(partitionKey, dict): - break - - # Match the count of hops we did to get the partitionKey with the length of - # partition key parts and validate that it's not a dict at that level - if (matchCount != expected_matchCount) or isinstance(partitionKey, dict): - return self._return_undefined_or_empty_partition_key(is_system_key) - - return partitionKey - - def _UpdateSessionIfRequired(self, request_headers, response_result, response_headers): - """ - Updates session if necessary. - - :param dict response_result: - :param dict response_headers: - :param dict response_headers - - :return: - None, but updates the client session if necessary. - - """ - - # if this request was made with consistency level as session, then update the session - if response_result is None or response_headers is None: - return - - is_session_consistency = False - if http_constants.HttpHeaders.ConsistencyLevel in request_headers: - if documents.ConsistencyLevel.Session == request_headers[http_constants.HttpHeaders.ConsistencyLevel]: - is_session_consistency = True - - if is_session_consistency: - # update session - self.session.update_session(response_result, response_headers) - - @staticmethod - def _return_undefined_or_empty_partition_key(is_system_key): - if is_system_key: - return _Empty - return _Undefined + async def _setup(self): + if not 'database_account' in self._setup_kwargs: + self._setup_kwargs['database_account'] = await self._global_endpoint_manager._GetDatabaseAccount(**self._setup_kwargs) + await self._global_endpoint_manager.force_refresh(self._setup_kwargs['database_account']) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_endpoint_manager_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_endpoint_manager_async.py index 2248619fb335..089daf5a225d 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_endpoint_manager_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_endpoint_manager_async.py @@ -23,7 +23,7 @@ database service. """ -import threading +import asyncio from six.moves.urllib.parse import urlparse @@ -54,7 +54,7 @@ def __init__(self, client): self.refresh_time_interval_in_ms, ) self.refresh_needed = False - self.refresh_lock = threading.RLock() + self.refresh_lock = asyncio.RLock() self.last_refresh_time = 0 def get_refresh_time_interval_in_ms_stub(self): # pylint: disable=no-self-use @@ -84,21 +84,21 @@ def get_ordered_read_endpoints(self): def can_use_multiple_write_locations(self, request): return self.location_cache.can_use_multiple_write_locations_for_request(request) - def force_refresh(self, database_account): + async def force_refresh(self, database_account): self.refresh_needed = True - self.refresh_endpoint_list(database_account) + await self.refresh_endpoint_list(database_account) - def refresh_endpoint_list(self, database_account, **kwargs): - with self.refresh_lock: + async def refresh_endpoint_list(self, database_account, **kwargs): + async with self.refresh_lock: # if refresh is not needed or refresh is already taking place, return if not self.refresh_needed: return try: - self._refresh_endpoint_list_private(database_account, **kwargs) + await self._refresh_endpoint_list_private(database_account, **kwargs) except Exception as e: raise e - def _refresh_endpoint_list_private(self, database_account=None, **kwargs): + async def _refresh_endpoint_list_private(self, database_account=None, **kwargs): if database_account: self.location_cache.perform_on_database_account_read(database_account) self.refresh_needed = False @@ -108,12 +108,12 @@ def _refresh_endpoint_list_private(self, database_account=None, **kwargs): and self.location_cache.current_time_millis() - self.last_refresh_time > self.refresh_time_interval_in_ms ): if not database_account: - database_account = self._GetDatabaseAccount(**kwargs) + database_account = await self._GetDatabaseAccount(**kwargs) self.location_cache.perform_on_database_account_read(database_account) self.last_refresh_time = self.location_cache.current_time_millis() self.refresh_needed = False - def _GetDatabaseAccount(self, **kwargs): + async def _GetDatabaseAccount(self, **kwargs): """Gets the database account. First tries by using the default endpoint, and if that doesn't work, @@ -121,7 +121,7 @@ def _GetDatabaseAccount(self, **kwargs): specified, to get the database account. """ try: - database_account = self._GetDatabaseAccountStub(self.DefaultEndpoint, **kwargs) + database_account = await self._GetDatabaseAccountStub(self.DefaultEndpoint, **kwargs) return database_account # If for any reason(non-globaldb related), we are not able to get the database # account from the above call to GetDatabaseAccount, we would try to get this @@ -133,7 +133,7 @@ def _GetDatabaseAccount(self, **kwargs): for location_name in self.PreferredLocations: locational_endpoint = _GlobalEndpointManager.GetLocationalEndpoint(self.DefaultEndpoint, location_name) try: - database_account = self._GetDatabaseAccountStub(locational_endpoint, **kwargs) + database_account = await self._GetDatabaseAccountStub(locational_endpoint, **kwargs) return database_account except exceptions.CosmosHttpResponseError: pass diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_retry_utility.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_retry_utility.py new file mode 100644 index 000000000000..39e2bd1264e3 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_retry_utility.py @@ -0,0 +1,196 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal methods for executing functions in the Azure Cosmos database service. +""" + +import time +import asyncio + +from azure.core.exceptions import AzureError, ClientAuthenticationError +from azure.core.pipeline.policies import AsyncRetryPolicy + +from .. import exceptions +from ..http_constants import HttpHeaders, StatusCodes, SubStatusCodes +from .._retry_utility import _configure_timeout +from .. import _endpoint_discovery_retry_policy +from .. import _resource_throttle_retry_policy +from .. import _default_retry_policy +from .. import _session_retry_policy + + +# pylint: disable=protected-access + + +async def ExecuteAsync(client, global_endpoint_manager, function, *args, **kwargs): + """Executes the function with passed parameters applying all retry policies + + :param object client: + Document client instance + :param object global_endpoint_manager: + Instance of _GlobalEndpointManager class + :param function function: + Function to be called wrapped with retries + :param (non-keyworded, variable number of arguments list) *args: + :param (keyworded, variable number of arguments list) **kwargs: + + """ + # instantiate all retry policies here to be applied for each request execution + endpointDiscovery_retry_policy = _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy( + client.connection_policy, global_endpoint_manager, *args + ) + + resourceThrottle_retry_policy = _resource_throttle_retry_policy.ResourceThrottleRetryPolicy( + client.connection_policy.RetryOptions.MaxRetryAttemptCount, + client.connection_policy.RetryOptions.FixedRetryIntervalInMilliseconds, + client.connection_policy.RetryOptions.MaxWaitTimeInSeconds, + ) + defaultRetry_policy = _default_retry_policy.DefaultRetryPolicy(*args) + + sessionRetry_policy = _session_retry_policy._SessionRetryPolicy( + client.connection_policy.EnableEndpointDiscovery, global_endpoint_manager, *args + ) + while True: + try: + client_timeout = kwargs.get('timeout') + start_time = time.time() + if args: + result = await ExecuteFunctionAsync(function, global_endpoint_manager, *args, **kwargs) + else: + result = await ExecuteFunctionAsync(function, *args, **kwargs) + if not client.last_response_headers: + client.last_response_headers = {} + + # setting the throttle related response headers before returning the result + client.last_response_headers[ + HttpHeaders.ThrottleRetryCount + ] = resourceThrottle_retry_policy.current_retry_attempt_count + client.last_response_headers[ + HttpHeaders.ThrottleRetryWaitTimeInMs + ] = resourceThrottle_retry_policy.cummulative_wait_time_in_milliseconds + + return result + except exceptions.CosmosHttpResponseError as e: + retry_policy = None + if e.status_code == StatusCodes.FORBIDDEN and e.sub_status == SubStatusCodes.WRITE_FORBIDDEN: + retry_policy = endpointDiscovery_retry_policy + elif e.status_code == StatusCodes.TOO_MANY_REQUESTS: + retry_policy = resourceThrottle_retry_policy + elif ( + e.status_code == StatusCodes.NOT_FOUND + and e.sub_status + and e.sub_status == SubStatusCodes.READ_SESSION_NOTAVAILABLE + ): + retry_policy = sessionRetry_policy + else: + retry_policy = defaultRetry_policy + + # If none of the retry policies applies or there is no retry needed, set the + # throttle related response hedaers and re-throw the exception back arg[0] + # is the request. It needs to be modified for write forbidden exception + if not retry_policy.ShouldRetry(e): + if not client.last_response_headers: + client.last_response_headers = {} + client.last_response_headers[ + HttpHeaders.ThrottleRetryCount + ] = resourceThrottle_retry_policy.current_retry_attempt_count + client.last_response_headers[ + HttpHeaders.ThrottleRetryWaitTimeInMs + ] = resourceThrottle_retry_policy.cummulative_wait_time_in_milliseconds + if args and args[0].should_clear_session_token_on_session_read_failure: + client.session.clear_session_token(client.last_response_headers) + raise + + # Wait for retry_after_in_milliseconds time before the next retry + await asyncio.sleep(retry_policy.retry_after_in_milliseconds / 1000.0) + if client_timeout: + kwargs['timeout'] = client_timeout - (time.time() - start_time) + if kwargs['timeout'] <= 0: + raise exceptions.CosmosClientTimeoutError() + + +async def ExecuteFunctionAsync(function, *args, **kwargs): + """Stub method so that it can be used for mocking purposes as well. + """ + return await function(*args, **kwargs) + + +class ConnectionRetryPolicy(AsyncRetryPolicy): + + def __init__(self, **kwargs): + clean_kwargs = {k: v for k, v in kwargs.items() if v is not None} + super(ConnectionRetryPolicy, self).__init__(**clean_kwargs) + + async def send(self, request): + """Sends the PipelineRequest object to the next policy. Uses retry settings if necessary. + Also enforces an absolute client-side timeout that spans multiple retry attempts. + + :param request: The PipelineRequest object + :type request: ~azure.core.pipeline.PipelineRequest + :return: Returns the PipelineResponse or raises error if maximum retries exceeded. + :rtype: ~azure.core.pipeline.PipelineResponse + :raises ~azure.core.exceptions.AzureError: Maximum retries exceeded. + :raises ~azure.cosmos.exceptions.CosmosClientTimeoutError: Specified timeout exceeded. + :raises ~azure.core.exceptions.ClientAuthenticationError: Authentication failed. + """ + absolute_timeout = request.context.options.pop('timeout', None) + per_request_timeout = request.context.options.pop('connection_timeout', 0) + + retry_error = None + retry_active = True + response = None + retry_settings = self.configure_retries(request.context.options) + while retry_active: + try: + start_time = time.time() + _configure_timeout(request, absolute_timeout, per_request_timeout) + + response = await self.next.send(request) + if self.is_retry(retry_settings, response): + retry_active = self.increment(retry_settings, response=response) + if retry_active: + await self.sleep(retry_settings, request.context.transport, response=response) + continue + break + except ClientAuthenticationError: # pylint:disable=try-except-raise + # the authentication policy failed such that the client's request can't + # succeed--we'll never have a response to it, so propagate the exception + raise + except exceptions.CosmosClientTimeoutError as timeout_error: + timeout_error.inner_exception = retry_error + timeout_error.response = response + timeout_error.history = retry_settings['history'] + raise + except AzureError as err: + retry_error = err + if self._is_method_retryable(retry_settings, request.http_request): + retry_active = self.increment(retry_settings, response=request, error=err) + if retry_active: + await self.sleep(retry_settings, request.context.transport) + continue + raise err + finally: + end_time = time.time() + if absolute_timeout: + absolute_timeout -= (end_time - start_time) + + self.update_context(response.context, retry_settings) + return response diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/container_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/container_async.py deleted file mode 100644 index 6120d533c918..000000000000 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/container_async.py +++ /dev/null @@ -1,802 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) 2014 Microsoft Corporation - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -"""Create, read, update and delete items in the Azure Cosmos DB SQL API service. -""" - -from typing import Any, Dict, List, Optional, Union, Iterable, cast # pylint: disable=unused-import - -import six -import asyncio -import time -from azure.core.tracing.decorator import distributed_trace # type: ignore - -from ._cosmos_client_connection_async import CosmosClientConnection -from .._base import build_options -from ..exceptions import CosmosResourceNotFoundError -from ..http_constants import StatusCodes -from ..offer import Offer -from ..scripts import ScriptsProxy -from ..partition_key import NonePartitionKeyValue - -__all__ = ("ContainerProxy",) - -# pylint: disable=protected-access -# pylint: disable=missing-client-constructor-parameter-credential,missing-client-constructor-parameter-kwargs - - -class ContainerProxy(object): - """An interface to interact with a specific DB Container. - - This class should not be instantiated directly. Instead, use the - :func:`DatabaseProxy.get_container_client` method to get an existing - container, or the :func:`Database.create_container` method to create a - new container. - - A container in an Azure Cosmos DB SQL API database is a collection of - documents, each of which is represented as an Item. - - :ivar str id: ID (name) of the container - :ivar str session_token: The session token for the container. - """ - - def __init__(self, client_connection, database_link, id, properties=None): # pylint: disable=redefined-builtin - # type: (CosmosClientConnection, str, str, Dict[str, Any]) -> None - self.client_connection = client_connection - self.id = id - self._properties = properties - self.container_link = u"{}/colls/{}".format(database_link, self.id) - self._is_system_key = None - self._scripts = None # type: Optional[ScriptsProxy] - - def __repr__(self): - # type () -> str - return "".format(self.container_link)[:1024] - - def _get_properties(self): - # type: () -> Dict[str, Any] - if self._properties is None: - self._properties = self.read() - return self._properties - - @property - def is_system_key(self): - # type: () -> bool - if self._is_system_key is None: - properties = self._get_properties() - self._is_system_key = ( - properties["partitionKey"]["systemKey"] if "systemKey" in properties["partitionKey"] else False - ) - return cast('bool', self._is_system_key) - - @property - def scripts(self): - # type: () -> ScriptsProxy - if self._scripts is None: - self._scripts = ScriptsProxy(self.client_connection, self.container_link, self.is_system_key) - return cast('ScriptsProxy', self._scripts) - - def _get_document_link(self, item_or_link): - # type: (Union[Dict[str, Any], str]) -> str - if isinstance(item_or_link, six.string_types): - return u"{}/docs/{}".format(self.container_link, item_or_link) - return item_or_link["_self"] - - def _get_conflict_link(self, conflict_or_link): - # type: (Union[Dict[str, Any], str]) -> str - if isinstance(conflict_or_link, six.string_types): - return u"{}/conflicts/{}".format(self.container_link, conflict_or_link) - return conflict_or_link["_self"] - - def _set_partition_key(self, partition_key): - if partition_key == NonePartitionKeyValue: - return CosmosClientConnection._return_undefined_or_empty_partition_key(self.is_system_key) - return partition_key - - @distributed_trace - def read( - self, - populate_query_metrics=None, # type: Optional[bool] - populate_partition_key_range_statistics=None, # type: Optional[bool] - populate_quota_info=None, # type: Optional[bool] - **kwargs # type: Any - ): - # type: (...) -> Dict[str, Any] - """Read the container properties. - - :param populate_query_metrics: Enable returning query metrics in response headers. - :param populate_partition_key_range_statistics: Enable returning partition key - range statistics in response headers. - :param populate_quota_info: Enable returning collection storage quota information in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword Callable response_hook: A callable invoked with the response metadata. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Raised if the container couldn't be retrieved. - This includes if the container does not exist. - :returns: Dict representing the retrieved container. - :rtype: dict[str, Any] - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - if populate_partition_key_range_statistics is not None: - request_options["populatePartitionKeyRangeStatistics"] = populate_partition_key_range_statistics - if populate_quota_info is not None: - request_options["populateQuotaInfo"] = populate_quota_info - - collection_link = self.container_link - self._properties = self.client_connection.ReadContainer( - collection_link, options=request_options, **kwargs - ) - - if response_hook: - response_hook(self.client_connection.last_response_headers, self._properties) - - return cast('Dict[str, Any]', self._properties) - - @distributed_trace - async def read_item( - self, - item, # type: Union[str, Dict[str, Any]] - partition_key, # type: Any - populate_query_metrics=None, # type: Optional[bool] - post_trigger_include=None, # type: Optional[str] - **kwargs # type: Any - ): - # type: (...) -> Dict[str, str] - """Get the item identified by `item`. - - :param item: The ID (name) or dict representing item to retrieve. - :param partition_key: Partition key for the item to retrieve. - :param populate_query_metrics: Enable returning query metrics in response headers. - :param post_trigger_include: trigger id to be used as post operation trigger. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: Dict representing the item to be retrieved. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item couldn't be retrieved. - :rtype: dict[str, Any] - - .. admonition:: Example: - - .. literalinclude:: ../samples/examples.py - :start-after: [START update_item] - :end-before: [END update_item] - :language: python - :dedent: 0 - :caption: Get an item from the database and update one of its properties: - :name: update_item - """ - doc_link = self._get_document_link(item) - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - - if partition_key is not None: - request_options["partitionKey"] = self._set_partition_key(partition_key) - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - if post_trigger_include is not None: - request_options["postTriggerInclude"] = post_trigger_include - - result = await self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return result - - @distributed_trace - def read_all_items( - self, - max_item_count=None, # type: Optional[int] - populate_query_metrics=None, # type: Optional[bool] - **kwargs # type: Any - ): - # type: (...) -> Iterable[Dict[str, Any]] - """List all the items in the container. - - :param max_item_count: Max number of items to be returned in the enumeration operation. - :param populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of items (dicts). - :rtype: Iterable[dict[str, Any]] - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if populate_query_metrics is not None: - feed_options["populateQueryMetrics"] = populate_query_metrics - - if hasattr(response_hook, "clear"): - response_hook.clear() - - items = self.client_connection.ReadItems( - collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, items) - return items - - @distributed_trace - def query_items_change_feed( - self, - partition_key_range_id=None, # type: Optional[str] - is_start_from_beginning=False, # type: bool - continuation=None, # type: Optional[str] - max_item_count=None, # type: Optional[int] - **kwargs # type: Any - ): - # type: (...) -> Iterable[Dict[str, Any]] - """Get a sorted list of items that were changed, in the order in which they were modified. - - :param partition_key_range_id: ChangeFeed requests can be executed against specific partition key ranges. - This is used to process the change feed in parallel across multiple consumers. - :param partition_key: partition key at which ChangeFeed requests are targetted. - :param is_start_from_beginning: Get whether change feed should start from - beginning (true) or from current (false). By default it's start from current (false). - :param continuation: e_tag value to be used as continuation for reading change feed. - :param max_item_count: Max number of items to be returned in the enumeration operation. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of items (dicts). - :rtype: Iterable[dict[str, Any]] - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if partition_key_range_id is not None: - feed_options["partitionKeyRangeId"] = partition_key_range_id - partition_key = kwargs.pop("partitionKey", None) - if partition_key is not None: - feed_options["partitionKey"] = partition_key - if is_start_from_beginning is not None: - feed_options["isStartFromBeginning"] = is_start_from_beginning - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if continuation is not None: - feed_options["continuation"] = continuation - - if hasattr(response_hook, "clear"): - response_hook.clear() - - result = self.client_connection.QueryItemsChangeFeed( - self.container_link, options=feed_options, response_hook=response_hook, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return result - - @distributed_trace - def query_items( - self, - query, # type: str - parameters=None, # type: Optional[List[Dict[str, object]]] - partition_key=None, # type: Optional[Any] - enable_cross_partition_query=None, # type: Optional[bool] - max_item_count=None, # type: Optional[int] - enable_scan_in_query=None, # type: Optional[bool] - populate_query_metrics=None, # type: Optional[bool] - **kwargs # type: Any - ): - # type: (...) -> Iterable[Dict[str, Any]] - """Return all results matching the given `query`. - - You can use any value for the container name in the FROM clause, but - often the container name is used. In the examples below, the container - name is "products," and is aliased as "p" for easier referencing in - the WHERE clause. - - :param query: The Azure Cosmos DB SQL query to execute. - :param parameters: Optional array of parameters to the query. - Each parameter is a dict() with 'name' and 'value' keys. - Ignored if no query is provided. - :param partition_key: Specifies the partition key value for the item. - :param enable_cross_partition_query: Allows sending of more than one request to - execute the query in the Azure Cosmos DB service. - More than one request is necessary if the query is not scoped to single partition key value. - :param max_item_count: Max number of items to be returned in the enumeration operation. - :param enable_scan_in_query: Allow scan on the queries which couldn't be served as - indexing was opted out on the requested paths. - :param populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of items (dicts). - :rtype: Iterable[dict[str, Any]] - - .. admonition:: Example: - - .. literalinclude:: ../samples/examples.py - :start-after: [START query_items] - :end-before: [END query_items] - :language: python - :dedent: 0 - :caption: Get all products that have not been discontinued: - :name: query_items - - .. literalinclude:: ../samples/examples.py - :start-after: [START query_items_param] - :end-before: [END query_items_param] - :language: python - :dedent: 0 - :caption: Parameterized query to get all products that have been discontinued: - :name: query_items_param - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if enable_cross_partition_query is not None: - feed_options["enableCrossPartitionQuery"] = enable_cross_partition_query - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if populate_query_metrics is not None: - feed_options["populateQueryMetrics"] = populate_query_metrics - if partition_key is not None: - feed_options["partitionKey"] = self._set_partition_key(partition_key) - if enable_scan_in_query is not None: - feed_options["enableScanInQuery"] = enable_scan_in_query - - if hasattr(response_hook, "clear"): - response_hook.clear() - - items = self.client_connection.QueryItems( - database_or_container_link=self.container_link, - query=query if parameters is None else dict(query=query, parameters=parameters), - options=feed_options, - partition_key=partition_key, - response_hook=response_hook, - **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, items) - return items - - @distributed_trace - def replace_item( - self, - item, # type: Union[str, Dict[str, Any]] - body, # type: Dict[str, Any] - populate_query_metrics=None, # type: Optional[bool] - pre_trigger_include=None, # type: Optional[str] - post_trigger_include=None, # type: Optional[str] - **kwargs # type: Any - ): - # type: (...) -> Dict[str, str] - """Replaces the specified item if it exists in the container. - - If the item does not already exist in the container, an exception is raised. - - :param item: The ID (name) or dict representing item to be replaced. - :param body: A dict-like object representing the item to replace. - :param populate_query_metrics: Enable returning query metrics in response headers. - :param pre_trigger_include: trigger id to be used as pre operation trigger. - :param post_trigger_include: trigger id to be used as post operation trigger. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A dict representing the item after replace went through. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The replace failed or the item with - given id does not exist. - :rtype: dict[str, Any] - """ - item_link = self._get_document_link(item) - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - request_options["disableIdGeneration"] = True - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - if pre_trigger_include is not None: - request_options["preTriggerInclude"] = pre_trigger_include - if post_trigger_include is not None: - request_options["postTriggerInclude"] = post_trigger_include - - result = self.client_connection.ReplaceItem( - document_link=item_link, new_document=body, options=request_options, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return result - - @distributed_trace - def upsert_item( - self, - body, # type: Dict[str, Any] - populate_query_metrics=None, # type: Optional[bool] - pre_trigger_include=None, # type: Optional[str] - post_trigger_include=None, # type: Optional[str] - **kwargs # type: Any - ): - # type: (...) -> Dict[str, str] - """Insert or update the specified item. - - If the item already exists in the container, it is replaced. If the item - does not already exist, it is inserted. - - :param body: A dict-like object representing the item to update or insert. - :param populate_query_metrics: Enable returning query metrics in response headers. - :param pre_trigger_include: trigger id to be used as pre operation trigger. - :param post_trigger_include: trigger id to be used as post operation trigger. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A dict representing the upserted item. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item could not be upserted. - :rtype: dict[str, Any] - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - request_options["disableIdGeneration"] = True - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - if pre_trigger_include is not None: - request_options["preTriggerInclude"] = pre_trigger_include - if post_trigger_include is not None: - request_options["postTriggerInclude"] = post_trigger_include - - result = self.client_connection.UpsertItem( - database_or_container_link=self.container_link, - document=body, - options=request_options, - **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return result - - @distributed_trace - def create_item( - self, - body, # type: Dict[str, Any] - populate_query_metrics=None, # type: Optional[bool] - pre_trigger_include=None, # type: Optional[str] - post_trigger_include=None, # type: Optional[str] - indexing_directive=None, # type: Optional[Any] - **kwargs # type: Any - ): - # type: (...) -> Dict[str, str] - """Create an item in the container. - - To update or replace an existing item, use the - :func:`ContainerProxy.upsert_item` method. - - :param body: A dict-like object representing the item to create. - :param populate_query_metrics: Enable returning query metrics in response headers. - :param pre_trigger_include: trigger id to be used as pre operation trigger. - :param post_trigger_include: trigger id to be used as post operation trigger. - :param indexing_directive: Indicate whether the document should be omitted from indexing. - :keyword bool enable_automatic_id_generation: Enable automatic id generation if no id present. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A dict representing the new item. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists. - :rtype: dict[str, Any] - """ - start = time.time() - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - - request_options["disableAutomaticIdGeneration"] = not kwargs.pop('enable_automatic_id_generation', False) - if populate_query_metrics: - request_options["populateQueryMetrics"] = populate_query_metrics - if pre_trigger_include is not None: - request_options["preTriggerInclude"] = pre_trigger_include - if post_trigger_include is not None: - request_options["postTriggerInclude"] = post_trigger_include - if indexing_directive is not None: - request_options["indexingDirective"] = indexing_directive - - result = self.client_connection.CreateItem( - database_or_container_link=self.container_link, document=body, options=request_options, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - print(f"Create item took {(time.time() - start) * 1000} ms") - print("ASYNC CONTAINER USED") - return result - - @distributed_trace - async def create_item_aio( - self, - body, # type: Dict[str, Any] - populate_query_metrics=None, # type: Optional[bool] - pre_trigger_include=None, # type: Optional[str] - post_trigger_include=None, # type: Optional[str] - indexing_directive=None, # type: Optional[Any] - **kwargs # type: Any - ): - # type: (...) -> Dict[str, str] - """Create an item in the container. - - To update or replace an existing item, use the - :func:`ContainerProxy.upsert_item` method. - - :param body: A dict-like object representing the item to create. - :param populate_query_metrics: Enable returning query metrics in response headers. - :param pre_trigger_include: trigger id to be used as pre operation trigger. - :param post_trigger_include: trigger id to be used as post operation trigger. - :param indexing_directive: Indicate whether the document should be omitted from indexing. - :keyword bool enable_automatic_id_generation: Enable automatic id generation if no id present. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A dict representing the new item. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists. - :rtype: dict[str, Any] - """ - start = time.time() - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - - request_options["disableAutomaticIdGeneration"] = not kwargs.pop('enable_automatic_id_generation', False) - if populate_query_metrics: - request_options["populateQueryMetrics"] = populate_query_metrics - if pre_trigger_include is not None: - request_options["preTriggerInclude"] = pre_trigger_include - if post_trigger_include is not None: - request_options["postTriggerInclude"] = post_trigger_include - if indexing_directive is not None: - request_options["indexingDirective"] = indexing_directive - - result = await self.client_connection.CreateItemAIO( - database_or_container_link=self.container_link, document=body, options=request_options, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) #what is this doing? can't find function - print(f"Create item took {(time.time() - start) * 1000} ms") - return result - - @distributed_trace - def delete_item( - self, - item, # type: Union[Dict[str, Any], str] - partition_key, # type: Any - populate_query_metrics=None, # type: Optional[bool] - pre_trigger_include=None, # type: Optional[str] - post_trigger_include=None, # type: Optional[str] - **kwargs # type: Any - ): - # type: (...) -> None - """Delete the specified item from the container. - - If the item does not already exist in the container, an exception is raised. - - :param item: The ID (name) or dict representing item to be deleted. - :param partition_key: Specifies the partition key value for the item. - :param populate_query_metrics: Enable returning query metrics in response headers. - :param pre_trigger_include: trigger id to be used as pre operation trigger. - :param post_trigger_include: trigger id to be used as post operation trigger. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The item wasn't deleted successfully. - :raises ~azure.cosmos.exceptions.CosmosResourceNotFoundError: The item does not exist in the container. - :rtype: None - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if partition_key is not None: - request_options["partitionKey"] = self._set_partition_key(partition_key) - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - if pre_trigger_include is not None: - request_options["preTriggerInclude"] = pre_trigger_include - if post_trigger_include is not None: - request_options["postTriggerInclude"] = post_trigger_include - - document_link = self._get_document_link(item) - result = self.client_connection.DeleteItem(document_link=document_link, options=request_options, **kwargs) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - - @distributed_trace - def read_offer(self, **kwargs): - # type: (Any) -> Offer - """Read the Offer object for this container. - - If no Offer already exists for the container, an exception is raised. - - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: Offer for the container. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: No offer exists for the container or - the offer could not be retrieved. - :rtype: ~azure.cosmos.Offer - """ - response_hook = kwargs.pop('response_hook', None) - properties = self._get_properties() - link = properties["_self"] - query_spec = { - "query": "SELECT * FROM root r WHERE r.resource=@link", - "parameters": [{"name": "@link", "value": link}], - } - offers = list(self.client_connection.QueryOffers(query_spec, **kwargs)) - if not offers: - raise CosmosResourceNotFoundError( - status_code=StatusCodes.NOT_FOUND, - message="Could not find Offer for container " + self.container_link) - - if response_hook: - response_hook(self.client_connection.last_response_headers, offers) - - return Offer(offer_throughput=offers[0]["content"]["offerThroughput"], properties=offers[0]) - - @distributed_trace - def replace_throughput(self, throughput, **kwargs): - # type: (int, Any) -> Offer - """Replace the container's throughput. - - If no Offer already exists for the container, an exception is raised. - - :param throughput: The throughput to be set (an integer). - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: Offer for the container, updated with new throughput. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: No offer exists for the container - or the offer could not be updated. - :rtype: ~azure.cosmos.Offer - """ - response_hook = kwargs.pop('response_hook', None) - properties = self._get_properties() - link = properties["_self"] - query_spec = { - "query": "SELECT * FROM root r WHERE r.resource=@link", - "parameters": [{"name": "@link", "value": link}], - } - offers = list(self.client_connection.QueryOffers(query_spec, **kwargs)) - if not offers: - raise CosmosResourceNotFoundError( - status_code=StatusCodes.NOT_FOUND, - message="Could not find Offer for container " + self.container_link) - new_offer = offers[0].copy() - new_offer["content"]["offerThroughput"] = throughput - data = self.client_connection.ReplaceOffer(offer_link=offers[0]["_self"], offer=offers[0], **kwargs) - - if response_hook: - response_hook(self.client_connection.last_response_headers, data) - - return Offer(offer_throughput=data["content"]["offerThroughput"], properties=data) - - @distributed_trace - def list_conflicts(self, max_item_count=None, **kwargs): - # type: (Optional[int], Any) -> Iterable[Dict[str, Any]] - """List all the conflicts in the container. - - :param max_item_count: Max number of items to be returned in the enumeration operation. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of conflicts (dicts). - :rtype: Iterable[dict[str, Any]] - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - - result = self.client_connection.ReadConflicts( - collection_link=self.container_link, feed_options=feed_options, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return result - - @distributed_trace - def query_conflicts( - self, - query, # type: str - parameters=None, # type: Optional[List[str]] - enable_cross_partition_query=None, # type: Optional[bool] - partition_key=None, # type: Optional[Any] - max_item_count=None, # type: Optional[int] - **kwargs # type: Any - ): - # type: (...) -> Iterable[Dict[str, Any]] - """Return all conflicts matching a given `query`. - - :param query: The Azure Cosmos DB SQL query to execute. - :param parameters: Optional array of parameters to the query. Ignored if no query is provided. - :param enable_cross_partition_query: Allows sending of more than one request to execute - the query in the Azure Cosmos DB service. - More than one request is necessary if the query is not scoped to single partition key value. - :param partition_key: Specifies the partition key value for the item. - :param max_item_count: Max number of items to be returned in the enumeration operation. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of conflicts (dicts). - :rtype: Iterable[dict[str, Any]] - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if enable_cross_partition_query is not None: - feed_options["enableCrossPartitionQuery"] = enable_cross_partition_query - if partition_key is not None: - feed_options["partitionKey"] = self._set_partition_key(partition_key) - - result = self.client_connection.QueryConflicts( - collection_link=self.container_link, - query=query if parameters is None else dict(query=query, parameters=parameters), - options=feed_options, - **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return result - - @distributed_trace - def get_conflict(self, conflict, partition_key, **kwargs): - # type: (Union[str, Dict[str, Any]], Any, Any) -> Dict[str, str] - """Get the conflict identified by `conflict`. - - :param conflict: The ID (name) or dict representing the conflict to retrieve. - :param partition_key: Partition key for the conflict to retrieve. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A dict representing the retrieved conflict. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given conflict couldn't be retrieved. - :rtype: dict[str, Any] - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if partition_key is not None: - request_options["partitionKey"] = self._set_partition_key(partition_key) - - result = self.client_connection.ReadConflict( - conflict_link=self._get_conflict_link(conflict), options=request_options, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return result - - @distributed_trace - def delete_conflict(self, conflict, partition_key, **kwargs): - # type: (Union[str, Dict[str, Any]], Any, Any) -> None - """Delete a specified conflict from the container. - - If the conflict does not already exist in the container, an exception is raised. - - :param conflict: The ID (name) or dict representing the conflict to be deleted. - :param partition_key: Partition key for the conflict to delete. - :keyword Callable response_hook: A callable invoked with the response metadata. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The conflict wasn't deleted successfully. - :raises ~azure.cosmos.exceptions.CosmosResourceNotFoundError: The conflict does not exist in the container. - :rtype: None - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if partition_key is not None: - request_options["partitionKey"] = self._set_partition_key(partition_key) - - result = self.client_connection.DeleteConflict( - conflict_link=self._get_conflict_link(conflict), options=request_options, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/cosmos_client.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/cosmos_client.py new file mode 100644 index 000000000000..26166d14bf6c --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/cosmos_client.py @@ -0,0 +1,163 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Create, read, and delete databases in the Azure Cosmos DB SQL API service. +""" + +from typing import Any, Dict, Optional, Union, cast, Iterable, List # pylint: disable=unused-import + +import six +from azure.core.tracing.decorator import distributed_trace # type: ignore + +from ..cosmos_client import _parse_connection_str, _build_auth +from ._cosmos_client_connection_async import CosmosClientConnection +from .._base import build_options +from ._retry_utility import ConnectionRetryPolicy +# from .database import DatabaseProxy +from ..documents import ConnectionPolicy, DatabaseAccount +from ..exceptions import CosmosResourceNotFoundError + +__all__ = ("CosmosClient",) + + +def _build_connection_policy(kwargs): + # type: (Dict[str, Any]) -> ConnectionPolicy + # pylint: disable=protected-access + policy = kwargs.pop('connection_policy', None) or ConnectionPolicy() + + # Connection config + policy.RequestTimeout = kwargs.pop('request_timeout', None) or \ + kwargs.pop('connection_timeout', None) or \ + policy.RequestTimeout + policy.ConnectionMode = kwargs.pop('connection_mode', None) or policy.ConnectionMode + policy.ProxyConfiguration = kwargs.pop('proxy_config', None) or policy.ProxyConfiguration + policy.EnableEndpointDiscovery = kwargs.pop('enable_endpoint_discovery', None) or policy.EnableEndpointDiscovery + policy.PreferredLocations = kwargs.pop('preferred_locations', None) or policy.PreferredLocations + policy.UseMultipleWriteLocations = kwargs.pop('multiple_write_locations', None) or \ + policy.UseMultipleWriteLocations + + # SSL config + verify = kwargs.pop('connection_verify', None) + policy.DisableSSLVerification = not bool(verify if verify is not None else True) + ssl = kwargs.pop('ssl_config', None) or policy.SSLConfiguration + if ssl: + ssl.SSLCertFile = kwargs.pop('connection_cert', None) or ssl.SSLCertFile + ssl.SSLCaCerts = verify or ssl.SSLCaCerts + policy.SSLConfiguration = ssl + + # Retry config + retry = kwargs.pop('retry_options', None) or policy.RetryOptions + total_retries = kwargs.pop('retry_total', None) + retry._max_retry_attempt_count = total_retries or retry._max_retry_attempt_count + retry._fixed_retry_interval_in_milliseconds = kwargs.pop('retry_fixed_interval', None) or \ + retry._fixed_retry_interval_in_milliseconds + max_backoff = kwargs.pop('retry_backoff_max', None) + retry._max_wait_time_in_seconds = max_backoff or retry._max_wait_time_in_seconds + policy.RetryOptions = retry + connection_retry = kwargs.pop('connection_retry_policy', None) or policy.ConnectionRetryConfiguration + if not connection_retry: + connection_retry = ConnectionRetryPolicy( + retry_total=total_retries, + retry_connect=kwargs.pop('retry_connect', None), + retry_read=kwargs.pop('retry_read', None), + retry_status=kwargs.pop('retry_status', None), + retry_backoff_max=max_backoff, + retry_on_status_codes=kwargs.pop('retry_on_status_codes', []), + retry_backoff_factor=kwargs.pop('retry_backoff_factor', 0.8), + ) + policy.ConnectionRetryConfiguration = connection_retry + + return policy + + + +class CosmosClient(object): + """A client-side logical representation of an Azure Cosmos DB account. + + Use this client to configure and execute requests to the Azure Cosmos DB service. + + :param str url: The URL of the Cosmos DB account. + :param credential: Can be the account key, or a dictionary of resource tokens. + :type credential: str or dict[str, str] + :param str consistency_level: Consistency level to use for the session. The default value is "Session". + + .. admonition:: Example: + + .. literalinclude:: ../samples/examples.py + :start-after: [START create_client] + :end-before: [END create_client] + :language: python + :dedent: 0 + :caption: Create a new instance of the Cosmos DB client: + :name: create_client + """ + + def __init__(self, url, credential, **kwargs): + # type: (str, Any, str, Any) -> None + """Instantiate a new CosmosClient.""" + auth = _build_auth(credential) + consistency_level = kwargs.get('consistency_level', 'Session') + connection_policy = _build_connection_policy(kwargs) + self.client_connection = CosmosClientConnection( + url, + auth=auth, + consistency_level=consistency_level, + connection_policy=connection_policy, + **kwargs + ) + + def __repr__(self): + # type () -> str + return "".format(self.client_connection.url_connection)[:1024] + + async def __aenter__(self): + await self.client_connection.pipeline_client.__aenter__() + await self.client_connection._setup() + return self + + async def __aexit__(self, *args): + return await self.client_connection.pipeline_client.__aexit__(*args) + + async def close(self): + await self.__aexit__() + + @classmethod + def from_connection_string(cls, conn_str, credential=None, consistency_level="Session", **kwargs): + # type: (str, Optional[Any], str, Any) -> CosmosClient + """Create a CosmosClient instance from a connection string. + + This can be retrieved from the Azure portal.For full list of optional + keyword arguments, see the CosmosClient constructor. + + :param str conn_str: The connection string. + :param credential: Alternative credentials to use instead of the key + provided in the connection string. + :type credential: str or dict(str, str) + :param str consistency_level: + Consistency level to use for the session. The default value is "Session". + """ + settings = _parse_connection_str(conn_str, credential) + return cls( + url=settings['AccountEndpoint'], + credential=credential or settings['AccountKey'], + consistency_level=consistency_level, + **kwargs + ) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/cosmos_client_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/cosmos_client_async.py deleted file mode 100644 index 879cafcf7f9f..000000000000 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/cosmos_client_async.py +++ /dev/null @@ -1,456 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) 2014 Microsoft Corporation - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -"""Create, read, and delete databases in the Azure Cosmos DB SQL API service. -""" - -from typing import Any, Dict, Optional, Union, cast, Iterable, List # pylint: disable=unused-import - -import six -from azure.core.tracing.decorator import distributed_trace # type: ignore - -from ._cosmos_client_connection_async import CosmosClientConnection -from .._base import build_options -from .._retry_utility import ConnectionRetryPolicy -from .database_async import DatabaseProxy -from ..documents import ConnectionPolicy, DatabaseAccount -from ..exceptions import CosmosResourceNotFoundError - -__all__ = ("CosmosClient",) - - -def _parse_connection_str(conn_str, credential): - # type: (str, Optional[Any]) -> Dict[str, str] - conn_str = conn_str.rstrip(";") - conn_settings = dict( # type: ignore # pylint: disable=consider-using-dict-comprehension - s.split("=", 1) for s in conn_str.split(";") - ) - if 'AccountEndpoint' not in conn_settings: - raise ValueError("Connection string missing setting 'AccountEndpoint'.") - if not credential and 'AccountKey' not in conn_settings: - raise ValueError("Connection string missing setting 'AccountKey'.") - return conn_settings - - -def _build_auth(credential): - # type: (Any) -> Dict[str, Any] - auth = {} - if isinstance(credential, six.string_types): - auth['masterKey'] = credential - elif isinstance(credential, dict): - if any(k for k in credential.keys() if k in ['masterKey', 'resourceTokens', 'permissionFeed']): - return credential # Backwards compatible - auth['resourceTokens'] = credential # type: ignore - elif hasattr(credential, '__iter__'): - auth['permissionFeed'] = credential - else: - raise TypeError( - "Unrecognized credential type. Please supply the master key as str, " - "or a dictionary or resource tokens, or a list of permissions.") - return auth - - -def _build_connection_policy(kwargs): - # type: (Dict[str, Any]) -> ConnectionPolicy - # pylint: disable=protected-access - policy = kwargs.pop('connection_policy', None) or ConnectionPolicy() - - # Connection config - policy.RequestTimeout = kwargs.pop('request_timeout', None) or \ - kwargs.pop('connection_timeout', None) or \ - policy.RequestTimeout - policy.ConnectionMode = kwargs.pop('connection_mode', None) or policy.ConnectionMode - policy.ProxyConfiguration = kwargs.pop('proxy_config', None) or policy.ProxyConfiguration - policy.EnableEndpointDiscovery = kwargs.pop('enable_endpoint_discovery', None) or policy.EnableEndpointDiscovery - policy.PreferredLocations = kwargs.pop('preferred_locations', None) or policy.PreferredLocations - policy.UseMultipleWriteLocations = kwargs.pop('multiple_write_locations', None) or \ - policy.UseMultipleWriteLocations - - # SSL config - verify = kwargs.pop('connection_verify', None) - policy.DisableSSLVerification = not bool(verify if verify is not None else True) - ssl = kwargs.pop('ssl_config', None) or policy.SSLConfiguration - if ssl: - ssl.SSLCertFile = kwargs.pop('connection_cert', None) or ssl.SSLCertFile - ssl.SSLCaCerts = verify or ssl.SSLCaCerts - policy.SSLConfiguration = ssl - - # Retry config - retry = kwargs.pop('retry_options', None) or policy.RetryOptions - total_retries = kwargs.pop('retry_total', None) - retry._max_retry_attempt_count = total_retries or retry._max_retry_attempt_count - retry._fixed_retry_interval_in_milliseconds = kwargs.pop('retry_fixed_interval', None) or \ - retry._fixed_retry_interval_in_milliseconds - max_backoff = kwargs.pop('retry_backoff_max', None) - retry._max_wait_time_in_seconds = max_backoff or retry._max_wait_time_in_seconds - policy.RetryOptions = retry - connection_retry = kwargs.pop('connection_retry_policy', None) or policy.ConnectionRetryConfiguration - if not connection_retry: - connection_retry = ConnectionRetryPolicy( - retry_total=total_retries, - retry_connect=kwargs.pop('retry_connect', None), - retry_read=kwargs.pop('retry_read', None), - retry_status=kwargs.pop('retry_status', None), - retry_backoff_max=max_backoff, - retry_on_status_codes=kwargs.pop('retry_on_status_codes', []), - retry_backoff_factor=kwargs.pop('retry_backoff_factor', 0.8), - ) - policy.ConnectionRetryConfiguration = connection_retry - - return policy - - - -class AsyncCosmosClient(object): - """A client-side logical representation of an Azure Cosmos DB account. - - Use this client to configure and execute requests to the Azure Cosmos DB service. - - :param str url: The URL of the Cosmos DB account. - :param credential: Can be the account key, or a dictionary of resource tokens. - :type credential: str or dict[str, str] - :param str consistency_level: Consistency level to use for the session. The default value is "Session". - :keyword int timeout: An absolute timeout in seconds, for the combined HTTP request and response processing. - :keyword int request_timeout: The HTTP request timeout in milliseconds. - :keyword str connection_mode: The connection mode for the client - currently only supports 'Gateway'. - :keyword proxy_config: Connection proxy configuration. - :paramtype proxy_config: ~azure.cosmos.ProxyConfiguration - :keyword ssl_config: Connection SSL configuration. - :paramtype ssl_config: ~azure.cosmos.SSLConfiguration - :keyword bool connection_verify: Whether to verify the connection, default value is True. - :keyword str connection_cert: An alternative certificate to verify the connection. - :keyword int retry_total: Maximum retry attempts. - :keyword int retry_backoff_max: Maximum retry wait time in seconds. - :keyword int retry_fixed_interval: Fixed retry interval in milliseconds. - :keyword int retry_read: Maximum number of socket read retry attempts. - :keyword int retry_connect: Maximum number of connection error retry attempts. - :keyword int retry_status: Maximum number of retry attempts on error status codes. - :keyword list[int] retry_on_status_codes: A list of specific status codes to retry on. - :keyword float retry_backoff_factor: Factor to calculate wait time between retry attempts. - :keyword bool enable_endpoint_discovery: Enable endpoint discovery for - geo-replicated database accounts. (Default: True) - :keyword list[str] preferred_locations: The preferred locations for geo-replicated database accounts. - - .. admonition:: Example: - - .. literalinclude:: ../samples/examples.py - :start-after: [START create_client] - :end-before: [END create_client] - :language: python - :dedent: 0 - :caption: Create a new instance of the Cosmos DB client: - :name: create_client - """ - - def __init__(self, url, credential, consistency_level="Session", **kwargs): - # type: (str, Any, str, Any) -> None - """Instantiate a new CosmosClient.""" - auth = _build_auth(credential) - connection_policy = _build_connection_policy(kwargs) - self.client_connection = CosmosClientConnection( - url, auth=auth, consistency_level=consistency_level, connection_policy=connection_policy, **kwargs - ) - - def __repr__(self): # pylint:disable=client-method-name-no-double-underscore - # type () -> str - return "".format(self.client_connection.url_connection)[:1024] - - def __enter__(self): - self.client_connection.pipeline_client.__enter__() - return self - - def __exit__(self, *args): - return self.client_connection.pipeline_client.__exit__(*args) - - @classmethod - def from_connection_string(cls, conn_str, credential=None, consistency_level="Session", **kwargs): - # type: (str, Optional[Any], str, Any) -> CosmosClient - """Create a CosmosClient instance from a connection string. - - This can be retrieved from the Azure portal.For full list of optional - keyword arguments, see the CosmosClient constructor. - - :param str conn_str: The connection string. - :param credential: Alternative credentials to use instead of the key - provided in the connection string. - :type credential: str or dict(str, str) - :param str consistency_level: - Consistency level to use for the session. The default value is "Session". - """ - settings = _parse_connection_str(conn_str, credential) - return cls( - url=settings['AccountEndpoint'], - credential=credential or settings['AccountKey'], - consistency_level=consistency_level, - **kwargs - ) - - @staticmethod - def _get_database_link(database_or_id): - # type: (Union[DatabaseProxy, str, Dict[str, str]]) -> str - if isinstance(database_or_id, six.string_types): - return "dbs/{}".format(database_or_id) - try: - return cast("DatabaseProxy", database_or_id).database_link - except AttributeError: - pass - database_id = cast("Dict[str, str]", database_or_id)["id"] - return "dbs/{}".format(database_id) - - @distributed_trace - def create_database( # pylint: disable=redefined-builtin - self, - id, # type: str - populate_query_metrics=None, # type: Optional[bool] - offer_throughput=None, # type: Optional[int] - **kwargs # type: Any - ): - # type: (...) -> DatabaseProxy - """ - Create a new database with the given ID (name). - - :param id: ID (name) of the database to create. - :param bool populate_query_metrics: Enable returning query metrics in response headers. - :param int offer_throughput: The provisioned throughput for this offer. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A DatabaseProxy instance representing the new database. - :rtype: ~azure.cosmos.DatabaseProxy - :raises ~azure.cosmos.exceptions.CosmosResourceExistsError: Database with the given ID already exists. - - .. admonition:: Example: - - .. literalinclude:: ../samples/examples.py - :start-after: [START create_database] - :end-before: [END create_database] - :language: python - :dedent: 0 - :caption: Create a database in the Cosmos DB account: - :name: create_database - """ - - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - if offer_throughput is not None: - request_options["offerThroughput"] = offer_throughput - - result = self.client_connection.CreateDatabase(database=dict(id=id), options=request_options, **kwargs) - if response_hook: - response_hook(self.client_connection.last_response_headers) - return DatabaseProxy(self.client_connection, id=result["id"], properties=result) - - @distributed_trace - def create_database_if_not_exists( # pylint: disable=redefined-builtin - self, - id, # type: str - populate_query_metrics=None, # type: Optional[bool] - offer_throughput=None, # type: Optional[int] - **kwargs # type: Any - ): - # type: (...) -> DatabaseProxy - """ - Create the database if it does not exist already. - - If the database already exists, the existing settings are returned. - - ..note:: - This function does not check or update existing database settings or - offer throughput if they differ from what is passed in. - - :param id: ID (name) of the database to read or create. - :param bool populate_query_metrics: Enable returning query metrics in response headers. - :param int offer_throughput: The provisioned throughput for this offer. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A DatabaseProxy instance representing the database. - :rtype: ~azure.cosmos.DatabaseProxy - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The database read or creation failed. - """ - try: - database_proxy = self.get_database_client(id) - database_proxy.read( - populate_query_metrics=populate_query_metrics, - **kwargs - ) - return database_proxy - except CosmosResourceNotFoundError: - return self.create_database( - id, - populate_query_metrics=populate_query_metrics, - offer_throughput=offer_throughput, - **kwargs - ) - - def get_database_client(self, database): - # type: (Union[str, DatabaseProxy, Dict[str, Any]]) -> DatabaseProxy - """Retrieve an existing database with the ID (name) `id`. - - :param database: The ID (name), dict representing the properties or - `DatabaseProxy` instance of the database to read. - :type database: str or dict(str, str) or ~azure.cosmos.DatabaseProxy - :returns: A `DatabaseProxy` instance representing the retrieved database. - :rtype: ~azure.cosmos.DatabaseProxy - """ - if isinstance(database, DatabaseProxy): - id_value = database.id - else: - try: - id_value = database["id"] - except TypeError: - id_value = database - - return DatabaseProxy(self.client_connection, id_value) - - @distributed_trace - def list_databases( - self, - max_item_count=None, # type: Optional[int] - populate_query_metrics=None, # type: Optional[bool] - **kwargs # type: Any - ): - # type: (...) -> Iterable[Dict[str, Any]] - """List the databases in a Cosmos DB SQL database account. - - :param int max_item_count: Max number of items to be returned in the enumeration operation. - :param bool populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of database properties (dicts). - :rtype: Iterable[dict[str, str]] - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if populate_query_metrics is not None: - feed_options["populateQueryMetrics"] = populate_query_metrics - - result = self.client_connection.ReadDatabases(options=feed_options, **kwargs) - if response_hook: - response_hook(self.client_connection.last_response_headers) - return result - - @distributed_trace - def query_databases( - self, - query=None, # type: Optional[str] - parameters=None, # type: Optional[List[str]] - enable_cross_partition_query=None, # type: Optional[bool] - max_item_count=None, # type: Optional[int] - populate_query_metrics=None, # type: Optional[bool] - **kwargs # type: Any - ): - # type: (...) -> Iterable[Dict[str, Any]] - """Query the databases in a Cosmos DB SQL database account. - - :param str query: The Azure Cosmos DB SQL query to execute. - :param list[str] parameters: Optional array of parameters to the query. Ignored if no query is provided. - :param bool enable_cross_partition_query: Allow scan on the queries which couldn't be - served as indexing was opted out on the requested paths. - :param int max_item_count: Max number of items to be returned in the enumeration operation. - :param bool populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of database properties (dicts). - :rtype: Iterable[dict[str, str]] - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if enable_cross_partition_query is not None: - feed_options["enableCrossPartitionQuery"] = enable_cross_partition_query - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if populate_query_metrics is not None: - feed_options["populateQueryMetrics"] = populate_query_metrics - - if query: - # This is currently eagerly evaluated in order to capture the headers - # from the call. - # (just returning a generator did not initiate the first network call, so - # the headers were misleading) - # This needs to change for "real" implementation - query = query if parameters is None else dict(query=query, parameters=parameters) # type: ignore - result = self.client_connection.QueryDatabases(query=query, options=feed_options, **kwargs) - else: - result = self.client_connection.ReadDatabases(options=feed_options, **kwargs) - if response_hook: - response_hook(self.client_connection.last_response_headers) - return result - - @distributed_trace - def delete_database( - self, - database, # type: Union[str, DatabaseProxy, Dict[str, Any]] - populate_query_metrics=None, # type: Optional[bool] - **kwargs # type: Any - ): - # type: (...) -> None - """Delete the database with the given ID (name). - - :param database: The ID (name), dict representing the properties or :class:`DatabaseProxy` - instance of the database to delete. - :type database: str or dict(str, str) or ~azure.cosmos.DatabaseProxy - :param bool populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the database couldn't be deleted. - :rtype: None - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - - database_link = self._get_database_link(database) - self.client_connection.DeleteDatabase(database_link, options=request_options, **kwargs) - if response_hook: - response_hook(self.client_connection.last_response_headers) - - @distributed_trace - def get_database_account(self, **kwargs): - # type: (Any) -> DatabaseAccount - """Retrieve the database account information. - - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A `DatabaseAccount` instance representing the Cosmos DB Database Account. - :rtype: ~azure.cosmos.DatabaseAccount - """ - response_hook = kwargs.pop('response_hook', None) - result = self.client_connection.GetDatabaseAccount(**kwargs) - if response_hook: - response_hook(self.client_connection.last_response_headers) - return result diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/database_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/database_async.py deleted file mode 100644 index cbb1e0ab6902..000000000000 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/database_async.py +++ /dev/null @@ -1,768 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) 2014 Microsoft Corporation - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -"""Interact with databases in the Azure Cosmos DB SQL API service. -""" - -from typing import Any, List, Dict, Union, cast, Iterable, Optional - -import warnings -import six -from azure.core.tracing.decorator import distributed_trace # type: ignore - -from ._cosmos_client_connection_async import CosmosClientConnection -from .._base import build_options -from .container_async import ContainerProxy -from ..offer import Offer -from ..http_constants import StatusCodes -from ..exceptions import CosmosResourceNotFoundError -from ..user import UserProxy -from ..documents import IndexingMode - -__all__ = ("DatabaseProxy",) - -# pylint: disable=protected-access -# pylint: disable=missing-client-constructor-parameter-credential,missing-client-constructor-parameter-kwargs - - -class DatabaseProxy(object): - """An interface to interact with a specific database. - - This class should not be instantiated directly. Instead use the - :func:`CosmosClient.get_database_client` method. - - A database contains one or more containers, each of which can contain items, - stored procedures, triggers, and user-defined functions. - - A database can also have associated users, each of which is configured with - a set of permissions for accessing certain containers, stored procedures, - triggers, user-defined functions, or items. - - :ivar id: The ID (name) of the database. - - An Azure Cosmos DB SQL API database has the following system-generated - properties. These properties are read-only: - - * `_rid`: The resource ID. - * `_ts`: When the resource was last updated. The value is a timestamp. - * `_self`: The unique addressable URI for the resource. - * `_etag`: The resource etag required for optimistic concurrency control. - * `_colls`: The addressable path of the collections resource. - * `_users`: The addressable path of the users resource. - """ - - def __init__(self, client_connection, id, properties=None): # pylint: disable=redefined-builtin - # type: (CosmosClientConnection, str, Dict[str, Any]) -> None - """ - :param ClientSession client_connection: Client from which this database was retrieved. - :param str id: ID (name) of the database. - """ - self.client_connection = client_connection - self.id = id - self.database_link = u"dbs/{}".format(self.id) - self._properties = properties - - def __repr__(self): - # type () -> str - return "".format(self.database_link)[:1024] - - @staticmethod - def _get_container_id(container_or_id): - # type: (Union[str, ContainerProxy, Dict[str, Any]]) -> str - if isinstance(container_or_id, six.string_types): - return container_or_id - try: - return cast("ContainerProxy", container_or_id).id - except AttributeError: - pass - return cast("Dict[str, str]", container_or_id)["id"] - - def _get_container_link(self, container_or_id): - # type: (Union[str, ContainerProxy, Dict[str, Any]]) -> str - return u"{}/colls/{}".format(self.database_link, self._get_container_id(container_or_id)) - - def _get_user_link(self, user_or_id): - # type: (Union[UserProxy, str, Dict[str, Any]]) -> str - if isinstance(user_or_id, six.string_types): - return u"{}/users/{}".format(self.database_link, user_or_id) - try: - return cast("UserProxy", user_or_id).user_link - except AttributeError: - pass - return u"{}/users/{}".format(self.database_link, cast("Dict[str, str]", user_or_id)["id"]) - - def _get_properties(self): - # type: () -> Dict[str, Any] - if self._properties is None: - self._properties = self.read() - return self._properties - - @distributed_trace - def read(self, populate_query_metrics=None, **kwargs): - # type: (Optional[bool], Any) -> Dict[str, Any] - """Read the database properties. - - :param bool populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword Callable response_hook: A callable invoked with the response metadata. - :rtype: Dict[Str, Any] - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the given database couldn't be retrieved. - """ - # TODO this helper function should be extracted from CosmosClient - from .cosmos_client_async import CosmosClient - - database_link = CosmosClient._get_database_link(self) - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - - self._properties = self.client_connection.ReadDatabase( - database_link, options=request_options, **kwargs - ) - - if response_hook: - response_hook(self.client_connection.last_response_headers, self._properties) - - return cast('Dict[str, Any]', self._properties) - - @distributed_trace - async def create_container( - self, - id, # type: str # pylint: disable=redefined-builtin - partition_key, # type: Any - indexing_policy=None, # type: Optional[Dict[str, Any]] - default_ttl=None, # type: Optional[int] - populate_query_metrics=None, # type: Optional[bool] - offer_throughput=None, # type: Optional[int] - unique_key_policy=None, # type: Optional[Dict[str, Any]] - conflict_resolution_policy=None, # type: Optional[Dict[str, Any]] - **kwargs # type: Any - ): - # type: (...) -> ContainerProxy - """Create a new container with the given ID (name). - - If a container with the given ID already exists, a CosmosResourceExistsError is raised. - - :param id: ID (name) of container to create. - :param partition_key: The partition key to use for the container. - :param indexing_policy: The indexing policy to apply to the container. - :param default_ttl: Default time to live (TTL) for items in the container. If unspecified, items do not expire. - :param populate_query_metrics: Enable returning query metrics in response headers. - :param offer_throughput: The provisioned throughput for this offer. - :param unique_key_policy: The unique key policy to apply to the container. - :param conflict_resolution_policy: The conflict resolution policy to apply to the container. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :keyword analytical_storage_ttl: Analytical store time to live (TTL) for items in the container. A value of - None leaves analytical storage off and a value of -1 turns analytical storage on with no TTL. Please - note that analytical storage can only be enabled on Synapse Link enabled accounts. - :returns: A `ContainerProxy` instance representing the new container. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The container creation failed. - :rtype: ~azure.cosmos.ContainerProxy - - .. admonition:: Example: - - .. literalinclude:: ../samples/examples.py - :start-after: [START create_container] - :end-before: [END create_container] - :language: python - :dedent: 0 - :caption: Create a container with default settings: - :name: create_container - - .. literalinclude:: ../samples/examples.py - :start-after: [START create_container_with_settings] - :end-before: [END create_container_with_settings] - :language: python - :dedent: 0 - :caption: Create a container with specific settings; in this case, a custom partition key: - :name: create_container_with_settings - """ - definition = dict(id=id) # type: Dict[str, Any] - if partition_key is not None: - definition["partitionKey"] = partition_key - if indexing_policy is not None: - if indexing_policy.get("indexingMode") is IndexingMode.Lazy: - warnings.warn( - "Lazy indexing mode has been deprecated. Mode will be set to consistent indexing by the backend.", - DeprecationWarning - ) - definition["indexingPolicy"] = indexing_policy - if default_ttl is not None: - definition["defaultTtl"] = default_ttl - if unique_key_policy is not None: - definition["uniqueKeyPolicy"] = unique_key_policy - if conflict_resolution_policy is not None: - definition["conflictResolutionPolicy"] = conflict_resolution_policy - - analytical_storage_ttl = kwargs.pop("analytical_storage_ttl", None) - if analytical_storage_ttl is not None: - definition["analyticalStorageTtl"] = analytical_storage_ttl - - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - if offer_throughput is not None: - request_options["offerThroughput"] = offer_throughput - - data = self.client_connection.CreateContainer( - database_link=self.database_link, collection=definition, options=request_options, **kwargs - ) - - if response_hook: - response_hook(self.client_connection.last_response_headers, data) - - return ContainerProxy(self.client_connection, self.database_link, data["id"], properties=data) - - @distributed_trace - def create_container_if_not_exists( - self, - id, # type: str # pylint: disable=redefined-builtin - partition_key, # type: Any - indexing_policy=None, # type: Optional[Dict[str, Any]] - default_ttl=None, # type: Optional[int] - populate_query_metrics=None, # type: Optional[bool] - offer_throughput=None, # type: Optional[int] - unique_key_policy=None, # type: Optional[Dict[str, Any]] - conflict_resolution_policy=None, # type: Optional[Dict[str, Any]] - **kwargs # type: Any - ): - # type: (...) -> ContainerProxy - """Create a container if it does not exist already. - - If the container already exists, the existing settings are returned. - Note: it does not check or update the existing container settings or offer throughput - if they differ from what was passed into the method. - - :param id: ID (name) of container to read or create. - :param partition_key: The partition key to use for the container. - :param indexing_policy: The indexing policy to apply to the container. - :param default_ttl: Default time to live (TTL) for items in the container. If unspecified, items do not expire. - :param populate_query_metrics: Enable returning query metrics in response headers. - :param offer_throughput: The provisioned throughput for this offer. - :param unique_key_policy: The unique key policy to apply to the container. - :param conflict_resolution_policy: The conflict resolution policy to apply to the container. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :keyword analytical_storage_ttl: Analytical store time to live (TTL) for items in the container. A value of - None leaves analytical storage off and a value of -1 turns analytical storage on with no TTL. Please - note that analytical storage can only be enabled on Synapse Link enabled accounts. - :returns: A `ContainerProxy` instance representing the container. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The container read or creation failed. - :rtype: ~azure.cosmos.ContainerProxy - """ - - analytical_storage_ttl = kwargs.pop("analytical_storage_ttl", None) - try: - container_proxy = self.get_container_client(id) - container_proxy.read( - populate_query_metrics=populate_query_metrics, - **kwargs - ) - return container_proxy - except CosmosResourceNotFoundError: - return self.create_container( - id=id, - partition_key=partition_key, - indexing_policy=indexing_policy, - default_ttl=default_ttl, - populate_query_metrics=populate_query_metrics, - offer_throughput=offer_throughput, - unique_key_policy=unique_key_policy, - conflict_resolution_policy=conflict_resolution_policy, - analytical_storage_ttl=analytical_storage_ttl - ) - - @distributed_trace - def delete_container( - self, - container, # type: Union[str, ContainerProxy, Dict[str, Any]] - populate_query_metrics=None, # type: Optional[bool] - **kwargs # type: Any - ): - # type: (...) -> None - """Delete a container. - - :param container: The ID (name) of the container to delete. You can either - pass in the ID of the container to delete, a :class:`ContainerProxy` instance or - a dict representing the properties of the container. - :param populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword Callable response_hook: A callable invoked with the response metadata. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the container couldn't be deleted. - :rtype: None - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - - collection_link = self._get_container_link(container) - result = self.client_connection.DeleteContainer(collection_link, options=request_options, **kwargs) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - - def get_container_client(self, container): - # type: (Union[str, ContainerProxy, Dict[str, Any]]) -> ContainerProxy - """Get a `ContainerProxy` for a container with specified ID (name). - - :param container: The ID (name) of the container, a :class:`ContainerProxy` instance, - or a dict representing the properties of the container to be retrieved. - :rtype: ~azure.cosmos.ContainerProxy - - .. admonition:: Example: - - .. literalinclude:: ../samples/examples.py - :start-after: [START get_container] - :end-before: [END get_container] - :language: python - :dedent: 0 - :caption: Get an existing container, handling a failure if encountered: - :name: get_container - """ - if isinstance(container, ContainerProxy): - id_value = container.id - else: - try: - id_value = container["id"] - except TypeError: - id_value = container - - return ContainerProxy(self.client_connection, self.database_link, id_value) - - @distributed_trace - def list_containers(self, max_item_count=None, populate_query_metrics=None, **kwargs): - # type: (Optional[int], Optional[bool], Any) -> Iterable[Dict[str, Any]] - """List the containers in the database. - - :param max_item_count: Max number of items to be returned in the enumeration operation. - :param populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of container properties (dicts). - :rtype: Iterable[dict[str, Any]] - - .. admonition:: Example: - - .. literalinclude:: ../samples/examples.py - :start-after: [START list_containers] - :end-before: [END list_containers] - :language: python - :dedent: 0 - :caption: List all containers in the database: - :name: list_containers - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if populate_query_metrics is not None: - feed_options["populateQueryMetrics"] = populate_query_metrics - - result = self.client_connection.ReadContainers( - database_link=self.database_link, options=feed_options, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return response_hook - - @distributed_trace - def query_containers( - self, - query=None, # type: Optional[str] - parameters=None, # type: Optional[List[str]] - max_item_count=None, # type: Optional[int] - populate_query_metrics=None, # type: Optional[bool] - **kwargs # type: Any - ): - # type: (...) -> Iterable[Dict[str, Any]] - """List the properties for containers in the current database. - - :param query: The Azure Cosmos DB SQL query to execute. - :param parameters: Optional array of parameters to the query. Ignored if no query is provided. - :param max_item_count: Max number of items to be returned in the enumeration operation. - :param populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of container properties (dicts). - :rtype: Iterable[dict[str, Any]] - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if populate_query_metrics is not None: - feed_options["populateQueryMetrics"] = populate_query_metrics - - result = self.client_connection.QueryContainers( - database_link=self.database_link, - query=query if parameters is None else dict(query=query, parameters=parameters), - options=feed_options, - **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return result - - @distributed_trace - def replace_container( - self, - container, # type: Union[str, ContainerProxy, Dict[str, Any]] - partition_key, # type: Any - indexing_policy=None, # type: Optional[Dict[str, Any]] - default_ttl=None, # type: Optional[int] - conflict_resolution_policy=None, # type: Optional[Dict[str, Any]] - populate_query_metrics=None, # type: Optional[bool] - **kwargs # type: Any - ): - # type: (...) -> ContainerProxy - """Reset the properties of the container. - - Property changes are persisted immediately. Any properties not specified - will be reset to their default values. - - :param container: The ID (name), dict representing the properties or - :class:`ContainerProxy` instance of the container to be replaced. - :param partition_key: The partition key to use for the container. - :param indexing_policy: The indexing policy to apply to the container. - :param default_ttl: Default time to live (TTL) for items in the container. - If unspecified, items do not expire. - :param conflict_resolution_policy: The conflict resolution policy to apply to the container. - :param populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource - has changed, and act according to the condition specified by the `match_condition` parameter. - :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag. - :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. - :keyword Callable response_hook: A callable invoked with the response metadata. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Raised if the container couldn't be replaced. - This includes if the container with given id does not exist. - :returns: A `ContainerProxy` instance representing the container after replace completed. - :rtype: ~azure.cosmos.ContainerProxy - - .. admonition:: Example: - - .. literalinclude:: ../samples/examples.py - :start-after: [START reset_container_properties] - :end-before: [END reset_container_properties] - :language: python - :dedent: 0 - :caption: Reset the TTL property on a container, and display the updated properties: - :name: reset_container_properties - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if populate_query_metrics is not None: - request_options["populateQueryMetrics"] = populate_query_metrics - - container_id = self._get_container_id(container) - container_link = self._get_container_link(container_id) - parameters = { - key: value - for key, value in { - "id": container_id, - "partitionKey": partition_key, - "indexingPolicy": indexing_policy, - "defaultTtl": default_ttl, - "conflictResolutionPolicy": conflict_resolution_policy, - }.items() - if value is not None - } - - container_properties = self.client_connection.ReplaceContainer( - container_link, collection=parameters, options=request_options, **kwargs - ) - - if response_hook: - response_hook(self.client_connection.last_response_headers, container_properties) - - return ContainerProxy( - self.client_connection, self.database_link, container_properties["id"], properties=container_properties - ) - - @distributed_trace - def list_users(self, max_item_count=None, **kwargs): - # type: (Optional[int], Any) -> Iterable[Dict[str, Any]] - """List all the users in the container. - - :param max_item_count: Max number of users to be returned in the enumeration operation. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of user properties (dicts). - :rtype: Iterable[dict[str, Any]] - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - - result = self.client_connection.ReadUsers( - database_link=self.database_link, options=feed_options, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return result - - @distributed_trace - def query_users(self, query, parameters=None, max_item_count=None, **kwargs): - # type: (str, Optional[List[str]], Optional[int], Any) -> Iterable[Dict[str, Any]] - """Return all users matching the given `query`. - - :param query: The Azure Cosmos DB SQL query to execute. - :param parameters: Optional array of parameters to the query. Ignored if no query is provided. - :param max_item_count: Max number of users to be returned in the enumeration operation. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: An Iterable of user properties (dicts). - :rtype: Iterable[str, Any] - """ - feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - - result = self.client_connection.QueryUsers( - database_link=self.database_link, - query=query if parameters is None else dict(query=query, parameters=parameters), - options=feed_options, - **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - return result - - def get_user_client(self, user): - # type: (Union[str, UserProxy, Dict[str, Any]]) -> UserProxy - """Get a `UserProxy` for a user with specified ID. - - :param user: The ID (name), dict representing the properties or :class:`UserProxy` - instance of the user to be retrieved. - :returns: A `UserProxy` instance representing the retrieved user. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the given user couldn't be retrieved. - :rtype: ~azure.cosmos.UserProxy - """ - if isinstance(user, UserProxy): - id_value = user.id - else: - try: - id_value = user["id"] - except TypeError: - id_value = user - - return UserProxy(client_connection=self.client_connection, id=id_value, database_link=self.database_link) - - @distributed_trace - def create_user(self, body, **kwargs): - # type: (Dict[str, Any], Any) -> UserProxy - """Create a new user in the container. - - To update or replace an existing user, use the - :func:`ContainerProxy.upsert_user` method. - - :param body: A dict-like object with an `id` key and value representing the user to be created. - The user ID must be unique within the database, and consist of no more than 255 characters. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A `UserProxy` instance representing the new user. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the given user couldn't be created. - :rtype: ~azure.cosmos.UserProxy - - .. admonition:: Example: - - .. literalinclude:: ../samples/examples.py - :start-after: [START create_user] - :end-before: [END create_user] - :language: python - :dedent: 0 - :caption: Create a database user: - :name: create_user - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - - user = self.client_connection.CreateUser( - database_link=self.database_link, user=body, options=request_options, **kwargs) - - if response_hook: - response_hook(self.client_connection.last_response_headers, user) - - return UserProxy( - client_connection=self.client_connection, id=user["id"], database_link=self.database_link, properties=user - ) - - @distributed_trace - def upsert_user(self, body, **kwargs): - # type: (Dict[str, Any], Any) -> UserProxy - """Insert or update the specified user. - - If the user already exists in the container, it is replaced. If the user - does not already exist, it is inserted. - - :param body: A dict-like object representing the user to update or insert. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A `UserProxy` instance representing the upserted user. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the given user could not be upserted. - :rtype: ~azure.cosmos.UserProxy - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - - user = self.client_connection.UpsertUser( - database_link=self.database_link, user=body, options=request_options, **kwargs - ) - - if response_hook: - response_hook(self.client_connection.last_response_headers, user) - - return UserProxy( - client_connection=self.client_connection, id=user["id"], database_link=self.database_link, properties=user - ) - - @distributed_trace - def replace_user( - self, - user, # type: Union[str, UserProxy, Dict[str, Any]] - body, # type: Dict[str, Any] - **kwargs # type: Any - ): - # type: (...) -> UserProxy - """Replaces the specified user if it exists in the container. - - :param user: The ID (name), dict representing the properties or :class:`UserProxy` - instance of the user to be replaced. - :param body: A dict-like object representing the user to replace. - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A `UserProxy` instance representing the user after replace went through. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: - If the replace failed or the user with given ID does not exist. - :rtype: ~azure.cosmos.UserProxy - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - - replaced_user = self.client_connection.ReplaceUser( - user_link=self._get_user_link(user), user=body, options=request_options, **kwargs - ) # type: Dict[str, str] - - if response_hook: - response_hook(self.client_connection.last_response_headers, replaced_user) - - return UserProxy( - client_connection=self.client_connection, - id=replaced_user["id"], - database_link=self.database_link, - properties=replaced_user - ) - - @distributed_trace - def delete_user(self, user, **kwargs): - # type: (Union[str, UserProxy, Dict[str, Any]], Any) -> None - """Delete the specified user from the container. - - :param user: The ID (name), dict representing the properties or :class:`UserProxy` - instance of the user to be deleted. - :keyword Callable response_hook: A callable invoked with the response metadata. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The user wasn't deleted successfully. - :raises ~azure.cosmos.exceptions.CosmosResourceNotFoundError: The user does not exist in the container. - :rtype: None - """ - request_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - - result = self.client_connection.DeleteUser( - user_link=self._get_user_link(user), options=request_options, **kwargs - ) - if response_hook: - response_hook(self.client_connection.last_response_headers, result) - - @distributed_trace - def read_offer(self, **kwargs): - # type: (Any) -> Offer - """Read the Offer object for this database. - - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: Offer for the database. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: - If no offer exists for the database or if the offer could not be retrieved. - :rtype: ~azure.cosmos.Offer - """ - response_hook = kwargs.pop('response_hook', None) - properties = self._get_properties() - link = properties["_self"] - query_spec = { - "query": "SELECT * FROM root r WHERE r.resource=@link", - "parameters": [{"name": "@link", "value": link}], - } - offers = list(self.client_connection.QueryOffers(query_spec, **kwargs)) - if not offers: - raise CosmosResourceNotFoundError( - status_code=StatusCodes.NOT_FOUND, - message="Could not find Offer for database " + self.database_link) - - if response_hook: - response_hook(self.client_connection.last_response_headers, offers) - - return Offer(offer_throughput=offers[0]["content"]["offerThroughput"], properties=offers[0]) - - @distributed_trace - def replace_throughput(self, throughput, **kwargs): - # type: (Optional[int], Any) -> Offer - """Replace the database-level throughput. - - :param throughput: The throughput to be set (an integer). - :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: Offer for the database, updated with new throughput. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: - If no offer exists for the database or if the offer could not be updated. - :rtype: ~azure.cosmos.Offer - """ - response_hook = kwargs.pop('response_hook', None) - properties = self._get_properties() - link = properties["_self"] - query_spec = { - "query": "SELECT * FROM root r WHERE r.resource=@link", - "parameters": [{"name": "@link", "value": link}], - } - offers = list(self.client_connection.QueryOffers(query_spec)) - if not offers: - raise CosmosResourceNotFoundError( - status_code=StatusCodes.NOT_FOUND, - message="Could not find Offer for collection " + self.database_link) - new_offer = offers[0].copy() - new_offer["content"]["offerThroughput"] = throughput - data = self.client_connection.ReplaceOffer(offer_link=offers[0]["_self"], offer=offers[0], **kwargs) - if response_hook: - response_hook(self.client_connection.last_response_headers, data) - return Offer(offer_throughput=data["content"]["offerThroughput"], properties=data) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py index 4a2e6cdcbc50..2954a3578faf 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py @@ -179,6 +179,9 @@ def __enter__(self): def __exit__(self, *args): return self.client_connection.pipeline_client.__exit__(*args) + def close(self): + self.__exit__() + @classmethod def from_connection_string(cls, conn_str, credential=None, consistency_level="Session", **kwargs): # type: (str, Optional[Any], str, Any) -> CosmosClient diff --git a/sdk/cosmos/azure-cosmos/samples/simon_testfile.py b/sdk/cosmos/azure-cosmos/samples/simon_testfile.py index c26bbc42d81c..5a4aacad43e7 100644 --- a/sdk/cosmos/azure-cosmos/samples/simon_testfile.py +++ b/sdk/cosmos/azure-cosmos/samples/simon_testfile.py @@ -6,24 +6,21 @@ from azure.core.tracing.decorator import distributed_trace import asyncio from azure.cosmos import partition_key, cosmos_client -from azure.cosmos.aio.cosmos_client_async import AsyncCosmosClient +from azure.cosmos.aio.cosmos_client import CosmosClient import azure.cosmos.exceptions as exceptions from azure.cosmos.partition_key import PartitionKey from azure.cosmos.database import DatabaseProxy -from azure.cosmos.aio.database_async import DatabaseProxy import config import heroes -def get_azure_data(): - endpoint = "https://simonmoreno-sql.documents.azure.com:443/" - key = 'd3KEBamwtPiQpuuyFSlXEOF98cuhL8oqW3jQygmAfTOPImEZPN2yYWFd4IE5pQNdBF70v8I7LldjXB6fimMbrg==' - return [endpoint, key] +endpoint = '' +key = '' def creation(): # - client = AsyncCosmosClient(get_azure_data()[0], get_azure_data()[1]) + client = CosmosClient(endpoint, key) #