[Tables] Cosmos workaround for Merge/Update (Azure#17888)

Closes Azure#13844
benbp · Apr 19, 2021 · f9cc68c · f9cc68c
1 parent 8920028
commit f9cc68c
Show file tree

Hide file tree

Showing 101 changed files with 10,598 additions and 5,847 deletions.
diff --git a/sdk/tables/azure-data-tables/CHANGELOG.md b/sdk/tables/azure-data-tables/CHANGELOG.md
@@ -1,7 +1,7 @@
 # Release History
 
 ## 12.0.0b7 (Unreleased)
-
+* Fixed issue with Cosmos merge operations
 
 ## 12.0.0b6 (2021-04-06)
 * Updated deserialization of datetime fields in entities to support preservation of the service format with additional decimal place.

diff --git a/sdk/tables/azure-data-tables/azure/data/tables/_base_client.py b/sdk/tables/azure-data-tables/azure/data/tables/_base_client.py
@@ -36,7 +36,7 @@
     AzureSasCredentialPolicy
 )
 
-from ._common_conversion import _to_utc_datetime
+from ._common_conversion import _to_utc_datetime, _is_cosmos_endpoint
 from ._shared_access_signature import QueryStringConstants
 from ._constants import (
     STORAGE_OAUTH_SCOPE,
@@ -47,6 +47,7 @@
 from ._models import LocationMode, BatchTransactionResult
 from ._authentication import SharedKeyCredentialPolicy
 from ._policies import (
+    CosmosPatchTransformPolicy,
     StorageHeadersPolicy,
     StorageContentValidation,
     StorageRequestHook,
@@ -80,7 +81,7 @@
 }
 
 
-class StorageAccountHostsMixin(object):
+class StorageAccountHostsMixin(object):  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
         parsed_url,  # type: Any
@@ -92,6 +93,7 @@ def __init__(
         self._location_mode = kwargs.get("_location_mode", LocationMode.PRIMARY)
         self._hosts = kwargs.get("_hosts")
         self.scheme = parsed_url.scheme
+        self._cosmos_endpoint = _is_cosmos_endpoint(parsed_url.hostname)
 
         if service not in ["blob", "queue", "file-share", "dfs", "table"]:
             raise ValueError("Invalid service: {}".format(service))
@@ -149,6 +151,9 @@ def __init__(
             HttpLoggingPolicy(**kwargs),
         ]
 
+        if self._cosmos_endpoint:
+            self._policies.insert(0, CosmosPatchTransformPolicy())
+
     def __enter__(self):
         self._client.__enter__()
         return self

diff --git a/sdk/tables/azure-data-tables/azure/data/tables/_common_conversion.py b/sdk/tables/azure-data-tables/azure/data/tables/_common_conversion.py
@@ -86,3 +86,18 @@ def _sign_string(key, string_to_sign, key_is_base64=True):
     digest = signed_hmac_sha256.digest()
     encoded_digest = _encode_base64(digest)
     return encoded_digest
+
+
+def _is_cosmos_endpoint(url):
+    if ".table.cosmodb." in url:
+        return True
+
+    if ".table.cosmos." in url:
+        return True
+
+    return False
+
+
+def _transform_patch_to_cosmos_post(request):
+    request.method = "POST"
+    request.headers["X-HTTP-Method"] = "MERGE"
diff --git a/sdk/tables/azure-data-tables/azure/data/tables/_policies.py b/sdk/tables/azure-data-tables/azure/data/tables/_policies.py
@@ -41,6 +41,7 @@
 )
 from azure.core.exceptions import AzureError, ServiceRequestError, ServiceResponseError
 
+from ._common_conversion import _transform_patch_to_cosmos_post
 from ._models import LocationMode
 
 try:
@@ -750,3 +751,12 @@ def get_backoff_time(self, settings):
         )
         random_range_end = self.backoff + self.random_jitter_range
         return random_generator.uniform(random_range_start, random_range_end)
