From 4d6fa0e5fa6a8a2f91c18ae48b0dcdc971749b46 Mon Sep 17 00:00:00 2001
From: Simon Moreno <30335873+simorenoh@users.noreply.github.com>
Date: Wed, 6 Apr 2022 11:12:16 -0400
Subject: [PATCH] [Cosmos] AAD authentication async client (#23717)

* working authentication to get database account

* working aad authentication for sync client with sample

* readme and changelog

* pylint and better comments on sample

* working async aad

* Delete access_cosmos_with_aad.py

snuck its way into the async PR

* Update _auth_policies.py

* small changes

* Update _cosmos_client_connection.py

* removing changes made in sync

* Update _auth_policy_async.py

* Update _auth_policy_async.py

* Update _auth_policy_async.py

* added licenses to samples
---
 sdk/cosmos/azure-cosmos/CHANGELOG.md          |   1 +
 .../azure/cosmos/_cosmos_client_connection.py |   2 -
 .../azure/cosmos/aio/_auth_policy_async.py    | 175 ++++++++++++++++++
 .../aio/_cosmos_client_connection_async.py    |  10 +
 .../samples/access_cosmos_with_aad.py         |   5 +
 .../samples/access_cosmos_with_aad_async.py   | 116 ++++++++++++
 .../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 +
 23 files changed, 392 insertions(+), 22 deletions(-)
 create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.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..260bf6ce75c8 100644
--- a/sdk/cosmos/azure-cosmos/CHANGELOG.md
+++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md
@@ -3,6 +3,7 @@
 ### 4.3.0b4 (Unreleased)
 
 #### Features Added
+- Added support for AAD authentication for the async client
 - Added support for AAD authentication for the sync client
 
 ### 4.3.0b3 (2022-03-10)
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 821ddc00adfc..ffafc2e397a4 100644
--- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py
+++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py
@@ -61,8 +61,6 @@
 from ._auth_policy import CosmosBearerTokenCredentialPolicy
 
 ClassType = TypeVar("ClassType")
-
-
 # pylint: disable=protected-access
 
 
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
new file mode 100644
index 000000000000..f5913f5cf56a
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_auth_policy_async.py
@@ -0,0 +1,175 @@
+# -------------------------------------------------------------------------
+# 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, Dict, Union
+from azure.core.pipeline.policies import AsyncHTTPPolicy
+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.
+
+    :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.
+        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.
+        """
+        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.
+    """
+
+    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/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py
index dab5b82b52a0..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,10 +58,12 @@
 from .. import _session
 from .. import _utils
 from ..partition_key import _Undefined, _Empty
+from ._auth_policy_async 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,12 +180,18 @@ 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(),
             retry_policy,
+            credentials_policy,
             CustomHookPolicy(**kwargs),
             NetworkTraceLoggingPolicy(**kwargs),
             DistributedTracingPolicy(**kwargs),
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
new file mode 100644
index 000000000000..4f6849a3f487
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos/samples/access_cosmos_with_aad_async.py
@@ -0,0 +1,116 @@
+# -------------------------------------------------------------------------
+# 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
+from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential
+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.
+# ----------------------------------------------------------------------------------------------------------
+# <configureConnectivity>
+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 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 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 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 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/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