[Cosmos] Typing part 2 (#33341)

* More typing

* Literals

* Fix import

* Add typing extensions

* Body must be Dict

* Fixed some samples typing

* Multihash pk typehints

* Remove ignored option

* Review feedback

* More review feedback

* Fix tests

* Fix instance check

Author: annatisch

sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policy.py

@@ -3,22 +3,47 @@
 # Licensed under the MIT License. See LICENSE.txt in the project root for
 # license information.
 # -------------------------------------------------------------------------
-from typing import MutableMapping
+from typing import TypeVar, Any, MutableMapping
 
+from azure.core.pipeline import PipelineRequest
 from azure.core.pipeline.policies import BearerTokenCredentialPolicy
-from azure.cosmos import http_constants
+from azure.core.pipeline.transport import HttpRequest as LegacyHttpRequest
+from azure.core.rest import HttpRequest
+
+from .http_constants import HttpHeaders
+
+HTTPRequestType = TypeVar("HTTPRequestType", HttpRequest, LegacyHttpRequest)
 
 
 class CosmosBearerTokenCredentialPolicy(BearerTokenCredentialPolicy):
 
     @staticmethod
     def _update_headers(headers: MutableMapping[str, str], token: str) -> None:
         """Updates the Authorization header with the 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 MutableMapping[str, str] headers: The HTTP Request headers
         :param str token: The OAuth token.
         """
-        headers[http_constants.HttpHeaders.Authorization] = f"type=aad&ver=1.0&sig={token}"
+        headers[HttpHeaders.Authorization] = f"type=aad&ver=1.0&sig={token}"
+
+    def on_request(self, request: PipelineRequest[HTTPRequestType]) -> 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
+        """
+        super().on_request(request)
+        self._update_headers(request.http_request.headers, self._token.token)
+
+    def authorize_request(self, request: PipelineRequest[HTTPRequestType], *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
+        """
+        super().authorize_request(request, *scopes, **kwargs)
+        self._update_headers(request.http_request.headers, self._token.token)

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -26,6 +26,7 @@
 from email.utils import formatdate
 import json
 import uuid
+import re
 import binascii
 from typing import Dict, Any, List, Mapping, Optional, Sequence, Union, Tuple, TYPE_CHECKING
 
@@ -61,6 +62,13 @@
     'priority_level': 'priorityLevel'
 }
 
+# Cosmos resource ID validation regex breakdown:
+# ^ Match start of string.
+# [^/\#?]{0,255} Match any character that is not /\#? for between 0-255 characters.
+# $ End of string
+_VALID_COSMOS_RESOURCE = re.compile(r"^[^/\\#?\t\r\n]{0,255}$")
+
+
 def _get_match_headers(kwargs: Dict[str, Any]) -> Tuple[Optional[str], Optional[str]]:
     if_match = kwargs.pop('if_match', None)
     if_none_match = kwargs.pop('if_none_match', None)
@@ -684,8 +692,20 @@ def validate_cache_staleness_value(max_integrated_cache_staleness: Any) -> None:
                          "integer greater than or equal to zero")
 
 
+def _validate_resource(resource: Mapping[str, Any]) -> None:
+    id_: Optional[str] = resource.get("id")
+    if id_:
+        try:
+            if _VALID_COSMOS_RESOURCE.match(id_) is None:
+                raise ValueError("Id contains illegal chars.")
+            if id_[-1] in [" ", "\n"]:
+                raise ValueError("Id ends with a space or newline.")
+        except TypeError as e:
+            raise TypeError("Id type must be a string.") from e
+
+
 def _stringify_auto_scale(offer: ThroughputProperties) -> str:
-    auto_scale_params = None
+    auto_scale_params: Optional[Dict[str, Union[None, int, Dict[str, Any]]]] = None
     max_throughput = offer.auto_scale_max_throughput
     increment_percent = offer.auto_scale_increment_percent
     auto_scale_params = {"maxThroughput": max_throughput}

sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py

@@ -23,24 +23,28 @@
 """
 
 
-class _Constants(object):
+from typing import Dict
+from typing_extensions import Literal
+
+
+class _Constants:
     """Constants used in the azure-cosmos package"""
 
-    UserConsistencyPolicy = "userConsistencyPolicy"
-    DefaultConsistencyLevel = "defaultConsistencyLevel"
+    UserConsistencyPolicy: Literal["userConsistencyPolicy"] = "userConsistencyPolicy"
+    DefaultConsistencyLevel: Literal["defaultConsistencyLevel"] = "defaultConsistencyLevel"
 
     # GlobalDB related constants
-    WritableLocations = "writableLocations"
-    ReadableLocations = "readableLocations"
-    Name = "name"
-    DatabaseAccountEndpoint = "databaseAccountEndpoint"
-    DefaultUnavailableLocationExpirationTime = 5 * 60 * 1000
+    WritableLocations: Literal["writableLocations"] = "writableLocations"
+    ReadableLocations: Literal["readableLocations"] = "readableLocations"
+    Name: Literal["name"] = "name"
+    DatabaseAccountEndpoint: Literal["databaseAccountEndpoint"] = "databaseAccountEndpoint"
+    DefaultUnavailableLocationExpirationTime: int = 5 * 60 * 1000
 
     # ServiceDocument Resource
-    EnableMultipleWritableLocations = "enableMultipleWriteLocations"
+    EnableMultipleWritableLocations: Literal["enableMultipleWriteLocations"] = "enableMultipleWriteLocations"
 
     # Error code translations
-    ERROR_TRANSLATIONS = {
+    ERROR_TRANSLATIONS: Dict[int, str] = {
         400: "BAD_REQUEST - Request being sent is invalid.",
         401: "UNAUTHORIZED - The input authorization token can't serve the request.",
         403: "FORBIDDEN",