From 5ae800e3f24a3213dd79d87ff3a97a23a0ed1f41 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Fri, 18 Mar 2022 17:53:01 -0400 Subject: [PATCH 01/14] working authentication to get database account --- sdk/cosmos/azure-cosmos/azure/cosmos/_base.py | 23 +++++++++++-------- .../azure/cosmos/_cosmos_client_connection.py | 19 ++++++++++++++- sdk/cosmos/azure-cosmos/azure/cosmos/auth.py | 14 +++++++++-- .../azure/cosmos/cosmos_client.py | 7 ++++-- 4 files changed, 49 insertions(+), 14 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py index 9a42583513ed..e70b8e0e38ff 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py @@ -71,6 +71,7 @@ 'query_version': 'queryVersion' } + def _get_match_headers(kwargs): # type: (Dict[str, Any]) -> Tuple(Optional[str], Optional[str]) if_match = kwargs.pop('if_match', None) @@ -112,14 +113,14 @@ def build_options(kwargs): def GetHeaders( # pylint: disable=too-many-statements,too-many-branches - cosmos_client_connection, - default_headers, - verb, - path, - resource_id, - resource_type, - options, - partition_key_range_id=None, + cosmos_client_connection, + default_headers, + verb, + path, + resource_id, + resource_type, + options, + partition_key_range_id=None, ): """Gets HTTP request headers. @@ -638,7 +639,7 @@ def ParsePaths(paths): newIndex += 1 # This will extract the token excluding the quote chars - token = path[currentIndex + 1 : newIndex] + token = path[currentIndex + 1: newIndex] tokens.append(token) currentIndex = newIndex + 1 else: @@ -657,3 +658,7 @@ def ParsePaths(paths): tokens.append(token) return tokens + + +def create_scope_from_url(url): + return url.replace(":443", "") + ".default" diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index a8da94b53818..355f3cac2fd2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -41,7 +41,8 @@ CustomHookPolicy, DistributedTracingPolicy, HttpLoggingPolicy, - ProxyPolicy) + ProxyPolicy, + BearerTokenCredentialPolicy) from . import _base as base from . import documents @@ -116,9 +117,11 @@ def __init__( self.master_key = None self.resource_tokens = None + self.aad_credentials = None if auth is not None: self.master_key = auth.get("masterKey") self.resource_tokens = auth.get("resourceTokens") + self.aad_credentials = auth.get("clientSecretCredential") if auth.get("permissionFeed"): self.resource_tokens = {} @@ -187,6 +190,10 @@ def __init__( HttpLoggingPolicy(**kwargs), ] + # if self.aad_credentials: + # scopes = base.create_scope_from_url(self.url_connection) + # policies.append(BearerTokenCredentialPolicy(self.aad_credentials, scopes)) + transport = kwargs.pop("transport", None) self.pipeline_client = PipelineClient(base_url=url_connection, transport=transport, policies=policies) @@ -217,6 +224,9 @@ def _set_client_consistency_level( :type consistency_level: Optional[str] :rtype: None """ + if database_account is None: + print("DATABASE ACCOUNT IS NONE") + consistency_level = documents.ConsistencyLevel.Session if consistency_level is None: # Set to default level present in account user_consistency_policy = database_account.ConsistencyPolicy @@ -2042,6 +2052,13 @@ def GetDatabaseAccount(self, url_connection=None, **kwargs): initial_headers = dict(self.default_headers) headers = base.GetHeaders(self, initial_headers, "get", "", "", "", {}) # path # id # type + if self.aad_credentials: + auth_prefix = "type=aad&ver=1.0&sig=" + scope = self.url_connection.replace(":443", "") + ".default" + aad_token = self.aad_credentials.get_token(scope) # Returns an AccessToken object + auth_token = auth_prefix + aad_token.token + headers[http_constants.HttpHeaders.Authorization] = auth_token + request_params = _request_object.RequestObject("databaseaccount", documents._OperationType.Read, url_connection) result, self.last_response_headers = self.__Get("", request_params, headers, **kwargs) database_account = documents.DatabaseAccount() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py b/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py index 9c8252bd9869..7b0bc1ada195 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py @@ -31,7 +31,7 @@ def GetAuthorizationHeader( - cosmos_client_connection, verb, path, resource_id_or_fullname, is_name_based, resource_type, headers + cosmos_client_connection, verb, path, resource_id_or_fullname, is_name_based, resource_type, headers ): """Gets the authorization header. @@ -58,6 +58,8 @@ def GetAuthorizationHeader( return __GetAuthorizationTokenUsingResourceTokens( cosmos_client_connection.resource_tokens, path, resource_id_or_fullname ) + # if cosmos_client_connection.aad_credentials: + # return __get_aad_token(cosmos_client_connection.aad_credentials, cosmos_client_connection.url_connection) return None @@ -138,9 +140,17 @@ def __GetAuthorizationTokenUsingResourceTokens(resource_tokens, path, resource_i # Get the last resource id or resource name from the path and get it's token from resource_tokens for i in range(len(path_parts), 1, -1): - segment = path_parts[i-1] + segment = path_parts[i - 1] sub_path = "/".join(path_parts[:i]) if not segment in resource_types and sub_path in resource_tokens: return resource_tokens[sub_path] return None + + +# def __get_aad_token(aad_credentials, url): +# auth_prefix = "type=aad&ver=1.0&sig=" +# scope = url.replace(":443", "") + ".default" +# aad_token = aad_credentials.get_token(scope) # Returns an AccessToken object +# auth_token = auth_prefix + aad_token.token +# return auth_token diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py index c581dcdf6baa..26491f2cb52a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py @@ -25,6 +25,7 @@ from typing import Any, Dict, Optional, Union, cast, Iterable, List # pylint: disable=unused-import from azure.core.tracing.decorator import distributed_trace # type: ignore +from azure.identity import ClientSecretCredential from ._cosmos_client_connection import CosmosClientConnection from ._base import build_options @@ -60,10 +61,12 @@ def _build_auth(credential): auth['resourceTokens'] = credential # type: ignore elif hasattr(credential, '__iter__'): auth['permissionFeed'] = credential + elif hasattr(credential, 'get_token'): + auth['clientSecretCredential'] = 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.") + "Unrecognized credential type. Please supply the master key as a string " + "or a dictionary, or resource tokens, or a list of permissions, or a ClientSecretCredential.") return auth From bc0ed5ae3d78c722e22eee0720fe6b014e0d3707 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 22 Mar 2022 00:25:22 -0400 Subject: [PATCH 02/14] working aad authentication for sync client with sample --- .../azure/cosmos/_cosmos_client_connection.py | 20 +- sdk/cosmos/azure-cosmos/azure/cosmos/auth.py | 179 ++++++++++++++++-- .../azure/cosmos/http_constants.py | 1 + .../samples/access_cosmos_with_aad.py | 96 ++++++++++ sdk/cosmos/azure-cosmos/samples/config.py | 3 + 5 files changed, 272 insertions(+), 27 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index 355f3cac2fd2..487d2e1dad6a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -41,8 +41,7 @@ CustomHookPolicy, DistributedTracingPolicy, HttpLoggingPolicy, - ProxyPolicy, - BearerTokenCredentialPolicy) + ProxyPolicy) from . import _base as base from . import documents @@ -59,6 +58,7 @@ from . import _session from . import _utils from .partition_key import _Undefined, _Empty +from .auth import CosmosBearerTokenCredentialPolicy ClassType = TypeVar("ClassType") @@ -190,9 +190,9 @@ def __init__( HttpLoggingPolicy(**kwargs), ] - # if self.aad_credentials: - # scopes = base.create_scope_from_url(self.url_connection) - # policies.append(BearerTokenCredentialPolicy(self.aad_credentials, scopes)) + if self.aad_credentials: + scopes = base.create_scope_from_url(self.url_connection) + policies.append(CosmosBearerTokenCredentialPolicy(self.aad_credentials, scopes)) transport = kwargs.pop("transport", None) self.pipeline_client = PipelineClient(base_url=url_connection, transport=transport, policies=policies) @@ -224,9 +224,6 @@ def _set_client_consistency_level( :type consistency_level: Optional[str] :rtype: None """ - if database_account is None: - print("DATABASE ACCOUNT IS NONE") - consistency_level = documents.ConsistencyLevel.Session if consistency_level is None: # Set to default level present in account user_consistency_policy = database_account.ConsistencyPolicy @@ -2052,13 +2049,6 @@ def GetDatabaseAccount(self, url_connection=None, **kwargs): initial_headers = dict(self.default_headers) headers = base.GetHeaders(self, initial_headers, "get", "", "", "", {}) # path # id # type - if self.aad_credentials: - auth_prefix = "type=aad&ver=1.0&sig=" - scope = self.url_connection.replace(":443", "") + ".default" - aad_token = self.aad_credentials.get_token(scope) # Returns an AccessToken object - auth_token = auth_prefix + aad_token.token - headers[http_constants.HttpHeaders.Authorization] = auth_token - request_params = _request_object.RequestObject("databaseaccount", documents._OperationType.Read, url_connection) result, self.last_response_headers = self.__Get("", request_params, headers, **kwargs) database_account = documents.DatabaseAccount() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py b/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py index 7b0bc1ada195..21b0647a3f0e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py @@ -28,7 +28,15 @@ import urllib.parse from . import http_constants +from typing import Any, Dict, Optional, ClassVar +import time +from azure.core.pipeline.policies._base import HTTPPolicy +from azure.core.exceptions import ServiceRequestError +from azure.core.credentials import AccessToken +from azure.core.pipeline import PipelineRequest, PipelineResponse + +TokenCredential = ClassVar def GetAuthorizationHeader( cosmos_client_connection, verb, path, resource_id_or_fullname, is_name_based, resource_type, headers @@ -51,20 +59,18 @@ def GetAuthorizationHeader( resource_id_or_fullname = resource_id_or_fullname.lower() if cosmos_client_connection.master_key: - return __GetAuthorizationTokenUsingMasterKey( + return __get_authorization_token_using_master_key( verb, resource_id_or_fullname, resource_type, headers, cosmos_client_connection.master_key ) if cosmos_client_connection.resource_tokens: - return __GetAuthorizationTokenUsingResourceTokens( + return __get_authorization_token_using_resource_token( cosmos_client_connection.resource_tokens, path, resource_id_or_fullname ) - # if cosmos_client_connection.aad_credentials: - # return __get_aad_token(cosmos_client_connection.aad_credentials, cosmos_client_connection.url_connection) return None -def __GetAuthorizationTokenUsingMasterKey(verb, resource_id_or_fullname, resource_type, headers, master_key): +def __get_authorization_token_using_master_key(verb, resource_id_or_fullname, resource_type, headers, master_key): """Gets the authorization token using `master_key. :param str verb: @@ -99,7 +105,7 @@ def __GetAuthorizationTokenUsingMasterKey(verb, resource_id_or_fullname, resourc return "type={type}&ver={ver}&sig={sig}".format(type=master_token, ver=token_version, sig=signature[:-1]) -def __GetAuthorizationTokenUsingResourceTokens(resource_tokens, path, resource_id_or_fullname): +def __get_authorization_token_using_resource_token(resource_tokens, path, resource_id_or_fullname): """Get the authorization token using `resource_tokens`. :param dict resource_tokens: @@ -148,9 +154,158 @@ def __GetAuthorizationTokenUsingResourceTokens(resource_tokens, path, resource_i return None -# def __get_aad_token(aad_credentials, url): -# auth_prefix = "type=aad&ver=1.0&sig=" -# scope = url.replace(":443", "") + ".default" -# aad_token = aad_credentials.get_token(scope) # Returns an AccessToken object -# auth_token = auth_prefix + aad_token.token -# return auth_token +# pylint:disable=too-few-public-methods +class _CosmosBearerTokenCredentialPolicyBase(object): + """Base class for a Cosmos Bearer Token Credential Policy. + Based on ~azure.core.pipeline.policies.BearerTokenCredentialPolicy. + Takes care of updating token as needed and updating the request headers. + + :param credential: The credential. + :type credential: ~azure.core.credentials.TokenCredential + :param str scopes: Lets you specify the type of access needed. + """ + + def __init__(self, credential, *scopes, **kwargs): # pylint:disable=unused-argument + # type: (TokenCredential, *str, **Any) -> None + super(_CosmosBearerTokenCredentialPolicyBase, self).__init__() + self._scopes = scopes + self._credential = credential + self._token = None # type: Optional[AccessToken] + + @staticmethod + def _enforce_https(request): + # type: (PipelineRequest) -> None + + # move 'enforce_https' from options to context so it persists + # across retries but isn't passed to a transport implementation + option = request.context.options.pop("enforce_https", None) + + # True is the default setting; we needn't preserve an explicit opt in to the default behavior + if option is False: + request.context["enforce_https"] = option + + enforce_https = request.context.get("enforce_https", True) + if enforce_https and not request.http_request.url.lower().startswith("https"): + raise ServiceRequestError( + "Bearer token authentication is not permitted for non-TLS protected (non-https) URLs." + ) + + @staticmethod + def _update_headers(headers, token): + # type: (Dict[str, str], str) -> None + """Updates the Authorization header with the Cosmos AAD bearer token. + + :param dict headers: The HTTP Request headers + :param str token: The OAuth token. + """ + headers[http_constants.HttpHeaders.Authorization] = "type=aad&ver=1.0&sig={}".format(token) + + @property + def _need_new_token(self): + # type: () -> bool + return not self._token or self._token.expires_on - time.time() < 300 + + +class CosmosBearerTokenCredentialPolicy(_CosmosBearerTokenCredentialPolicyBase, HTTPPolicy): + """Adds a Cosmos bearer token Authorization header to requests. + + :param credential: The credential. + :type credential: ~azure.core.TokenCredential + :param str scopes: Lets you specify the type of access needed. + :raises: :class:`~azure.core.exceptions.ServiceRequestError` + """ + + def on_request(self, request): + # type: (PipelineRequest) -> None + """Called before the policy sends a request. + + The base implementation authorizes the request with a bearer token. + + :param ~azure.core.pipeline.PipelineRequest request: the request + """ + self._enforce_https(request) + + if self._token is None or self._need_new_token: + self._token = self._credential.get_token(*self._scopes) + self._update_headers(request.http_request.headers, self._token.token) + + def authorize_request(self, request, *scopes, **kwargs): + # type: (PipelineRequest, *str, **Any) -> None + """Acquire a token from the credential and authorize the request with it. + + Keyword arguments are passed to the credential's get_token method. The token will be cached and used to + authorize future requests. + + :param ~azure.core.pipeline.PipelineRequest request: the request + :param str scopes: required scopes of authentication + """ + self._token = self._credential.get_token(*scopes, **kwargs) + self._update_headers(request.http_request.headers, self._token.token) + + def send(self, request): + # type: (PipelineRequest) -> PipelineResponse + """Authorize request with a bearer token and send it to the next policy + + :param request: The pipeline request object + :type request: ~azure.core.pipeline.PipelineRequest + """ + self.on_request(request) + try: + response = self.next.send(request) + self.on_response(request, response) + except Exception: # pylint:disable=broad-except + handled = self.on_exception(request) + if not handled: + raise + else: + if response.http_response.status_code == 401: + self._token = None # any cached token is invalid + if "WWW-Authenticate" in response.http_response.headers: + request_authorized = self.on_challenge(request, response) + if request_authorized: + try: + response = self.next.send(request) + self.on_response(request, response) + except Exception: # pylint:disable=broad-except + handled = self.on_exception(request) + if not handled: + raise + + return response + + def on_challenge(self, request, response): + # type: (PipelineRequest, PipelineResponse) -> bool + """Authorize request according to an authentication challenge + + This method is called when the resource provider responds 401 with a WWW-Authenticate header. + + :param ~azure.core.pipeline.PipelineRequest request: the request which elicited an authentication challenge + :param ~azure.core.pipeline.PipelineResponse response: the resource provider's response + :returns: a bool indicating whether the policy should send the request + """ + # pylint:disable=unused-argument,no-self-use + return False + + def on_response(self, request, response): + # type: (PipelineRequest, PipelineResponse) -> None + """Executed after the request comes back from the next policy. + + :param request: Request to be modified after returning from the policy. + :type request: ~azure.core.pipeline.PipelineRequest + :param response: Pipeline response object + :type response: ~azure.core.pipeline.PipelineResponse + """ + + def on_exception(self, request): + # type: (PipelineRequest) -> bool + """Executed when an exception is raised while executing the next policy. + + This method is executed inside the exception handler. + + :param request: The Pipeline request object + :type request: ~azure.core.pipeline.PipelineRequest + :return: False by default, override with True to stop the exception. + :rtype: bool + """ + # pylint: disable=no-self-use,unused-argument + return False diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py b/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py index b658af7389f0..499b22ad05ae 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py @@ -375,6 +375,7 @@ class SubStatusCodes(object): REDUNDANT_COLLECTION_PUT = 1009 SHARED_THROUGHPUT_DATABASE_QUOTA_EXCEEDED = 1010 SHARED_THROUGHPUT_OFFER_GROW_NOT_NEEDED = 1011 + AAD_REQUEST_NOT_AUTHORIZED = 5300 # 404: LSN in session token is higher READ_SESSION_NOTAVAILABLE = 1002 diff --git a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py new file mode 100644 index 000000000000..bfe377571bd2 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py @@ -0,0 +1,96 @@ +import azure.cosmos.cosmos_client as cosmos_client +import azure.cosmos.exceptions as exceptions +from azure.cosmos.partition_key import PartitionKey +from azure.identity import ClientSecretCredential +import config + +# ---------------------------------------------------------------------------------------------------------- +# Prerequistes - +# +# 1. An Azure Cosmos account - +# https://docs.microsoft.com/azure/cosmos-db/create-sql-api-python#create-a-database-account +# +# 2. Microsoft Azure Cosmos +# pip install azure-cosmos>=4.3.0b4 +# ---------------------------------------------------------------------------------------------------------- +# Sample - demonstrates how to authenticate and use your database account using AAD credentials +# Read more about operations allowed for this authorization method: https://aka.ms/cosmos-native-rbac +# ---------------------------------------------------------------------------------------------------------- +# Note: +# This sample creates a Container to your database account. +# Each time a Container is created the account will be billed for 1 hour of usage based on +# the provisioned throughput (RU/s) of that account. +# ---------------------------------------------------------------------------------------------------------- +# +HOST = config.settings["host"] +MASTER_KEY = config.settings["master_key"] + +TENANT_ID = config.settings["tenant_id"] +CLIENT_ID = config.settings["client_id"] +CLIENT_SECRET = config.settings["client_secret"] + +DATABASE_ID = config.settings["database_id"] +CONTAINER_ID = config.settings["container_id"] +PARTITION_KEY = PartitionKey(path="/id") + + +def get_test_item(num): + test_item = { + 'id': 'Item_' + str(num), + 'test_object': True, + 'lastName': 'Smith' + } + return test_item + + +def create_sample_resources(): + print("creating sample resources") + client = cosmos_client.CosmosClient(HOST, MASTER_KEY) + db = client.create_database(DATABASE_ID) + db.create_container(id=CONTAINER_ID, partition_key=PARTITION_KEY) + + +def delete_sample_resources(): + print("deleting sample resources") + client = cosmos_client.CosmosClient(HOST, MASTER_KEY) + client.delete_database(DATABASE_ID) + + +def run_sample(): + # Since Azure Cosmos DB data plane SDK does not cover management operations, we have to create our resources + # with a master key authenticated client for this sample. + create_sample_resources() + + # With this done, you can use your AAD service principal id and secret to create your ClientSecretCredential. + aad_credentials = ClientSecretCredential( + tenant_id=TENANT_ID, + client_id=CLIENT_ID, + client_secret=CLIENT_SECRET) + + # Use your credentials to authenticate your client. + aad_client = cosmos_client.CosmosClient(HOST, aad_credentials) + + # Do any R/W data operations with your authorized AAD client. + db = aad_client.get_database_client(DATABASE_ID) + container = db.get_container_client(CONTAINER_ID) + + print(container.read()) + container.create_item(get_test_item(0)) + print(container.read_item(item='Item_0', partition_key='Item_0')) + print(list(container.query_items(query='select * from c', partition_key='Item_0'))) + container.delete_item(item='Item_0', partition_key='Item_0') + + # Attempting to do management operations will return a 403 Forbidden exception. + try: + aad_client.delete_database(DATABASE_ID) + except exceptions.CosmosHttpResponseError as e: + assert e.status_code == 403 + print("403 error assertion success") + + # To clean up the sample, we use a master key client again to get access to deleting containers and databases. + delete_sample_resources() + print("end of sample") + + +if __name__ == "__main__": + run_sample() diff --git a/sdk/cosmos/azure-cosmos/samples/config.py b/sdk/cosmos/azure-cosmos/samples/config.py index a85ac445a84c..a8752970bdd4 100644 --- a/sdk/cosmos/azure-cosmos/samples/config.py +++ b/sdk/cosmos/azure-cosmos/samples/config.py @@ -5,4 +5,7 @@ 'master_key': os.environ.get('ACCOUNT_KEY', '[YOUR KEY]'), 'database_id': os.environ.get('COSMOS_DATABASE', '[YOUR DATABASE]'), 'container_id': os.environ.get('COSMOS_CONTAINER', '[YOUR CONTAINER]'), + 'tenant_id': os.environ.get('TENANT_ID', '[YOUR CONTAINER]'), + 'client_id': os.environ.get('COSMOS_CONTAINER', '[YOUR CONTAINER]'), + 'client_secret': os.environ.get('COSMOS_CONTAINER', '[YOUR CONTAINER]'), } From 3fa7683a3cd0a6a0fd909b08bc2014915f2308a2 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 22 Mar 2022 00:48:57 -0400 Subject: [PATCH 03/14] readme and changelog --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 5 +++ sdk/cosmos/azure-cosmos/README.md | 32 +++++++++++++------ sdk/cosmos/azure-cosmos/azure/cosmos/_base.py | 1 + 3 files changed, 29 insertions(+), 9 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 10a8015983f9..b8f38da36e46 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -1,5 +1,10 @@ ## Release History +### 4.3.0b4 (Unreleased) + +#### Features Added +- Added support for AAD authentication for the sync client + ### 4.3.0b3 (Unreleased) #### Features Added diff --git a/sdk/cosmos/azure-cosmos/README.md b/sdk/cosmos/azure-cosmos/README.md index 50fb349f10dd..a2651b232cdf 100644 --- a/sdk/cosmos/azure-cosmos/README.md +++ b/sdk/cosmos/azure-cosmos/README.md @@ -76,6 +76,28 @@ KEY = os.environ['ACCOUNT_KEY'] client = CosmosClient(URL, credential=KEY) ``` +### AAD Authentication + +You can also authenticate a client utilizing your service principal's AAD credentials and the azure identity package: +```Python +from azure.cosmos import CosmosClient +from azure.identity import ClientSecretCredential + +import os +URL = os.environ['ACCOUNT_URI'] +TENANT_ID = os.environ['TENANT_ID'] +CLIENT_ID = os.environ['CLIENT_ID'] +CLIENT_SECRET = os.environ['CLIENT_SECRET'] + +AAD_CREDENTIALS = ClientSecretCredential( + tenant_id=TENANT_ID, + client_id=CLIENT_ID, + client_secret=CLIENT_SECRET) + +client = CosmosClient(URL, AAD_CREDENTIALS) +``` +More information on allowed operations for AAD authenticated clients: [RBAC Permission Model](https://aka.ms/cosmos-native-rbac) + ## Key concepts Once you've initialized a [CosmosClient][ref_cosmosclient], you can interact with the primary resource types in Cosmos DB: @@ -111,7 +133,7 @@ Currently the features below are **not supported**. For alternatives options, ch * Change Feed: Processor * Change Feed: Read multiple partitions key values * Change Feed: Read specific time -* Change Feed: Read from the beggining +* Change Feed: Read from the beginning * Change Feed: Pull model * Cross-partition ORDER BY for mixed types * Cross partition queries do not handle partition splits (410 Gone errors) @@ -126,10 +148,6 @@ Currently the features below are **not supported**. For alternatives options, ch * Get the connection string * Get the minimum RU/s of a container -### Security Limitations: - -* AAD support - ## Workarounds ### Bulk processing Limitation Workaround @@ -140,10 +158,6 @@ If you want to use Python SDK to perform bulk inserts to Cosmos DB, the best alt Typically you can use [Azure Portal](https://portal.azure.com/), [Azure Cosmos DB Resource Provider REST API](https://docs.microsoft.com/rest/api/cosmos-db-resource-provider), [Azure CLI](https://docs.microsoft.com/cli/azure/azure-cli-reference-for-cosmos-db) or [PowerShell](https://docs.microsoft.com/azure/cosmos-db/manage-with-powershell) for the control plane unsupported limitations. -### AAD Support Workaround - -A possible workaround is to use managed identities to [programmatically](https://docs.microsoft.com/azure/cosmos-db/managed-identity-based-authentication) get the keys. - ## Boolean Data Type While the Python language [uses](https://docs.python.org/3/library/stdtypes.html?highlight=boolean#truth-value-testing) "True" and "False" for boolean types, Cosmos DB [accepts](https://docs.microsoft.com/azure/cosmos-db/sql-query-is-bool) "true" and "false" only. In other words, the Python language uses Boolean values with the first uppercase letter and all other lowercase letters, while Cosmos DB and its SQL language use only lowercase letters for those same Boolean values. How to deal with this challenge? diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py index e70b8e0e38ff..aa2e881bd34e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py @@ -661,4 +661,5 @@ def ParsePaths(paths): def create_scope_from_url(url): + # Need a better way to do this, regex? return url.replace(":443", "") + ".default" From 37a27e00978841710d445bfada70452484798d2c Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 22 Mar 2022 12:18:56 -0400 Subject: [PATCH 04/14] pylint and better comments on sample --- sdk/cosmos/azure-cosmos/azure/cosmos/auth.py | 9 +++------ sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py | 1 - .../azure-cosmos/samples/access_cosmos_with_aad.py | 10 ++++++---- 3 files changed, 9 insertions(+), 11 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py b/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py index 21b0647a3f0e..a218a700023f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py @@ -26,15 +26,12 @@ from hashlib import sha256 import hmac import urllib.parse - -from . import http_constants -from typing import Any, Dict, Optional, ClassVar import time - -from azure.core.pipeline.policies._base import HTTPPolicy +from typing import Any, Dict, Optional, ClassVar from azure.core.exceptions import ServiceRequestError from azure.core.credentials import AccessToken from azure.core.pipeline import PipelineRequest, PipelineResponse +from . import http_constants TokenCredential = ClassVar @@ -206,7 +203,7 @@ def _need_new_token(self): return not self._token or self._token.expires_on - time.time() < 300 -class CosmosBearerTokenCredentialPolicy(_CosmosBearerTokenCredentialPolicyBase, HTTPPolicy): +class CosmosBearerTokenCredentialPolicy(_CosmosBearerTokenCredentialPolicyBase): """Adds a Cosmos bearer token Authorization header to requests. :param credential: The credential. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py index 26491f2cb52a..9a767a3afe4f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py @@ -25,7 +25,6 @@ from typing import Any, Dict, Optional, Union, cast, Iterable, List # pylint: disable=unused-import from azure.core.tracing.decorator import distributed_trace # type: ignore -from azure.identity import ClientSecretCredential from ._cosmos_client_connection import CosmosClientConnection from ._base import build_options diff --git a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py index bfe377571bd2..2e64275ec610 100644 --- a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py +++ b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py @@ -74,10 +74,12 @@ def run_sample(): db = aad_client.get_database_client(DATABASE_ID) container = db.get_container_client(CONTAINER_ID) - print(container.read()) - container.create_item(get_test_item(0)) - print(container.read_item(item='Item_0', partition_key='Item_0')) - print(list(container.query_items(query='select * from c', partition_key='Item_0'))) + print("Container info: " + str(container.read())) + container.create_item(get_test_item(879)) + print("Point read result: " + str(container.read_item(item='Item_0', partition_key='Item_0'))) + query_results = list(container.query_items(query='select * from c', partition_key='Item_0')) + assert len(query_results) == 1 + print("Query result: " + str(query_results[0])) container.delete_item(item='Item_0', partition_key='Item_0') # Attempting to do management operations will return a 403 Forbidden exception. From c1a55b99f6b5daf5c371af518ac1a7bdbaa9e64e Mon Sep 17 00:00:00 2001 From: simorenoh Date: Mon, 28 Mar 2022 19:47:23 -0400 Subject: [PATCH 05/14] working async aad --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + .../azure/cosmos/_auth_policies.py | 168 ++++++++++++++++++ sdk/cosmos/azure-cosmos/azure/cosmos/_base.py | 5 +- .../azure/cosmos/_cosmos_client_connection.py | 12 +- .../aio/_cosmos_client_connection_async.py | 10 ++ sdk/cosmos/azure-cosmos/azure/cosmos/auth.py | 163 ----------------- .../samples/access_cosmos_with_aad_async.py | 111 ++++++++++++ .../azure-cosmos/samples/examples_async.py | 7 +- 8 files changed, 306 insertions(+), 171 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py create mode 100644 sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 1b91b0609950..aa8976a544b7 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -4,6 +4,7 @@ #### Features Added - Added support for AAD authentication for the sync client +- Added support for AAD authentication for the async client ### 4.3.0b3 (2022-03-10) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py new file mode 100644 index 000000000000..e04dc884a976 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py @@ -0,0 +1,168 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- +import asyncio +import time + +from typing import Any, Awaitable, Optional, Union +from azure.core.pipeline.policies import AsyncHTTPPolicy +from azure.core.pipeline._tools_async import await_result +from azure.core.credentials import AccessToken +from azure.core.pipeline import PipelineRequest, PipelineResponse +from azure.cosmos import http_constants + + +class _AsyncCosmosBearerTokenCredentialPolicyBase(object): + """Base class for a Bearer Token Credential Policy. + + :param credential: The credential. + :type credential: ~azure.core.credentials.TokenCredential + :param str scopes: Lets you specify the type of access needed. + """ + + def __init__(self, credential, *scopes, **kwargs): # pylint:disable=unused-argument + # type: (TokenCredential, *str, **Any) -> None + super(_AsyncCosmosBearerTokenCredentialPolicyBase, self).__init__() + self._scopes = scopes + self._credential = credential + self._token = None # type: Optional[AccessToken] + self._lock = asyncio.Lock() + + @staticmethod + def _enforce_https(request): + # type: (PipelineRequest) -> None + + # move 'enforce_https' from options to context so it persists + # across retries but isn't passed to a transport implementation + option = request.context.options.pop("enforce_https", None) + + # True is the default setting; we needn't preserve an explicit opt in to the default behavior + if option is False: + request.context["enforce_https"] = option + + enforce_https = request.context.get("enforce_https", True) + if enforce_https and not request.http_request.url.lower().startswith("https"): + raise ValueError( + "Bearer token authentication is not permitted for non-TLS protected (non-https) URLs." + ) + + @staticmethod + def _update_headers(headers, token): + # type: (Dict[str, str], str) -> None + """Updates the Authorization header with the cosmos signature and bearer token. + + :param dict headers: The HTTP Request headers + :param str token: The OAuth token. + """ + headers[http_constants.HttpHeaders.Authorization] = "type=aad&ver=1.0&sig={}".format(token) + + @property + def _need_new_token(self) -> bool: + return not self._token or self._token.expires_on - time.time() < 300 + + +class AsyncCosmosBearerTokenCredentialPolicy(_AsyncCosmosBearerTokenCredentialPolicyBase, AsyncHTTPPolicy): + """Adds a bearer token Authorization header to requests. + + :param credential: The credential. + :type credential: ~azure.core.TokenCredential + :param str scopes: Lets you specify the type of access needed. + :raises ValueError: If https_enforce does not match with endpoint being used. + """ + # def __init__(self, credential: "AsyncTokenCredential", *scopes: str, **kwargs: "Any") -> None: + # # pylint:disable=unused-argument + # super(_CosmosBearerTokenCredentialPolicyBase).__init__() + # self._lock = asyncio.Lock() + + async def on_request(self, request: "PipelineRequest") -> None: # pylint:disable=invalid-overridden-method + """Adds a bearer token Authorization header to request and sends request to next policy. + + :param request: The pipeline request object to be modified. + :type request: ~azure.core.pipeline.PipelineRequest + :raises: :class:`~azure.core.exceptions.ServiceRequestError` + """ + self._enforce_https(request) # pylint:disable=protected-access + + if self._token is None or self._need_new_token: + async with self._lock: + # double check because another coroutine may have acquired a token while we waited to acquire the lock + if self._token is None or self._need_new_token: + self._token = await self._credential.get_token(*self._scopes) + self._update_headers(request.http_request.headers, self._token.token) + + async def authorize_request(self, request: "PipelineRequest", *scopes: str, **kwargs: "Any") -> None: + """Acquire a token from the credential and authorize the request with it. + + Keyword arguments are passed to the credential's get_token method. The token will be cached and used to + authorize future requests. + + :param ~azure.core.pipeline.PipelineRequest request: the request + :param str scopes: required scopes of authentication + """ + async with self._lock: + self._token = await self._credential.get_token(*scopes, **kwargs) + self._update_headers(request.http_request.headers, self._token.token) + + async def send(self, request: "PipelineRequest") -> "PipelineResponse": + """Authorize request with a bearer token and send it to the next policy + + :param request: The pipeline request object + :type request: ~azure.core.pipeline.PipelineRequest + """ + await await_result(self.on_request, request) + try: + response = await self.next.send(request) + await await_result(self.on_response, request, response) + except Exception: # pylint:disable=broad-except + handled = await await_result(self.on_exception, request) + if not handled: + raise + else: + if response.http_response.status_code == 401: + self._token = None # any cached token is invalid + if "WWW-Authenticate" in response.http_response.headers: + request_authorized = await self.on_challenge(request, response) + if request_authorized: + try: + response = await self.next.send(request) + await await_result(self.on_response, request, response) + except Exception: # pylint:disable=broad-except + handled = await await_result(self.on_exception, request) + if not handled: + raise + + return response + + async def on_challenge(self, request: "PipelineRequest", response: "PipelineResponse") -> bool: + """Authorize request according to an authentication challenge + + This method is called when the resource provider responds 401 with a WWW-Authenticate header. + + :param ~azure.core.pipeline.PipelineRequest request: the request which elicited an authentication challenge + :param ~azure.core.pipeline.PipelineResponse response: the resource provider's response + :returns: a bool indicating whether the policy should send the request + """ + # pylint:disable=unused-argument,no-self-use + return False + + def on_response(self, request: "PipelineRequest", response: "PipelineResponse") -> "Union[None, Awaitable[None]]": + """Executed after the request comes back from the next policy. + + :param request: Request to be modified after returning from the policy. + :type request: ~azure.core.pipeline.PipelineRequest + :param response: Pipeline response object + :type response: ~azure.core.pipeline.PipelineResponse + """ + + def on_exception(self, request: "PipelineRequest") -> None: + """Executed when an exception is raised while executing the next policy. + + This method is executed inside the exception handler. + + :param request: The Pipeline request object + :type request: ~azure.core.pipeline.PipelineRequest + """ + # pylint: disable=no-self-use,unused-argument + return diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py index 5b5d14b94bf6..c9a8af14ecb6 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py @@ -30,6 +30,7 @@ from typing import Dict, Any from urllib.parse import quote as urllib_quote +from urllib.parse import urlsplit from azure.core import MatchConditions @@ -664,8 +665,8 @@ def ParsePaths(paths): def create_scope_from_url(url): - # Need a better way to do this, regex? - return url.replace(":443", "") + ".default" + parsed_url = urlsplit(url) + return parsed_url.scheme + "://" + parsed_url.hostname + "/.default" def validate_cache_staleness_value(max_integrated_cache_staleness): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index c4e30f1d9ce3..b30fb1bb4288 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -58,7 +58,7 @@ from . import _session from . import _utils from .partition_key import _Undefined, _Empty -from .auth import CosmosBearerTokenCredentialPolicy +from ._auth_policies import AsyncCosmosBearerTokenCredentialPolicy ClassType = TypeVar("ClassType") @@ -179,11 +179,17 @@ def __init__( self._user_agent = _utils.get_user_agent() + credentials_policy = None + if self.aad_credentials: + scopes = base.create_scope_from_url(self.url_connection) + credentials_policy = AsyncCosmosBearerTokenCredentialPolicy(self.aad_credentials, scopes) + policies = [ HeadersPolicy(**kwargs), ProxyPolicy(proxies=proxies), UserAgentPolicy(base_user_agent=self._user_agent, **kwargs), ContentDecodePolicy(), + credentials_policy, retry_policy, CustomHookPolicy(**kwargs), NetworkTraceLoggingPolicy(**kwargs), @@ -191,10 +197,6 @@ def __init__( HttpLoggingPolicy(**kwargs), ] - if self.aad_credentials: - scopes = base.create_scope_from_url(self.url_connection) - policies.append(CosmosBearerTokenCredentialPolicy(self.aad_credentials, scopes)) - transport = kwargs.pop("transport", None) self.pipeline_client = PipelineClient(base_url=url_connection, transport=transport, policies=policies) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index dab5b82b52a0..ea07785be15e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -58,10 +58,12 @@ from .. import _session from .. import _utils from ..partition_key import _Undefined, _Empty +from .._auth_policies import AsyncCosmosBearerTokenCredentialPolicy ClassType = TypeVar("ClassType") # pylint: disable=protected-access + class CosmosClientConnection(object): # pylint: disable=too-many-public-methods,too-many-instance-attributes """Represents a document client. @@ -113,9 +115,11 @@ def __init__( self.master_key = None self.resource_tokens = None + self.aad_credentials = None if auth is not None: self.master_key = auth.get("masterKey") self.resource_tokens = auth.get("resourceTokens") + self.aad_credentials = auth.get("clientSecretCredential") if auth.get("permissionFeed"): self.resource_tokens = {} @@ -176,11 +180,17 @@ def __init__( self._user_agent = _utils.get_user_agent_async() + credentials_policy = None + if self.aad_credentials: + scopes = base.create_scope_from_url(self.url_connection) + credentials_policy = AsyncCosmosBearerTokenCredentialPolicy(self.aad_credentials, scopes) + policies = [ HeadersPolicy(**kwargs), ProxyPolicy(proxies=proxies), UserAgentPolicy(base_user_agent=self._user_agent, **kwargs), ContentDecodePolicy(), + credentials_policy, retry_policy, CustomHookPolicy(**kwargs), NetworkTraceLoggingPolicy(**kwargs), diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py b/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py index a218a700023f..d3c2cb0d8017 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/auth.py @@ -26,14 +26,8 @@ from hashlib import sha256 import hmac import urllib.parse -import time -from typing import Any, Dict, Optional, ClassVar -from azure.core.exceptions import ServiceRequestError -from azure.core.credentials import AccessToken -from azure.core.pipeline import PipelineRequest, PipelineResponse from . import http_constants -TokenCredential = ClassVar def GetAuthorizationHeader( cosmos_client_connection, verb, path, resource_id_or_fullname, is_name_based, resource_type, headers @@ -149,160 +143,3 @@ def __get_authorization_token_using_resource_token(resource_tokens, path, resour return resource_tokens[sub_path] return None - - -# pylint:disable=too-few-public-methods -class _CosmosBearerTokenCredentialPolicyBase(object): - """Base class for a Cosmos Bearer Token Credential Policy. - Based on ~azure.core.pipeline.policies.BearerTokenCredentialPolicy. - Takes care of updating token as needed and updating the request headers. - - :param credential: The credential. - :type credential: ~azure.core.credentials.TokenCredential - :param str scopes: Lets you specify the type of access needed. - """ - - def __init__(self, credential, *scopes, **kwargs): # pylint:disable=unused-argument - # type: (TokenCredential, *str, **Any) -> None - super(_CosmosBearerTokenCredentialPolicyBase, self).__init__() - self._scopes = scopes - self._credential = credential - self._token = None # type: Optional[AccessToken] - - @staticmethod - def _enforce_https(request): - # type: (PipelineRequest) -> None - - # move 'enforce_https' from options to context so it persists - # across retries but isn't passed to a transport implementation - option = request.context.options.pop("enforce_https", None) - - # True is the default setting; we needn't preserve an explicit opt in to the default behavior - if option is False: - request.context["enforce_https"] = option - - enforce_https = request.context.get("enforce_https", True) - if enforce_https and not request.http_request.url.lower().startswith("https"): - raise ServiceRequestError( - "Bearer token authentication is not permitted for non-TLS protected (non-https) URLs." - ) - - @staticmethod - def _update_headers(headers, token): - # type: (Dict[str, str], str) -> None - """Updates the Authorization header with the Cosmos AAD bearer token. - - :param dict headers: The HTTP Request headers - :param str token: The OAuth token. - """ - headers[http_constants.HttpHeaders.Authorization] = "type=aad&ver=1.0&sig={}".format(token) - - @property - def _need_new_token(self): - # type: () -> bool - return not self._token or self._token.expires_on - time.time() < 300 - - -class CosmosBearerTokenCredentialPolicy(_CosmosBearerTokenCredentialPolicyBase): - """Adds a Cosmos bearer token Authorization header to requests. - - :param credential: The credential. - :type credential: ~azure.core.TokenCredential - :param str scopes: Lets you specify the type of access needed. - :raises: :class:`~azure.core.exceptions.ServiceRequestError` - """ - - def on_request(self, request): - # type: (PipelineRequest) -> None - """Called before the policy sends a request. - - The base implementation authorizes the request with a bearer token. - - :param ~azure.core.pipeline.PipelineRequest request: the request - """ - self._enforce_https(request) - - if self._token is None or self._need_new_token: - self._token = self._credential.get_token(*self._scopes) - self._update_headers(request.http_request.headers, self._token.token) - - def authorize_request(self, request, *scopes, **kwargs): - # type: (PipelineRequest, *str, **Any) -> None - """Acquire a token from the credential and authorize the request with it. - - Keyword arguments are passed to the credential's get_token method. The token will be cached and used to - authorize future requests. - - :param ~azure.core.pipeline.PipelineRequest request: the request - :param str scopes: required scopes of authentication - """ - self._token = self._credential.get_token(*scopes, **kwargs) - self._update_headers(request.http_request.headers, self._token.token) - - def send(self, request): - # type: (PipelineRequest) -> PipelineResponse - """Authorize request with a bearer token and send it to the next policy - - :param request: The pipeline request object - :type request: ~azure.core.pipeline.PipelineRequest - """ - self.on_request(request) - try: - response = self.next.send(request) - self.on_response(request, response) - except Exception: # pylint:disable=broad-except - handled = self.on_exception(request) - if not handled: - raise - else: - if response.http_response.status_code == 401: - self._token = None # any cached token is invalid - if "WWW-Authenticate" in response.http_response.headers: - request_authorized = self.on_challenge(request, response) - if request_authorized: - try: - response = self.next.send(request) - self.on_response(request, response) - except Exception: # pylint:disable=broad-except - handled = self.on_exception(request) - if not handled: - raise - - return response - - def on_challenge(self, request, response): - # type: (PipelineRequest, PipelineResponse) -> bool - """Authorize request according to an authentication challenge - - This method is called when the resource provider responds 401 with a WWW-Authenticate header. - - :param ~azure.core.pipeline.PipelineRequest request: the request which elicited an authentication challenge - :param ~azure.core.pipeline.PipelineResponse response: the resource provider's response - :returns: a bool indicating whether the policy should send the request - """ - # pylint:disable=unused-argument,no-self-use - return False - - def on_response(self, request, response): - # type: (PipelineRequest, PipelineResponse) -> None - """Executed after the request comes back from the next policy. - - :param request: Request to be modified after returning from the policy. - :type request: ~azure.core.pipeline.PipelineRequest - :param response: Pipeline response object - :type response: ~azure.core.pipeline.PipelineResponse - """ - - def on_exception(self, request): - # type: (PipelineRequest) -> bool - """Executed when an exception is raised while executing the next policy. - - This method is executed inside the exception handler. - - :param request: The Pipeline request object - :type request: ~azure.core.pipeline.PipelineRequest - :return: False by default, override with True to stop the exception. - :rtype: bool - """ - # pylint: disable=no-self-use,unused-argument - return False diff --git a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py new file mode 100644 index 000000000000..0fc6b73048df --- /dev/null +++ b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py @@ -0,0 +1,111 @@ +import azure.cosmos.aio.cosmos_client as cosmos_client +import azure.cosmos.exceptions as exceptions +from azure.cosmos.partition_key import PartitionKey +from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential, EnvironmentCredential +import config +import asyncio + +# ---------------------------------------------------------------------------------------------------------- +# Prerequistes - +# +# 1. An Azure Cosmos account - +# https://docs.microsoft.com/azure/cosmos-db/create-sql-api-python#create-a-database-account +# +# 2. Microsoft Azure Cosmos +# pip install azure-cosmos>=4.3.0b4 +# ---------------------------------------------------------------------------------------------------------- +# Sample - demonstrates how to authenticate and use your database account using AAD credentials +# Read more about operations allowed for this authorization method: https://aka.ms/cosmos-native-rbac +# ---------------------------------------------------------------------------------------------------------- +# Note: +# This sample creates a Container to your database account. +# Each time a Container is created the account will be billed for 1 hour of usage based on +# the provisioned throughput (RU/s) of that account. +# ---------------------------------------------------------------------------------------------------------- +# +HOST = config.settings["host"] +MASTER_KEY = config.settings["master_key"] + +TENANT_ID = config.settings["tenant_id"] +CLIENT_ID = config.settings["client_id"] +CLIENT_SECRET = config.settings["client_secret"] + +DATABASE_ID = config.settings["database_id"] +CONTAINER_ID = config.settings["container_id"] +PARTITION_KEY = PartitionKey(path="/id") + + +def get_test_item(num): + test_item = { + 'id': 'Item_' + str(num), + 'test_object': True, + 'lastName': 'Smith' + } + return test_item + + +async def create_sample_resources(): + print("creating sample resources") + async with cosmos_client.CosmosClient(HOST, MASTER_KEY) as client: + db = await client.create_database(DATABASE_ID) + await db.create_container(id=CONTAINER_ID, partition_key=PARTITION_KEY) + + +async def delete_sample_resources(): + print("deleting sample resources") + async with cosmos_client.CosmosClient(HOST, MASTER_KEY) as client: + await client.delete_database(DATABASE_ID) + + +async def run_sample(): + # Since Azure Cosmos DB data plane SDK does not cover management operations, we have to create our resources + # with a master key authenticated client for this sample. + await create_sample_resources() + + # With this done, you can use your AAD service principal id and secret to create your ClientSecretCredential. + # The async ClientSecretCredentials, like the async client, also have a context manager, + # and as such should be used with the `async with` keywords. + async with ClientSecretCredential( + tenant_id=TENANT_ID, + client_id=CLIENT_ID, + client_secret=CLIENT_SECRET) as aad_credentials: + + # Use your credentials to authenticate your client. + async with cosmos_client.CosmosClient(HOST, aad_credentials) as aad_client: + print("Showed ClientSecretCredential, now showing DefaultAzureCredential") + + # You can also utilize DefaultAzureCredential rather than directly passing in the id's and secrets. + # This is the recommended method of authentication, and uses environment variables rather than in-code strings. + async with DefaultAzureCredential() as aad_credentials: + + # Use your credentials to authenticate your client. + async with cosmos_client.CosmosClient(HOST, aad_credentials) as aad_client: + + # Do any R/W data operations with your authorized AAD client. + db = aad_client.get_database_client(DATABASE_ID) + container = db.get_container_client(CONTAINER_ID) + + print("Container info: " + str(container.read())) + await container.create_item(get_test_item(879)) + print("Point read result: " + str(container.read_item(item='Item_0', partition_key='Item_0'))) + query_results = [item async for item in + container.query_items(query='select * from c', partition_key='Item_0')] + assert len(query_results) == 1 + print("Query result: " + str(query_results[0])) + await container.delete_item(item='Item_0', partition_key='Item_0') + + # Attempting to do management operations will return a 403 Forbidden exception. + try: + await aad_client.delete_database(DATABASE_ID) + except exceptions.CosmosHttpResponseError as e: + assert e.status_code == 403 + print("403 error assertion success") + + # To clean up the sample, we use a master key client again to get access to deleting containers/ databases. + await delete_sample_resources() + print("end of sample") + + +if __name__ == "__main__": + loop = asyncio.get_event_loop() + loop.run_until_complete(run_sample()) diff --git a/sdk/cosmos/azure-cosmos/samples/examples_async.py b/sdk/cosmos/azure-cosmos/samples/examples_async.py index e0c1693e3da7..1da9eea1b980 100644 --- a/sdk/cosmos/azure-cosmos/samples/examples_async.py +++ b/sdk/cosmos/azure-cosmos/samples/examples_async.py @@ -1,3 +1,4 @@ +import asyncio from azure.cosmos import exceptions, CosmosClient, PartitionKey from azure.cosmos.aio import CosmosClient @@ -76,7 +77,7 @@ async def examples_async(): # [START list_containers] database = client.get_database_client(database_name) - for container in database.list_containers(): + async for container in database.list_containers(): print("Container ID: {}".format(container['id'])) # [END list_containers] @@ -160,3 +161,7 @@ async def examples_async(): except exceptions.CosmosHttpResponseError as failure: print("Failed to create user. Status code:{}".format(failure.status_code)) # [END create_user] + +if __name__ == "__main__": + loop = asyncio.get_event_loop() + loop.run_until_complete(examples_async()) From e19ae40c5ebb08758560107edbb1e1c0a7c87de0 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Mon, 28 Mar 2022 19:54:57 -0400 Subject: [PATCH 06/14] Delete access_cosmos_with_aad.py snuck its way into the async PR --- .../samples/access_cosmos_with_aad.py | 98 ------------------- 1 file changed, 98 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py diff --git a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py deleted file mode 100644 index 2e64275ec610..000000000000 --- a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py +++ /dev/null @@ -1,98 +0,0 @@ -import azure.cosmos.cosmos_client as cosmos_client -import azure.cosmos.exceptions as exceptions -from azure.cosmos.partition_key import PartitionKey -from azure.identity import ClientSecretCredential -import config - -# ---------------------------------------------------------------------------------------------------------- -# Prerequistes - -# -# 1. An Azure Cosmos account - -# https://docs.microsoft.com/azure/cosmos-db/create-sql-api-python#create-a-database-account -# -# 2. Microsoft Azure Cosmos -# pip install azure-cosmos>=4.3.0b4 -# ---------------------------------------------------------------------------------------------------------- -# Sample - demonstrates how to authenticate and use your database account using AAD credentials -# Read more about operations allowed for this authorization method: https://aka.ms/cosmos-native-rbac -# ---------------------------------------------------------------------------------------------------------- -# Note: -# This sample creates a Container to your database account. -# Each time a Container is created the account will be billed for 1 hour of usage based on -# the provisioned throughput (RU/s) of that account. -# ---------------------------------------------------------------------------------------------------------- -# -HOST = config.settings["host"] -MASTER_KEY = config.settings["master_key"] - -TENANT_ID = config.settings["tenant_id"] -CLIENT_ID = config.settings["client_id"] -CLIENT_SECRET = config.settings["client_secret"] - -DATABASE_ID = config.settings["database_id"] -CONTAINER_ID = config.settings["container_id"] -PARTITION_KEY = PartitionKey(path="/id") - - -def get_test_item(num): - test_item = { - 'id': 'Item_' + str(num), - 'test_object': True, - 'lastName': 'Smith' - } - return test_item - - -def create_sample_resources(): - print("creating sample resources") - client = cosmos_client.CosmosClient(HOST, MASTER_KEY) - db = client.create_database(DATABASE_ID) - db.create_container(id=CONTAINER_ID, partition_key=PARTITION_KEY) - - -def delete_sample_resources(): - print("deleting sample resources") - client = cosmos_client.CosmosClient(HOST, MASTER_KEY) - client.delete_database(DATABASE_ID) - - -def run_sample(): - # Since Azure Cosmos DB data plane SDK does not cover management operations, we have to create our resources - # with a master key authenticated client for this sample. - create_sample_resources() - - # With this done, you can use your AAD service principal id and secret to create your ClientSecretCredential. - aad_credentials = ClientSecretCredential( - tenant_id=TENANT_ID, - client_id=CLIENT_ID, - client_secret=CLIENT_SECRET) - - # Use your credentials to authenticate your client. - aad_client = cosmos_client.CosmosClient(HOST, aad_credentials) - - # Do any R/W data operations with your authorized AAD client. - db = aad_client.get_database_client(DATABASE_ID) - container = db.get_container_client(CONTAINER_ID) - - print("Container info: " + str(container.read())) - container.create_item(get_test_item(879)) - print("Point read result: " + str(container.read_item(item='Item_0', partition_key='Item_0'))) - query_results = list(container.query_items(query='select * from c', partition_key='Item_0')) - assert len(query_results) == 1 - print("Query result: " + str(query_results[0])) - container.delete_item(item='Item_0', partition_key='Item_0') - - # Attempting to do management operations will return a 403 Forbidden exception. - try: - aad_client.delete_database(DATABASE_ID) - except exceptions.CosmosHttpResponseError as e: - assert e.status_code == 403 - print("403 error assertion success") - - # To clean up the sample, we use a master key client again to get access to deleting containers and databases. - delete_sample_resources() - print("end of sample") - - -if __name__ == "__main__": - run_sample() From aaaf3fcb2c326fa6dfd0b06f7ab8c974d3432991 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Mon, 28 Mar 2022 20:18:52 -0400 Subject: [PATCH 07/14] Update _auth_policies.py --- sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py index e04dc884a976..619112c91b33 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py @@ -71,10 +71,6 @@ class AsyncCosmosBearerTokenCredentialPolicy(_AsyncCosmosBearerTokenCredentialPo :param str scopes: Lets you specify the type of access needed. :raises ValueError: If https_enforce does not match with endpoint being used. """ - # def __init__(self, credential: "AsyncTokenCredential", *scopes: str, **kwargs: "Any") -> None: - # # pylint:disable=unused-argument - # super(_CosmosBearerTokenCredentialPolicyBase).__init__() - # self._lock = asyncio.Lock() async def on_request(self, request: "PipelineRequest") -> None: # pylint:disable=invalid-overridden-method """Adds a bearer token Authorization header to request and sends request to next policy. From 083b115f89ab5b9ca454386524242b691867825a Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 29 Mar 2022 11:59:32 -0400 Subject: [PATCH 08/14] small changes --- .../azure/cosmos/_cosmos_client_connection.py | 8 -------- .../{_auth_policies.py => aio/_auth_policy_async.py} | 0 .../cosmos/aio/_cosmos_client_connection_async.py | 4 ++-- .../samples/access_cosmos_with_aad_async.py | 12 ++++++------ 4 files changed, 8 insertions(+), 16 deletions(-) rename sdk/cosmos/azure-cosmos/azure/cosmos/{_auth_policies.py => aio/_auth_policy_async.py} (100%) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index b30fb1bb4288..f7048592a145 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -58,11 +58,8 @@ from . import _session from . import _utils from .partition_key import _Undefined, _Empty -from ._auth_policies import AsyncCosmosBearerTokenCredentialPolicy ClassType = TypeVar("ClassType") - - # pylint: disable=protected-access @@ -179,17 +176,12 @@ def __init__( self._user_agent = _utils.get_user_agent() - credentials_policy = None - if self.aad_credentials: - scopes = base.create_scope_from_url(self.url_connection) - credentials_policy = AsyncCosmosBearerTokenCredentialPolicy(self.aad_credentials, scopes) policies = [ HeadersPolicy(**kwargs), ProxyPolicy(proxies=proxies), UserAgentPolicy(base_user_agent=self._user_agent, **kwargs), ContentDecodePolicy(), - credentials_policy, retry_policy, CustomHookPolicy(**kwargs), NetworkTraceLoggingPolicy(**kwargs), diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py similarity index 100% rename from sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policies.py rename to sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index ea07785be15e..408cd0acd1e9 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -58,7 +58,7 @@ from .. import _session from .. import _utils from ..partition_key import _Undefined, _Empty -from .._auth_policies import AsyncCosmosBearerTokenCredentialPolicy +from ._auth_policy_async import AsyncCosmosBearerTokenCredentialPolicy ClassType = TypeVar("ClassType") # pylint: disable=protected-access @@ -190,8 +190,8 @@ def __init__( ProxyPolicy(proxies=proxies), UserAgentPolicy(base_user_agent=self._user_agent, **kwargs), ContentDecodePolicy(), - credentials_policy, retry_policy, + credentials_policy, CustomHookPolicy(**kwargs), NetworkTraceLoggingPolicy(**kwargs), DistributedTracingPolicy(**kwargs), diff --git a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py index 0fc6b73048df..90df67f03735 100644 --- a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py +++ b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py @@ -1,7 +1,7 @@ -import azure.cosmos.aio.cosmos_client as cosmos_client +from azure.cosmos.aio import CosmosClient import azure.cosmos.exceptions as exceptions from azure.cosmos.partition_key import PartitionKey -from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential, EnvironmentCredential +from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential import config import asyncio @@ -46,14 +46,14 @@ def get_test_item(num): async def create_sample_resources(): print("creating sample resources") - async with cosmos_client.CosmosClient(HOST, MASTER_KEY) as client: + async with CosmosClient(HOST, MASTER_KEY) as client: db = await client.create_database(DATABASE_ID) await db.create_container(id=CONTAINER_ID, partition_key=PARTITION_KEY) async def delete_sample_resources(): print("deleting sample resources") - async with cosmos_client.CosmosClient(HOST, MASTER_KEY) as client: + async with CosmosClient(HOST, MASTER_KEY) as client: await client.delete_database(DATABASE_ID) @@ -71,7 +71,7 @@ async def run_sample(): client_secret=CLIENT_SECRET) as aad_credentials: # Use your credentials to authenticate your client. - async with cosmos_client.CosmosClient(HOST, aad_credentials) as aad_client: + async with CosmosClient(HOST, aad_credentials) as aad_client: print("Showed ClientSecretCredential, now showing DefaultAzureCredential") # You can also utilize DefaultAzureCredential rather than directly passing in the id's and secrets. @@ -79,7 +79,7 @@ async def run_sample(): async with DefaultAzureCredential() as aad_credentials: # Use your credentials to authenticate your client. - async with cosmos_client.CosmosClient(HOST, aad_credentials) as aad_client: + async with CosmosClient(HOST, aad_credentials) as aad_client: # Do any R/W data operations with your authorized AAD client. db = aad_client.get_database_client(DATABASE_ID) From b6c192ebe2676e9e2237d7442887b5a17a070d30 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 29 Mar 2022 12:00:39 -0400 Subject: [PATCH 09/14] Update _cosmos_client_connection.py --- .../azure-cosmos/azure/cosmos/_cosmos_client_connection.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index f7048592a145..2375052779b0 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -114,11 +114,9 @@ def __init__( self.master_key = None self.resource_tokens = None - self.aad_credentials = None if auth is not None: self.master_key = auth.get("masterKey") self.resource_tokens = auth.get("resourceTokens") - self.aad_credentials = auth.get("clientSecretCredential") if auth.get("permissionFeed"): self.resource_tokens = {} @@ -176,7 +174,6 @@ def __init__( self._user_agent = _utils.get_user_agent() - policies = [ HeadersPolicy(**kwargs), ProxyPolicy(proxies=proxies), From c2cb7ee9a29e320ec6411f1fe321daa2966fd233 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 29 Mar 2022 12:01:56 -0400 Subject: [PATCH 10/14] removing changes made in sync --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 - sdk/cosmos/azure-cosmos/README.md | 22 ---------------------- 2 files changed, 23 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index aa8976a544b7..b403f760971d 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -3,7 +3,6 @@ ### 4.3.0b4 (Unreleased) #### Features Added -- Added support for AAD authentication for the sync client - Added support for AAD authentication for the async client ### 4.3.0b3 (2022-03-10) diff --git a/sdk/cosmos/azure-cosmos/README.md b/sdk/cosmos/azure-cosmos/README.md index b17b51d13b1d..ab46e9dc19f2 100644 --- a/sdk/cosmos/azure-cosmos/README.md +++ b/sdk/cosmos/azure-cosmos/README.md @@ -76,28 +76,6 @@ KEY = os.environ['ACCOUNT_KEY'] client = CosmosClient(URL, credential=KEY) ``` -### AAD Authentication - -You can also authenticate a client utilizing your service principal's AAD credentials and the azure identity package: -```Python -from azure.cosmos import CosmosClient -from azure.identity import ClientSecretCredential - -import os -URL = os.environ['ACCOUNT_URI'] -TENANT_ID = os.environ['TENANT_ID'] -CLIENT_ID = os.environ['CLIENT_ID'] -CLIENT_SECRET = os.environ['CLIENT_SECRET'] - -AAD_CREDENTIALS = ClientSecretCredential( - tenant_id=TENANT_ID, - client_id=CLIENT_ID, - client_secret=CLIENT_SECRET) - -client = CosmosClient(URL, AAD_CREDENTIALS) -``` -More information on allowed operations for AAD authenticated clients: [RBAC Permission Model](https://aka.ms/cosmos-native-rbac) - ## Key concepts Once you've initialized a [CosmosClient][ref_cosmosclient], you can interact with the primary resource types in Cosmos DB: From 12460f9262184d54e78a03092bfbfd9d02e75ad8 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 5 Apr 2022 13:07:08 -0400 Subject: [PATCH 11/14] Update _auth_policy_async.py --- .../azure-cosmos/azure/cosmos/aio/_auth_policy_async.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py index 619112c91b33..9c2fad41c0a4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py @@ -143,7 +143,7 @@ async def on_challenge(self, request: "PipelineRequest", response: "PipelineResp # pylint:disable=unused-argument,no-self-use return False - def on_response(self, request: "PipelineRequest", response: "PipelineResponse") -> "Union[None, Awaitable[None]]": + def on_response(self, request: PipelineRequest, response: PipelineResponse) -> Union[None, Awaitable[None]]: """Executed after the request comes back from the next policy. :param request: Request to be modified after returning from the policy. @@ -152,7 +152,7 @@ def on_response(self, request: "PipelineRequest", response: "PipelineResponse") :type response: ~azure.core.pipeline.PipelineResponse """ - def on_exception(self, request: "PipelineRequest") -> None: + def on_exception(self, request: PipelineRequest) -> None: """Executed when an exception is raised while executing the next policy. This method is executed inside the exception handler. From 414ebb45b26696c272e771165ca36157b7afe0ef Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 5 Apr 2022 13:12:34 -0400 Subject: [PATCH 12/14] Update _auth_policy_async.py --- sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py index 9c2fad41c0a4..8a268387c281 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py @@ -52,6 +52,9 @@ def _enforce_https(request): def _update_headers(headers, token): # type: (Dict[str, str], str) -> None """Updates the Authorization header with the cosmos signature and bearer token. + This is the main method that differentiates this policy from core's BearerTokenCredentialPolicy and works + to properly sign the authorization header for Cosmos' REST API. For more information: + https://docs.microsoft.com/rest/api/cosmos-db/access-control-on-cosmosdb-resources#authorization-header :param dict headers: The HTTP Request headers :param str token: The OAuth token. From fe9706f81789d57b0250cf70b6185f9197fbc8bf Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 5 Apr 2022 20:05:32 -0400 Subject: [PATCH 13/14] Update _auth_policy_async.py --- .../azure/cosmos/aio/_auth_policy_async.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py index 8a268387c281..f5913f5cf56a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py @@ -6,14 +6,22 @@ import asyncio import time -from typing import Any, Awaitable, Optional, Union +from typing import Any, Awaitable, Optional, Dict, Union from azure.core.pipeline.policies import AsyncHTTPPolicy -from azure.core.pipeline._tools_async import await_result from azure.core.credentials import AccessToken from azure.core.pipeline import PipelineRequest, PipelineResponse from azure.cosmos import http_constants +async def await_result(func, *args, **kwargs): + """If func returns an awaitable, await it.""" + result = func(*args, **kwargs) + if hasattr(result, "__await__"): + # type ignore on await: https://github.com/python/mypy/issues/7587 + return await result # type: ignore + return result + + class _AsyncCosmosBearerTokenCredentialPolicyBase(object): """Base class for a Bearer Token Credential Policy. From aea3f8a5867aac3275f69faafc0d7b8c256e40dd Mon Sep 17 00:00:00 2001 From: simorenoh Date: Wed, 6 Apr 2022 10:48:03 -0400 Subject: [PATCH 14/14] added licenses to samples --- .../samples/access_cosmos_with_aad.py | 5 ++++ .../samples/access_cosmos_with_aad_async.py | 5 ++++ .../access_cosmos_with_resource_token.py | 5 ++++ ...access_cosmos_with_resource_token_async.py | 5 ++++ .../samples/change_feed_management.py | 5 ++++ .../samples/change_feed_management_async.py | 5 ++++ sdk/cosmos/azure-cosmos/samples/config.py | 5 ++++ .../samples/container_management.py | 5 ++++ .../samples/container_management_async.py | 5 ++++ .../samples/database_management.py | 5 ++++ .../samples/database_management_async.py | 5 ++++ .../samples/document_management.py | 5 ++++ .../samples/document_management_async.py | 5 ++++ sdk/cosmos/azure-cosmos/samples/examples.py | 5 ++++ .../azure-cosmos/samples/examples_async.py | 5 ++++ .../azure-cosmos/samples/index_management.py | 5 ++++ .../samples/index_management_async.py | 5 ++++ .../nonpartitioned_container_operations.py | 25 ++++--------------- ...npartitioned_container_operations_async.py | 5 ++++ 19 files changed, 95 insertions(+), 20 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py index 33f080fce9f5..5033749712d3 100644 --- a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py +++ b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- from azure.cosmos import CosmosClient import azure.cosmos.exceptions as exceptions from azure.cosmos.partition_key import PartitionKey diff --git a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py index 90df67f03735..4f6849a3f487 100644 --- a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py +++ b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- from azure.cosmos.aio import CosmosClient import azure.cosmos.exceptions as exceptions from azure.cosmos.partition_key import PartitionKey diff --git a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_resource_token.py b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_resource_token.py index 99d09bc1e499..c4ae0e353c29 100644 --- a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_resource_token.py +++ b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_resource_token.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions from azure.cosmos.partition_key import PartitionKey diff --git a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_resource_token_async.py b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_resource_token_async.py index 367b4249b827..47319c9da30f 100644 --- a/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_resource_token_async.py +++ b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_resource_token_async.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.aio.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions from azure.cosmos.partition_key import PartitionKey diff --git a/sdk/cosmos/azure-cosmos/samples/change_feed_management.py b/sdk/cosmos/azure-cosmos/samples/change_feed_management.py index a3149bec21b9..ea0db9df9446 100644 --- a/sdk/cosmos/azure-cosmos/samples/change_feed_management.py +++ b/sdk/cosmos/azure-cosmos/samples/change_feed_management.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.documents as documents import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions diff --git a/sdk/cosmos/azure-cosmos/samples/change_feed_management_async.py b/sdk/cosmos/azure-cosmos/samples/change_feed_management_async.py index 027e1a773d93..b6918258544a 100644 --- a/sdk/cosmos/azure-cosmos/samples/change_feed_management_async.py +++ b/sdk/cosmos/azure-cosmos/samples/change_feed_management_async.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.aio.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions import azure.cosmos.documents as documents diff --git a/sdk/cosmos/azure-cosmos/samples/config.py b/sdk/cosmos/azure-cosmos/samples/config.py index a69ee67c0e88..9dfcf802423f 100644 --- a/sdk/cosmos/azure-cosmos/samples/config.py +++ b/sdk/cosmos/azure-cosmos/samples/config.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import os settings = { diff --git a/sdk/cosmos/azure-cosmos/samples/container_management.py b/sdk/cosmos/azure-cosmos/samples/container_management.py index 6ec53a68f136..c826f9b4aa12 100644 --- a/sdk/cosmos/azure-cosmos/samples/container_management.py +++ b/sdk/cosmos/azure-cosmos/samples/container_management.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions from azure.cosmos.partition_key import PartitionKey diff --git a/sdk/cosmos/azure-cosmos/samples/container_management_async.py b/sdk/cosmos/azure-cosmos/samples/container_management_async.py index 96d46124965e..4b3ab86a843d 100644 --- a/sdk/cosmos/azure-cosmos/samples/container_management_async.py +++ b/sdk/cosmos/azure-cosmos/samples/container_management_async.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.aio.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions from azure.cosmos.partition_key import PartitionKey diff --git a/sdk/cosmos/azure-cosmos/samples/database_management.py b/sdk/cosmos/azure-cosmos/samples/database_management.py index 31ccce6b2ba5..5e747c5d2485 100644 --- a/sdk/cosmos/azure-cosmos/samples/database_management.py +++ b/sdk/cosmos/azure-cosmos/samples/database_management.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions diff --git a/sdk/cosmos/azure-cosmos/samples/database_management_async.py b/sdk/cosmos/azure-cosmos/samples/database_management_async.py index 97da4530bc56..fbb4aa910167 100644 --- a/sdk/cosmos/azure-cosmos/samples/database_management_async.py +++ b/sdk/cosmos/azure-cosmos/samples/database_management_async.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.aio.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions diff --git a/sdk/cosmos/azure-cosmos/samples/document_management.py b/sdk/cosmos/azure-cosmos/samples/document_management.py index 20319aeed14d..7ebc0b869ed9 100644 --- a/sdk/cosmos/azure-cosmos/samples/document_management.py +++ b/sdk/cosmos/azure-cosmos/samples/document_management.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions from azure.cosmos.partition_key import PartitionKey diff --git a/sdk/cosmos/azure-cosmos/samples/document_management_async.py b/sdk/cosmos/azure-cosmos/samples/document_management_async.py index 77eb9f36d69c..fb5e86c462ca 100644 --- a/sdk/cosmos/azure-cosmos/samples/document_management_async.py +++ b/sdk/cosmos/azure-cosmos/samples/document_management_async.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.aio.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions from azure.cosmos.partition_key import PartitionKey diff --git a/sdk/cosmos/azure-cosmos/samples/examples.py b/sdk/cosmos/azure-cosmos/samples/examples.py index 4edf2ca5d321..c28195014a16 100644 --- a/sdk/cosmos/azure-cosmos/samples/examples.py +++ b/sdk/cosmos/azure-cosmos/samples/examples.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- # These examples are ingested by the documentation system, and are # displayed in the SDK reference documentation. When editing these # example snippets, take into consideration how this might affect diff --git a/sdk/cosmos/azure-cosmos/samples/examples_async.py b/sdk/cosmos/azure-cosmos/samples/examples_async.py index 5bdb16a5c37b..97edcd7c9519 100644 --- a/sdk/cosmos/azure-cosmos/samples/examples_async.py +++ b/sdk/cosmos/azure-cosmos/samples/examples_async.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import asyncio from azure.cosmos import exceptions, CosmosClient, PartitionKey from azure.cosmos.aio import CosmosClient diff --git a/sdk/cosmos/azure-cosmos/samples/index_management.py b/sdk/cosmos/azure-cosmos/samples/index_management.py index 82105cdab54b..d43e86a16386 100644 --- a/sdk/cosmos/azure-cosmos/samples/index_management.py +++ b/sdk/cosmos/azure-cosmos/samples/index_management.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.documents as documents import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions diff --git a/sdk/cosmos/azure-cosmos/samples/index_management_async.py b/sdk/cosmos/azure-cosmos/samples/index_management_async.py index 327909957e52..9c160fa508f8 100644 --- a/sdk/cosmos/azure-cosmos/samples/index_management_async.py +++ b/sdk/cosmos/azure-cosmos/samples/index_management_async.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.documents as documents import azure.cosmos.aio.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions diff --git a/sdk/cosmos/azure-cosmos/samples/nonpartitioned_container_operations.py b/sdk/cosmos/azure-cosmos/samples/nonpartitioned_container_operations.py index bab2cb083e66..cdde98ec40bf 100644 --- a/sdk/cosmos/azure-cosmos/samples/nonpartitioned_container_operations.py +++ b/sdk/cosmos/azure-cosmos/samples/nonpartitioned_container_operations.py @@ -1,23 +1,8 @@ -#The MIT License (MIT) -#Copyright (c) 2018 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. +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions diff --git a/sdk/cosmos/azure-cosmos/samples/nonpartitioned_container_operations_async.py b/sdk/cosmos/azure-cosmos/samples/nonpartitioned_container_operations_async.py index 2756af6bf129..02ecbb7bf61e 100644 --- a/sdk/cosmos/azure-cosmos/samples/nonpartitioned_container_operations_async.py +++ b/sdk/cosmos/azure-cosmos/samples/nonpartitioned_container_operations_async.py @@ -1,3 +1,8 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- import azure.cosmos.aio.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions import requests