diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py index 0c06891a549c..46b7d46918d5 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py @@ -175,7 +175,7 @@ def GetHeaders( # pylint: disable=too-many-statements,too-many-branches if options.get("consistencyLevel"): consistency_level = options["consistencyLevel"] headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level - elif default_client_consistency_level is not None: + elif default_client_consistency_level is not None: # Why not just check for `default_client_consistency_level` consistency_level = default_client_consistency_level headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/__init__.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py new file mode 100644 index 000000000000..986cb130b9ae --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py @@ -0,0 +1,186 @@ +# 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. + +"""Asynchronous request in the Azure Cosmos database service. +""" + +import json +import time + +from six.moves.urllib.parse import urlparse +import six +from azure.core.exceptions import DecodeError # type: ignore + +from .. import exceptions +from .. import http_constants +from .. import _retry_utility +from .._synchronized_request import _request_body_from_data + + +async def _Request(global_endpoint_manager, request_params, connection_policy, pipeline_client, request, **kwargs): + """Makes one http request using the requests module. + + :param _GlobalEndpointManager global_endpoint_manager: + :param dict request_params: + contains the resourceType, operationType, endpointOverride, + useWriteEndpoint, useAlternateWriteEndpoint information + :param documents.ConnectionPolicy connection_policy: + :param azure.core.PipelineClient pipeline_client: + Pipeline client to process the request + :param azure.core.HttpRequest request: + The request object to send through the pipeline + :return: tuple of (result, headers) + :rtype: tuple of (dict, dict) + + """ + # pylint: disable=protected-access + + connection_timeout = connection_policy.RequestTimeout + connection_timeout = kwargs.pop("connection_timeout", connection_timeout / 1000.0) + + # Every request tries to perform a refresh + client_timeout = kwargs.get('timeout') + start_time = time.time() + global_endpoint_manager.refresh_endpoint_list(None, **kwargs) + if client_timeout is not None: + kwargs['timeout'] = client_timeout - (time.time() - start_time) + if kwargs['timeout'] <= 0: + raise exceptions.CosmosClientTimeoutError() + + if request_params.endpoint_override: + base_url = request_params.endpoint_override + else: + base_url = global_endpoint_manager.resolve_service_endpoint(request_params) + if base_url != pipeline_client._base_url: + request.url = request.url.replace(pipeline_client._base_url, base_url) + + parse_result = urlparse(request.url) + + # The requests library now expects header values to be strings only starting 2.11, + # and will raise an error on validation if they are not, so casting all header values to strings. + request.headers.update({header: str(value) for header, value in request.headers.items()}) + + # We are disabling the SSL verification for local emulator(localhost/127.0.0.1) or if the user + # has explicitly specified to disable SSL verification. + is_ssl_enabled = ( + parse_result.hostname != "localhost" + and parse_result.hostname != "127.0.0.1" + and not connection_policy.DisableSSLVerification + ) + + if connection_policy.SSLConfiguration or "connection_cert" in kwargs: + ca_certs = connection_policy.SSLConfiguration.SSLCaCerts + cert_files = (connection_policy.SSLConfiguration.SSLCertFile, connection_policy.SSLConfiguration.SSLKeyFile) + response = await _PipelineRunFunction( + pipeline_client, + request, + connection_timeout=connection_timeout, + connection_verify=kwargs.pop("connection_verify", ca_certs), + connection_cert=kwargs.pop("connection_cert", cert_files), + **kwargs + ) + else: + response = await _PipelineRunFunction( + pipeline_client, + request, + connection_timeout=connection_timeout, + # If SSL is disabled, verify = false + connection_verify=kwargs.pop("connection_verify", is_ssl_enabled), + **kwargs + ) + + response = response.http_response + headers = dict(response.headers) + + data = response.body() + if data and not six.PY2: + # python 3 compatible: convert data from byte to unicode string + data = data.decode("utf-8") + + if response.status_code == 404: + raise exceptions.CosmosResourceNotFoundError(message=data, response=response) + if response.status_code == 409: + raise exceptions.CosmosResourceExistsError(message=data, response=response) + if response.status_code == 412: + raise exceptions.CosmosAccessConditionFailedError(message=data, response=response) + if response.status_code >= 400: + raise exceptions.CosmosHttpResponseError(message=data, response=response) + + result = None + if data: + try: + result = json.loads(data) + except Exception as e: + raise DecodeError( + message="Failed to decode JSON data: {}".format(e), + response=response, + error=e) + + return result, headers + + +async def _PipelineRunFunction(pipeline_client, request, **kwargs): + # pylint: disable=protected-access + + return await pipeline_client._pipeline.run(request, **kwargs) + +async def AsynchronousRequest( + client, + request_params, + global_endpoint_manager, + connection_policy, + pipeline_client, + request, + request_data, + **kwargs +): + """Performs one asynchronous http request according to the parameters. + + :param object client: Document client instance + :param dict request_params: + :param _GlobalEndpointManager global_endpoint_manager: + :param documents.ConnectionPolicy connection_policy: + :param azure.core.PipelineClient pipeline_client: PipelineClient to process the request. + :param str method: + :param str path: + :param (str, unicode, file-like stream object, dict, list or None) request_data: + :param dict query_params: + :param dict headers: + :return: tuple of (result, headers) + :rtype: tuple of (dict dict) + """ + request.data = _request_body_from_data(request_data) + if request.data and isinstance(request.data, six.string_types): + request.headers[http_constants.HttpHeaders.ContentLength] = len(request.data) + elif request.data is None: + 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( + client, + global_endpoint_manager, + _Request, + request_params, + connection_policy, + pipeline_client, + request, + **kwargs + ) 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 new file mode 100644 index 000000000000..c23ac5bde52e --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -0,0 +1,2667 @@ +# 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. + +# disable (too-many-lines) check +# pylint: disable=C0302 + +"""Document client class for the Azure Cosmos database service. +""" +# https://github.com/PyCQA/pylint/issues/3112 +# Currently pylint is locked to 2.3.3 and this is fixed in 2.4.4 +from typing import Dict, Any, Optional # pylint: disable=unused-import +import six +import asyncio +from urllib3.util.retry import Retry +from azure.core.async_paging import AsyncItemPaged +from azure.core import AsyncPipelineClient +from azure.core import PipelineClient +from azure.core.exceptions import raise_with_traceback # type: ignore +from azure.core.pipeline.policies import ( + AsyncHTTPPolicy, + ContentDecodePolicy, + HeadersPolicy, + UserAgentPolicy, + NetworkTraceLoggingPolicy, + CustomHookPolicy, + DistributedTracingPolicy, + HttpLoggingPolicy, + ProxyPolicy) + +from .. import _base as base +from .. import documents +from ..documents import ConnectionPolicy +from .. import _constants as constants +from .. import http_constants +from .. import _query_iterable as query_iterable +from .. import _runtime_constants as runtime_constants +from .. import _request_object +from .. import _synchronized_request as synchronized_request +from . import _asynchronous_request as asynchronous_request +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 .. 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 + """Represents a document client. + + Provides a client-side logical representation of the Azure Cosmos + service. This client is used to configure and execute requests against the + service. + + The service client encapsulates the endpoint and credentials used to access + the Azure Cosmos service. + """ + + class _QueryCompatibilityMode: + Default = 0 + Query = 1 + SqlQuery = 2 + + # default number precisions + _DefaultNumberHashPrecision = 3 + _DefaultNumberRangePrecision = -1 + + # default string precision + _DefaultStringHashPrecision = 3 + _DefaultStringRangePrecision = -1 + + def __init__( + self, + url_connection, # type: str + auth, # type: Dict[str, Any] + connection_policy=None, # type: Optional[ConnectionPolicy] + consistency_level=documents.ConsistencyLevel.Session, # type: str + **kwargs # type: Any + ): + # type: (...) -> None + """ + :param str url_connection: + The URL for connecting to the DB server. + :param dict auth: + Contains 'masterKey' or 'resourceTokens', where + auth['masterKey'] is the default authorization key to use to + create the client, and auth['resourceTokens'] is the alternative + authorization key. + :param documents.ConnectionPolicy connection_policy: + The connection policy for the client. + :param documents.ConsistencyLevel consistency_level: + 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 + self.resource_tokens = None + if auth is not None: + self.master_key = auth.get("masterKey") + self.resource_tokens = auth.get("resourceTokens") + + if auth.get("permissionFeed"): + self.resource_tokens = {} + for permission_feed in auth["permissionFeed"]: + resource_parts = permission_feed["resource"].split("/") + id_ = resource_parts[-1] + self.resource_tokens[id_] = permission_feed["_token"] + + self.connection_policy = connection_policy or ConnectionPolicy() + + self.partition_resolvers = {} # type: Dict[str, Any] + + self.partition_key_definition_cache = {} # type: Dict[str, Any] + + self.default_headers = { + http_constants.HttpHeaders.CacheControl: "no-cache", + http_constants.HttpHeaders.Version: http_constants.Versions.CurrentVersion, + # For single partition query with aggregate functions we would try to accumulate the results on the SDK. + # We need to set continuation as not expected. + http_constants.HttpHeaders.IsContinuationExpected: False, + } + + if consistency_level is not None: + self.default_headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level + + # Keeps the latest response headers from server. + self.last_response_headers = None + + if consistency_level == documents.ConsistencyLevel.Session: + # create a session - this is maintained only if the default consistency level + # on the client is set to session, or if the user explicitly sets it as a property + # via setter + self.session = _session.Session(self.url_connection) + else: + self.session = None # type: ignore + + self._useMultipleWriteLocations = False + self._global_endpoint_manager = global_endpoint_manager._GlobalEndpointManager(self) + + retry_policy = None + if isinstance(self.connection_policy.ConnectionRetryConfiguration, AsyncHTTPPolicy): + retry_policy = self.connection_policy.ConnectionRetryConfiguration + elif isinstance(self.connection_policy.ConnectionRetryConfiguration, int): + retry_policy = ConnectionRetryPolicy(total=self.connection_policy.ConnectionRetryConfiguration) + elif isinstance(self.connection_policy.ConnectionRetryConfiguration, Retry): + # Convert a urllib3 retry policy to a Pipeline policy + retry_policy = ConnectionRetryPolicy( + retry_total=self.connection_policy.ConnectionRetryConfiguration.total, + retry_connect=self.connection_policy.ConnectionRetryConfiguration.connect, + retry_read=self.connection_policy.ConnectionRetryConfiguration.read, + retry_status=self.connection_policy.ConnectionRetryConfiguration.status, + retry_backoff_max=self.connection_policy.ConnectionRetryConfiguration.BACKOFF_MAX, + retry_on_status_codes=list(self.connection_policy.ConnectionRetryConfiguration.status_forcelist), + retry_backoff_factor=self.connection_policy.ConnectionRetryConfiguration.backoff_factor + ) + else: + TypeError("Unsupported retry policy. Must be an azure.cosmos.ConnectionRetryPolicy, int, or urllib3.Retry") + + proxies = kwargs.pop('proxies', {}) + if self.connection_policy.ProxyConfiguration and self.connection_policy.ProxyConfiguration.Host: + host = self.connection_policy.ProxyConfiguration.Host + url = six.moves.urllib.parse.urlparse(host) + proxy = host if url.port else host + ":" + str(self.connection_policy.ProxyConfiguration.Port) + proxies.update({url.scheme : proxy}) + + policies = [ + HeadersPolicy(**kwargs), + ProxyPolicy(proxies=proxies), + UserAgentPolicy(base_user_agent=_utils.get_user_agent(), **kwargs), + ContentDecodePolicy(), + retry_policy, + CustomHookPolicy(**kwargs), + NetworkTraceLoggingPolicy(**kwargs), + DistributedTracingPolicy(**kwargs), + 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) + + # Query compatibility mode. + # Allows to specify compatibility mode used by client when making query requests. Should be removed when + # application/sql is no longer supported. + self._query_compatibility_mode = CosmosClientConnection._QueryCompatibilityMode.Default + + # 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 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 new file mode 100644 index 000000000000..2248619fb335 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_endpoint_manager_async.py @@ -0,0 +1,174 @@ +# 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 class for global endpoint manager implementation in the Azure Cosmos +database service. +""" + +import threading + +from six.moves.urllib.parse import urlparse + +from .. import _constants as constants +from .. import exceptions +from .._location_cache import LocationCache + +# pylint: disable=protected-access + + +class _GlobalEndpointManager(object): + """ + This internal class implements the logic for endpoint management for + geo-replicated database accounts. + """ + + def __init__(self, client): + self.Client = client + self.EnableEndpointDiscovery = client.connection_policy.EnableEndpointDiscovery + self.PreferredLocations = client.connection_policy.PreferredLocations + self.DefaultEndpoint = client.url_connection + self.refresh_time_interval_in_ms = self.get_refresh_time_interval_in_ms_stub() + self.location_cache = LocationCache( + self.PreferredLocations, + self.DefaultEndpoint, + self.EnableEndpointDiscovery, + client.connection_policy.UseMultipleWriteLocations, + self.refresh_time_interval_in_ms, + ) + self.refresh_needed = False + self.refresh_lock = threading.RLock() + self.last_refresh_time = 0 + + def get_refresh_time_interval_in_ms_stub(self): # pylint: disable=no-self-use + return constants._Constants.DefaultUnavailableLocationExpirationTime + + def get_write_endpoint(self): + return self.location_cache.get_write_endpoint() + + def get_read_endpoint(self): + return self.location_cache.get_read_endpoint() + + def resolve_service_endpoint(self, request): + return self.location_cache.resolve_service_endpoint(request) + + def mark_endpoint_unavailable_for_read(self, endpoint): + self.location_cache.mark_endpoint_unavailable_for_read(endpoint) + + def mark_endpoint_unavailable_for_write(self, endpoint): + self.location_cache.mark_endpoint_unavailable_for_write(endpoint) + + def get_ordered_write_endpoints(self): + return self.location_cache.get_ordered_write_endpoints() + + def get_ordered_read_endpoints(self): + return self.location_cache.get_ordered_read_endpoints() + + 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): + self.refresh_needed = True + self.refresh_endpoint_list(database_account) + + def refresh_endpoint_list(self, database_account, **kwargs): + 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) + except Exception as e: + raise e + + 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 + + if ( + self.location_cache.should_refresh_endpoints() + 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) + 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): + """Gets the database account. + + First tries by using the default endpoint, and if that doesn't work, + use the endpoints for the preferred locations in the order they are + specified, to get the database account. + """ + try: + database_account = 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 + # information from any of the preferred locations that the user might have + # specified (by creating a locational endpoint) and keeping eating the exception + # until we get the database account and return None at the end, if we are not able + # to get that info from any endpoints + except exceptions.CosmosHttpResponseError: + for location_name in self.PreferredLocations: + locational_endpoint = _GlobalEndpointManager.GetLocationalEndpoint(self.DefaultEndpoint, location_name) + try: + database_account = self._GetDatabaseAccountStub(locational_endpoint, **kwargs) + return database_account + except exceptions.CosmosHttpResponseError: + pass + + return None + + def _GetDatabaseAccountStub(self, endpoint, **kwargs): + """Stub for getting database account from the client. + + This can be used for mocking purposes as well. + """ + return self.Client.GetDatabaseAccount(endpoint, **kwargs) + + @staticmethod + def GetLocationalEndpoint(default_endpoint, location_name): + # For default_endpoint like 'https://contoso.documents.azure.com:443/' parse it to + # generate URL format. This default_endpoint should be global endpoint(and cannot + # be a locational endpoint) and we agreed to document that + endpoint_url = urlparse(default_endpoint) + + # hostname attribute in endpoint_url will return 'contoso.documents.azure.com' + if endpoint_url.hostname is not None: + hostname_parts = str(endpoint_url.hostname).lower().split(".") + if hostname_parts is not None: + # global_database_account_name will return 'contoso' + global_database_account_name = hostname_parts[0] + + # Prepare the locational_database_account_name as contoso-EastUS for location_name 'East US' + locational_database_account_name = global_database_account_name + "-" + location_name.replace(" ", "") + + # Replace 'contoso' with 'contoso-EastUS' and return locational_endpoint + # as https://contoso-EastUS.documents.azure.com:443/ + locational_endpoint = default_endpoint.lower().replace( + global_database_account_name, locational_database_account_name, 1 + ) + return locational_endpoint + + return None diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/container_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/container_async.py new file mode 100644 index 000000000000..6120d533c918 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/container_async.py @@ -0,0 +1,802 @@ +# 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_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/cosmos_client_async.py new file mode 100644 index 000000000000..879cafcf7f9f --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/cosmos_client_async.py @@ -0,0 +1,456 @@ +# 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 new file mode 100644 index 000000000000..cbb1e0ab6902 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/database_async.py @@ -0,0 +1,768 @@ +# 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/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 7286b2b8d0b4..1e8b9f51d483 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -25,6 +25,8 @@ 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 import CosmosClientConnection @@ -496,6 +498,7 @@ def create_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) @@ -514,6 +517,61 @@ def create_item( ) if response_hook: response_hook(self.client_connection.last_response_headers, result) + print(f"Create item took {(time.time() - start) * 1000} ms") + 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 diff --git a/sdk/cosmos/azure-cosmos/samples/heroes.py b/sdk/cosmos/azure-cosmos/samples/heroes.py new file mode 100644 index 000000000000..65b4f5ece427 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/samples/heroes.py @@ -0,0 +1,97 @@ +import uuid + +def get_generic_hero(): + hero_item = { + 'id': 'Generic_Hero_' + str(uuid.uuid4()), + 'lastName': 'Smith', + 'parents': None, + 'children': None, + 'address': { + 'state': 'FL', + 'city': 'Miami' + }, + 'saved': ['block'], + 'professional': False, + 'company': None + } + return hero_item + +def get_batman(): + hero_item = { + 'id': 'Batman', + 'lastName': 'Wayne', + 'parents': None, + 'children': None, + 'address': { + 'state': 'WA', + 'city': 'Gotham' + }, + 'saved': ['state', 'city'], + 'professional': True, + 'company': 'DC' + } + return hero_item + +def get_flash(): + hero_item = { + 'id': 'Flash', + 'lastName': 'Allen', + 'parents': None, + 'children': None, + 'address': { + 'state': 'NY', + 'city': 'New York' + }, + 'saved': ['world','country'], + 'professional': True, + 'company': 'DC' + } + return hero_item + +def get_superman(): + hero_item = { + 'id': 'Superman', + 'lastName': 'Kent', + 'parents': None, + 'children': None, + 'address': { + 'state': 'WA', + 'city': 'Metropolis' + }, + 'saved': ['universe','world','country', 'state'], + 'professional': True, + 'company': 'DC' + } + return hero_item + +def get_spider(): + hero_item = { + 'id': 'Spiderman', + 'lastName': 'Parker', + 'parents': None, + 'children': None, + 'address': { + 'state': 'NY', + 'city': 'New York' + }, + 'saved': ['galaxy','world','country'], + 'professional': True, + 'company': 'Marvel' + } + return hero_item + +def get_iron(): + hero_item = { + 'id': 'Ironman', + 'lastName': 'Stark', + 'parents': None, + 'children': None, + 'address': { + 'state': 'NY', + 'city': 'New York' + }, + 'saved': ['galaxy','world','country'], + 'professional': True, + 'company': 'Marvel' + } + return hero_item \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/samples/simon_testfile.py b/sdk/cosmos/azure-cosmos/samples/simon_testfile.py new file mode 100644 index 000000000000..c26bbc42d81c --- /dev/null +++ b/sdk/cosmos/azure-cosmos/samples/simon_testfile.py @@ -0,0 +1,169 @@ +import sys +sys.path.append(r"C:\Users\simonmoreno\Repos\azure-sdk-for-python\sdk\cosmos\azure-cosmos") + + +from azure.cosmos import container +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 +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] + +def creation(): + + # + client = AsyncCosmosClient(get_azure_data()[0], get_azure_data()[1]) + # + database_name = 'MockHeroesDatabase' + database = client.create_database_if_not_exists(id=database_name) + # + + container_name = 'mockHeroesContainer' + container = database.create_container_if_not_exists( + id=container_name, + partition_key=PartitionKey(path="/lastName"), + offer_throughput=400 + ) + + real_heroes = [heroes.get_superman(), heroes.get_batman(), heroes.get_flash(), heroes.get_spider(), heroes.get_iron()] + generics = [heroes.get_generic_hero(), heroes.get_generic_hero(), heroes.get_generic_hero()] + + for hero in real_heroes: + container.create_item(body=hero) + + for generic in generics: + container.create_item(body=generic) + + for hero in real_heroes: + response = container.read_item(item=hero['id'], partition_key=hero['lastName']) + request_charge = container.client_connection.last_response_headers['x-ms-request-charge'] #! + if hero['id'] == 'Superman': print(container.client_connection.last_response_headers) + print('Read item with id {0}. Operation consumed {1} request units'.format(response['id'], (request_charge))) + + query = "SELECT * FROM c WHERE c.lastName IN ('Kent', 'Parker')" + + items = list(container.query_items( + query=query, + enable_cross_partition_query=True #! + )) + + request_charge = container.client_connection.last_response_headers['x-ms-request-charge'] #! + print('Query returned {0} items. Operation consumed {1} request units'.format(len(items), request_charge)) + +def clean_heroes(): + client = AsyncCosmosClient(get_azure_data()[0], get_azure_data()[1]) + database_name = 'MockHeroesDatabase' + database = client.get_database_client(database_name) + container_name = 'mockHeroesContainer' + container = database.get_container_client(container_name) + real_heroes = [heroes.get_superman(), heroes.get_batman(), heroes.get_flash(), heroes.get_spider(), heroes.get_iron()] + for h in real_heroes: + response = container.delete_item(h['id'], partition_key=h['lastName']) + print(response) + +def destroy(): + client = AsyncCosmosClient(get_azure_data()[0], get_azure_data()[1]) + database_name = 'MockHeroesDatabase' + response = client.delete_database(database_name) + print(f"Database with name {database_name} has been deleted.") + print(response) + +async def createaio(): + client = AsyncCosmosClient(get_azure_data()[0], get_azure_data()[1]) + database_name = 'MockHeroesDatabase' + database = client.create_database_if_not_exists(id=database_name) + + container_name = 'mockHeroesContainer' + container = database.create_container_if_not_exists( + id=container_name, + partition_key=PartitionKey(path="/lastName"), + offer_throughput=400 + ) + + real_heroes = [heroes.get_superman(), heroes.get_batman(), heroes.get_flash(), heroes.get_spider(), heroes.get_iron()] + generics = [heroes.get_generic_hero(), heroes.get_generic_hero(), heroes.get_generic_hero()] + + for hero in real_heroes: + await container.create_item_aio(body=hero) + + # for generic in generics: + # container.create_item_aio(body=generic) + +def get_db(): + client = AsyncCosmosClient(get_azure_data()[0], get_azure_data()[1]) + database_name = 'MockHeroesDatabase' + res = client.get_database_client(database_name).list_containers() + r= client.get_database_client('lols').list_containers() + x = client.get_database_client.three + + print(res) + print("OOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOO") + print(r) + print("PPPPPPPPPPPPPPPPPPPPPPPPPPPPPPPP") + + # for i in res: + # print(i) + for i in res: + print("CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC") + print(i) + +# asyncio.run(createaio()) + +# creation() + +import uuid + + +endpoint = "https://simonmoreno-sql.documents.azure.com:443/" +key = 'd3KEBamwtPiQpuuyFSlXEOF98cuhL8oqW3jQygmAfTOPImEZPN2yYWFd4IE5pQNdBF70v8I7LldjXB6fimMbrg==' + +def get_test_item(): + async_item = { + 'id': 'Async_' + str(uuid.uuid4()), + 'address': { + 'state': 'WA', + 'city': 'Redmond', + 'street': '1 Microsoft Way' + }, + 'test_object': True + } + return async_item + +def create_test(): + client = cosmos_client.CosmosClient(endpoint, key) + db = client.create_database(id="AsyncDB") + container = db.create_container( + id="AsyncContainer", + partition_key=PartitionKey(path="/id")) + ids = [] + for i in range(20): + body = get_test_item() + print(body.get("id")) + ids.append(body.get("id")) + container.create_item(body=body) + return ids + +async def async_read_test(): + # ids = create_test() + client = AsyncCosmosClient(endpoint, key) + db = client.get_database_client(id="AsyncDB") + container = db.get_container_client(id="AsyncContainer") + print(container.read()) + + + +asyncio.run(async_read_test()) \ No newline at end of file