+
+
+class CosmosPatchTransformPolicy(SansIOHTTPPolicy):
+    """Policy to transform PATCH requests into POST requests with the "X-HTTP-Method":"MERGE" header set."""
+
+    def on_request(self, request):
+        # type: (PipelineRequest) -> Union[None, Awaitable[None]]
+        if request.http_request.method == "PATCH":
+            _transform_patch_to_cosmos_post(request.http_request)
diff --git a/sdk/tables/azure-data-tables/azure/data/tables/_table_batch.py b/sdk/tables/azure-data-tables/azure/data/tables/_table_batch.py
@@ -9,6 +9,7 @@
     Dict,
 )
 
+from ._common_conversion import _is_cosmos_endpoint, _transform_patch_to_cosmos_post
 from ._models import UpdateMode
 from ._serialize import _get_match_headers, _add_entity_properties
 
@@ -467,6 +468,8 @@ def _batch_merge_entity(
         request = self._client._client.patch(  # pylint: disable=protected-access
             url, query_parameters, header_parameters, **body_content_kwargs
         )
+        if _is_cosmos_endpoint(url):
+            _transform_patch_to_cosmos_post(request)
         self._requests.append(request)
 
     _batch_merge_entity.metadata = {

diff --git a/sdk/tables/azure-data-tables/azure/data/tables/aio/_base_client_async.py b/sdk/tables/azure-data-tables/azure/data/tables/aio/_base_client_async.py
@@ -40,6 +40,7 @@
 from .._generated.aio._configuration import AzureTableConfiguration
 from .._models import BatchErrorException, BatchTransactionResult
 from .._policies import (
+    CosmosPatchTransformPolicy,
     StorageContentValidation,
     StorageRequestHook,
     StorageHosts,
@@ -106,7 +107,6 @@ def _configure_policies(self, **kwargs):
 
         kwargs.setdefault("connection_timeout", CONNECTION_TIMEOUT)
         kwargs.setdefault("read_timeout", READ_TIMEOUT)
-
         self._policies = [
             StorageHeadersPolicy(**kwargs),
             ProxyPolicy(**kwargs),
@@ -124,6 +124,9 @@ def _configure_policies(self, **kwargs):
             HttpLoggingPolicy(**kwargs),
         ]
 
+        if self._cosmos_endpoint:
+            self._policies.insert(0, CosmosPatchTransformPolicy())
+
     async def _batch_send(
         self,
         entities,  # type: List[TableEntity]

diff --git a/sdk/tables/azure-data-tables/azure/data/tables/aio/_table_batch_async.py b/sdk/tables/azure-data-tables/azure/data/tables/aio/_table_batch_async.py
@@ -8,7 +8,7 @@
 
 from azure.core.pipeline import PipelineResponse
 
-
+from .._common_conversion import _is_cosmos_endpoint, _transform_patch_to_cosmos_post
 from .._models import UpdateMode
 from .._serialize import (
     _get_match_headers,
@@ -458,6 +458,8 @@ def _batch_merge_entity(
         request = self._client._client.patch(  # pylint: disable=protected-access
             url, query_parameters, header_parameters, **body_content_kwargs
         )
+        if _is_cosmos_endpoint(url):
+            _transform_patch_to_cosmos_post(request)
         self._requests.append(request)
 
     _batch_merge_entity.metadata = {

diff --git a/sdk/tables/azure-data-tables/azure/data/tables/aio/_table_client_async.py b/sdk/tables/azure-data-tables/azure/data/tables/aio/_table_client_async.py
@@ -22,6 +22,7 @@
 from azure.core.tracing.decorator_async import distributed_trace_async
 
 from .._base_client import parse_connection_str
+from .._common_conversion import _is_cosmos_endpoint
 from .._constants import CONNECTION_TIMEOUT
 from .._entity import TableEntity
 from .._generated.aio import AzureTable
@@ -71,6 +72,9 @@ def __init__(
             **kwargs
         )
         loop = kwargs.pop("loop", None)
+
+        self._cosmos_endpoint = _is_cosmos_endpoint(account_url)
+
         super(TableClient, self).__init__(
             account_url,
             table_name=table_name,

diff --git a/.../azure-data-tables/tests/recordings/test_retry.test_retry_callback_and_retry_context.yaml b/.../azure-data-tables/tests/recordings/test_retry.test_retry_callback_and_retry_context.yaml
@@ -15,11 +15,11 @@ interactions:
       DataServiceVersion:
       - '3.0'
       Date:
-      - Mon, 05 Apr 2021 20:52:56 GMT
+      - Wed, 07 Apr 2021 22:25:39 GMT
       User-Agent:
-      - azsdk-python-data-tables/12.0.0b6 Python/3.7.4 (Windows-10-10.0.19041-SP0)
+      - azsdk-python-data-tables/12.0.0b7 Python/3.9.0rc1 (Windows-10-10.0.19041-SP0)
       x-ms-date:
-      - Mon, 05 Apr 2021 20:52:56 GMT
+      - Wed, 07 Apr 2021 22:25:39 GMT
       x-ms-version:
       - '2019-02-02'
     method: POST
@@ -33,7 +33,7 @@ interactions:
       content-type:
       - application/json;odata=minimalmetadata;streaming=true;charset=utf-8
       date:
-      - Mon, 05 Apr 2021 20:52:59 GMT
+      - Wed, 07 Apr 2021 22:25:39 GMT
       location:
       - https://fake_table_account.table.core.windows.net/Tables('uttablee75e13f0')
       server:
@@ -63,26 +63,26 @@ interactions:
       DataServiceVersion:
       - '3.0'
       Date:
-      - Mon, 05 Apr 2021 20:52:56 GMT
+      - Wed, 07 Apr 2021 22:25:39 GMT
       User-Agent:
-      - azsdk-python-data-tables/12.0.0b6 Python/3.7.4 (Windows-10-10.0.19041-SP0)
+      - azsdk-python-data-tables/12.0.0b7 Python/3.9.0rc1 (Windows-10-10.0.19041-SP0)
       x-ms-date:
-      - Mon, 05 Apr 2021 20:52:56 GMT
+      - Wed, 07 Apr 2021 22:25:39 GMT
       x-ms-version:
       - '2019-02-02'
     method: POST
     uri: https://fake_table_account.table.core.windows.net/Tables
   response:
     body:
       string: '{"odata.error":{"code":"TableAlreadyExists","message":{"lang":"en-US","value":"The
-        table specified already exists.\nRequestId:3a9bd1f0-0002-000c-7e5d-2a4d2c000000\nTime:2021-04-05T20:53:02.8814645Z"}}}'
+        table specified already exists.\nRequestId:4faf3acc-6002-00a3-01fc-2bbfe1000000\nTime:2021-04-07T22:25:41.7333309Z"}}}'
     headers:
       cache-control:
       - no-cache
       content-type:
       - application/json;odata=minimalmetadata;streaming=true;charset=utf-8
       date:
-      - Mon, 05 Apr 2021 20:53:02 GMT
+      - Wed, 07 Apr 2021 22:25:40 GMT
       server:
       - Windows-Azure-Table/1.0 Microsoft-HTTPAPI/2.0
       transfer-encoding:
@@ -106,11 +106,11 @@ interactions:
       Content-Length:
       - '0'
       Date:
-      - Mon, 05 Apr 2021 20:53:03 GMT
+      - Wed, 07 Apr 2021 22:25:41 GMT
       User-Agent:
-      - azsdk-python-data-tables/12.0.0b6 Python/3.7.4 (Windows-10-10.0.19041-SP0)
+      - azsdk-python-data-tables/12.0.0b7 Python/3.9.0rc1 (Windows-10-10.0.19041-SP0)
       x-ms-date:
-      - Mon, 05 Apr 2021 20:53:03 GMT
+      - Wed, 07 Apr 2021 22:25:41 GMT
       x-ms-version:
       - '2019-02-02'
     method: DELETE
@@ -124,7 +124,7 @@ interactions:
       content-length:
       - '0'
       date:
-      - Mon, 05 Apr 2021 20:53:03 GMT
+      - Wed, 07 Apr 2021 22:25:41 GMT
       server:
       - Windows-Azure-Table/1.0 Microsoft-HTTPAPI/2.0
       x-content-type-options:
@@ -146,26 +146,26 @@ interactions:
       Content-Length:
       - '0'
       Date:
-      - Mon, 05 Apr 2021 20:53:03 GMT
+      - Wed, 07 Apr 2021 22:25:41 GMT
       User-Agent:
-      - azsdk-python-data-tables/12.0.0b6 Python/3.7.4 (Windows-10-10.0.19041-SP0)
+      - azsdk-python-data-tables/12.0.0b7 Python/3.9.0rc1 (Windows-10-10.0.19041-SP0)
       x-ms-date:
-      - Mon, 05 Apr 2021 20:53:03 GMT
+      - Wed, 07 Apr 2021 22:25:41 GMT
       x-ms-version:
       - '2019-02-02'
     method: DELETE
     uri: https://fake_table_account.table.core.windows.net/Tables('uttablee75e13f0')
   response:
     body:
       string: '{"odata.error":{"code":"ResourceNotFound","message":{"lang":"en-US","value":"The
-        specified resource does not exist.\nRequestId:3a9bd24c-0002-000c-4e5d-2a4d2c000000\nTime:2021-04-05T20:53:03.7921084Z"}}}'
+        specified resource does not exist.\nRequestId:4faf3af7-6002-00a3-26fc-2bbfe1000000\nTime:2021-04-07T22:25:42.0355353Z"}}}'
     headers:
       cache-control:
       - no-cache
       content-type:
       - application/json;odata=minimalmetadata;streaming=true;charset=utf-8
       date:
-      - Mon, 05 Apr 2021 20:53:03 GMT
+      - Wed, 07 Apr 2021 22:25:41 GMT
       server:
       - Windows-Azure-Table/1.0 Microsoft-HTTPAPI/2.0
       transfer-encoding:

diff --git a/sdk/tables/azure-data-tables/tests/recordings/test_retry.test_retry_on_server_error.yaml b/sdk/tables/azure-data-tables/tests/recordings/test_retry.test_retry_on_server_error.yaml
@@ -15,11 +15,11 @@ interactions:
       DataServiceVersion:
       - '3.0'
       Date:
-      - Mon, 05 Apr 2021 20:53:04 GMT
+      - Wed, 07 Apr 2021 22:32:50 GMT
       User-Agent:
-      - azsdk-python-data-tables/12.0.0b6 Python/3.7.4 (Windows-10-10.0.19041-SP0)
+      - azsdk-python-data-tables/12.0.0b7 Python/3.9.0rc1 (Windows-10-10.0.19041-SP0)
       x-ms-date:
-      - Mon, 05 Apr 2021 20:53:04 GMT
+      - Wed, 07 Apr 2021 22:32:50 GMT
       x-ms-version:
       - '2019-02-02'
     method: POST
@@ -33,7 +33,7 @@ interactions:
       content-type:
       - application/json;odata=minimalmetadata;streaming=true;charset=utf-8
       date:
-      - Mon, 05 Apr 2021 20:53:04 GMT
+      - Wed, 07 Apr 2021 22:32:51 GMT
       location:
       - https://fake_table_account.table.core.windows.net/Tables('uttable270d0f94')
       server:
@@ -63,26 +63,26 @@ interactions:
       DataServiceVersion:
       - '3.0'
       Date:
-      - Mon, 05 Apr 2021 20:53:04 GMT
+      - Wed, 07 Apr 2021 22:32:50 GMT
       User-Agent:
-      - azsdk-python-data-tables/12.0.0b6 Python/3.7.4 (Windows-10-10.0.19041-SP0)
+      - azsdk-python-data-tables/12.0.0b7 Python/3.9.0rc1 (Windows-10-10.0.19041-SP0)
       x-ms-date:
-      - Mon, 05 Apr 2021 20:53:04 GMT
+      - Wed, 07 Apr 2021 22:32:50 GMT
       x-ms-version:
       - '2019-02-02'
     method: POST
     uri: https://fake_table_account.table.core.windows.net/Tables
   response:
     body:
       string: '{"odata.error":{"code":"TableAlreadyExists","message":{"lang":"en-US","value":"The
-        table specified already exists.\nRequestId:71db8f83-4002-0032-1b5d-2ada53000000\nTime:2021-04-05T20:53:22.8853531Z"}}}'
+        table specified already exists.\nRequestId:0451134d-a002-0077-1dfd-2b0fb0000000\nTime:2021-04-07T22:33:07.4938348Z"}}}'
     headers:
       cache-control:
       - no-cache
       content-type:
       - application/json;odata=minimalmetadata;streaming=true;charset=utf-8
       date:
-      - Mon, 05 Apr 2021 20:53:22 GMT
+      - Wed, 07 Apr 2021 22:33:07 GMT
       server:
       - Windows-Azure-Table/1.0 Microsoft-HTTPAPI/2.0
       transfer-encoding:
@@ -106,11 +106,11 @@ interactions:
       Content-Length:
       - '0'
       Date:
-      - Mon, 05 Apr 2021 20:53:23 GMT
+      - Wed, 07 Apr 2021 22:33:07 GMT
       User-Agent:
-      - azsdk-python-data-tables/12.0.0b6 Python/3.7.4 (Windows-10-10.0.19041-SP0)
+      - azsdk-python-data-tables/12.0.0b7 Python/3.9.0rc1 (Windows-10-10.0.19041-SP0)
       x-ms-date:
-      - Mon, 05 Apr 2021 20:53:23 GMT
+      - Wed, 07 Apr 2021 22:33:07 GMT
       x-ms-version:
       - '2019-02-02'
     method: DELETE
@@ -124,7 +124,7 @@ interactions:
       content-length:
       - '0'
       date:
-      - Mon, 05 Apr 2021 20:53:22 GMT
+      - Wed, 07 Apr 2021 22:33:07 GMT
       server:
       - Windows-Azure-Table/1.0 Microsoft-HTTPAPI/2.0
       x-content-type-options:
@@ -146,26 +146,26 @@ interactions:
       Content-Length:
       - '0'
       Date:
-      - Mon, 05 Apr 2021 20:53:23 GMT
+      - Wed, 07 Apr 2021 22:33:07 GMT
       User-Agent:
-      - azsdk-python-data-tables/12.0.0b6 Python/3.7.4 (Windows-10-10.0.19041-SP0)
+      - azsdk-python-data-tables/12.0.0b7 Python/3.9.0rc1 (Windows-10-10.0.19041-SP0)
       x-ms-date:
-      - Mon, 05 Apr 2021 20:53:23 GMT
+      - Wed, 07 Apr 2021 22:33:07 GMT
       x-ms-version:
       - '2019-02-02'
     method: DELETE
     uri: https://fake_table_account.table.core.windows.net/Tables('uttable270d0f94')
   response:
     body:
       string: '{"odata.error":{"code":"ResourceNotFound","message":{"lang":"en-US","value":"The
-        specified resource does not exist.\nRequestId:71db9005-4002-0032-195d-2ada53000000\nTime:2021-04-05T20:53:23.8910718Z"}}}'
+        specified resource does not exist.\nRequestId:0451138c-a002-0077-5afd-2b0fb0000000\nTime:2021-04-07T22:33:07.7860414Z"}}}'
     headers:
       cache-control:
       - no-cache
       content-type:
       - application/json;odata=minimalmetadata;streaming=true;charset=utf-8
       date:
-      - Mon, 05 Apr 2021 20:53:23 GMT
+      - Wed, 07 Apr 2021 22:33:07 GMT
       server:
       - Windows-Azure-Table/1.0 Microsoft-HTTPAPI/2.0
       transfer-encoding: