Track feature IDs for retry modes, CBOR, and request compression (#9555)

Author: ashovlin

awscli/botocore/client.py

@@ -51,7 +51,7 @@
 from botocore.model import ServiceModel
 from botocore.paginate import Paginator
 from botocore.retries import adaptive, standard
-from botocore.useragent import UserAgentString
+from botocore.useragent import UserAgentString, register_feature_id
 from botocore.utils import (
     CachedProperty,
     EventbridgeSignerSetter,
@@ -228,6 +228,9 @@ def _register_retries(self, client):
         elif retry_mode == 'adaptive':
             self._register_v2_standard_retries(client)
             self._register_v2_adaptive_retries(client)
+        else:
+            return
+        register_feature_id(f'RETRY_MODE_{retry_mode.upper()}')
 
     def _register_v2_standard_retries(self, client):
         max_attempts = client.meta.config.retries.get('max_attempts')

awscli/botocore/compress.py

@@ -22,6 +22,7 @@
 from gzip import compress as gzip_compress
 
 from botocore.compat import urlencode
+from botocore.useragent import register_feature_id
 from botocore.utils import determine_content_length
 
 logger = logging.getLogger(__name__)
@@ -88,6 +89,7 @@ def _get_body_size(body):
 
 
 def _gzip_compress_body(body):
+    register_feature_id('GZIP_REQUEST_COMPRESSION')
     if isinstance(body, str):
         return gzip_compress(body.encode('utf-8'))
     elif isinstance(body, (bytes, bytearray)):

awscli/botocore/serialize.py

@@ -49,6 +49,7 @@
 
 from botocore import validate
 from botocore.compat import formatdate, json
+from botocore.useragent import register_feature_id
 from botocore.utils import (
     has_header,
     is_json_value_header,
@@ -1162,6 +1163,10 @@ def _add_xml_namespace(self, shape, structure_node):
 class RpcV2CBORSerializer(BaseRpcV2Serializer, CBORSerializer):
     TIMESTAMP_FORMAT = 'unixtimestamp'
 
+    def serialize_to_request(self, parameters, operation_model):
+        register_feature_id('PROTOCOL_RPC_V2_CBOR')
+        return super().serialize_to_request(parameters, operation_model)
+
     def _serialize_body_params(self, parameters, input_shape):
         body = bytearray()
         self._serialize_data_item(body, parameters, input_shape)

awscli/botocore/useragent.py

@@ -56,7 +56,11 @@
 _USERAGENT_FEATURE_MAPPINGS = {
     'WAITER': 'B',
     'PAGINATOR': 'C',
+    'RETRY_MODE_STANDARD': 'E',
+    'RETRY_MODE_ADAPTIVE': 'F',
     'S3_TRANSFER': 'G',
+    'PROTOCOL_RPC_V2_CBOR': 'M',
+    'GZIP_REQUEST_COMPRESSION': 'L',
     'ENDPOINT_OVERRIDE': 'N',
     'SIGV4A_SIGNING': 'S',
 }

tests/functional/botocore/test_compress.py

@@ -18,6 +18,10 @@
 from botocore.config import Config
 
 from tests import ALL_SERVICES, ClientHTTPStubber, patch_load_service_model
+from tests.functional.botocore.test_useragent import (
+    get_captured_ua_strings,
+    parse_registered_feature_ids,
+)
 
 FAKE_MODEL = {
     "version": "2.0",
@@ -127,3 +131,24 @@ def test_compression(patched_session, monkeypatch):
         serialized_body = f"{additional_params}&{serialized_params}"
         actual_body = gzip.decompress(http_stubber.requests[0].body)
         assert serialized_body.encode('utf-8') == actual_body
+
+
+def test_user_agent_has_gzip_feature_id(patched_session, monkeypatch):
+    patch_load_service_model(
+        patched_session, monkeypatch, FAKE_MODEL, FAKE_RULESET
+    )
+    client = patched_session.create_client(
+        "otherservice",
+        region_name="us-west-2",
+        config=Config(request_min_compression_size_bytes=100),
+    )
+    with ClientHTTPStubber(client, strict=True) as http_stubber:
+        http_stubber.add_response(status=200, body=b"<response/>")
+        params_list = [
+            {"MockOpParam": f"MockOpParamValue{i}"} for i in range(1, 21)
+        ]
+        # The mock operation registers `'GZIP_REQUEST_COMPRESSION': 'L'`
+        client.mock_operation(MockOpParamList=params_list)
+    ua_string = get_captured_ua_strings(http_stubber)[0]
+    feature_list = parse_registered_feature_ids(ua_string)
+    assert 'L' in feature_list

tests/functional/botocore/test_retry.py

@@ -17,6 +17,10 @@
 from botocore.exceptions import ClientError
 
 from tests import BaseSessionTest, ClientHTTPStubber, mock
+from tests.functional.botocore.test_useragent import (
+    get_captured_ua_strings,
+    parse_registered_feature_ids,
+)
 
 
 class TestRetry(BaseSessionTest):
@@ -46,6 +50,18 @@ def assert_will_retry_n_times(
                 yield
             self.assertEqual(len(http_stubber.requests), num_responses)
 
+    def _get_feature_id_lists_from_retries(self, client):
+        with ClientHTTPStubber(client) as http_stubber:
+            # Add two failed responses followed by a success
+            http_stubber.add_response(status=502, body=b'{}')
+            http_stubber.add_response(status=502, body=b'{}')
+            http_stubber.add_response(status=200, body=b'{}')
+            client.list_tables()
+        ua_strings = get_captured_ua_strings(http_stubber)
+        return [
+            parse_registered_feature_ids(ua_string) for ua_string in ua_strings
+        ]
+
     def create_client_with_retry_config(
         self, service, retry_mode='standard', max_attempts=None
     ):
@@ -161,3 +177,19 @@ def test_can_exhaust_default_retry_quota(self):
         self.assertTrue(
             e.exception.response['ResponseMetadata'].get('RetryQuotaReached')
         )
+
+    def test_user_agent_has_standard_mode_feature_id(self):
+        client = self.create_client_with_retry_config(
+            'dynamodb', retry_mode='standard'
+        )
+        feature_lists = self._get_feature_id_lists_from_retries(client)
+        # Confirm all requests register `'RETRY_MODE_STANDARD': 'E'`
+        assert all('E' in feature_list for feature_list in feature_lists)
+
+    def test_user_agent_has_adaptive_mode_feature_id(self):
+        client = self.create_client_with_retry_config(
+            'dynamodb', retry_mode='adaptive'
+        )
+        feature_lists = self._get_feature_id_lists_from_retries(client)
+        # Confirm all requests register `'RETRY_MODE_ADAPTIVE': 'F'`
+        assert all('F' in feature_list for feature_list in feature_lists)

tests/functional/botocore/test_rpc_v2_cbor.py

@@ -0,0 +1,90 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+from tests import ClientHTTPStubber, patch_load_service_model
+from tests.functional.botocore.test_useragent import (
+    get_captured_ua_strings,
+    parse_registered_feature_ids,
+)
+
+MOCK_SERVICE_MODEL = {
+    "version": "1.0",
+    "documentation": "",
+    "metadata": {
+        "apiVersion": "2020-02-02",
+        "endpointPrefix": "mockservice",
+        "protocols": ["smithy-rpc-v2-cbor"],
+        "protocol": "smithy-rpc-v2-cbor",
+        "serviceFullName": "Mock Service",
+        "serviceId": "mockservice",
+        "signatureVersion": "v4",
+        "signingName": "mockservice",
+        "targetPrefix": "mockservice",
+        "uid": "mockservice-2020-02-02",
+    },
+    "operations": {
+        "MockOperation": {
+            "name": "MockOperation",
+            "http": {"method": "GET", "requestUri": "/"},
+            "input": {"shape": "MockOperationRequest"},
+            "documentation": "",
+        },
+    },
+    "shapes": {
+        "MockOpParam": {
+            "type": "string",
+        },
+        "MockOperationRequest": {
+            "type": "structure",
+            "required": ["MockOpParam"],
+            "members": {
+                "MockOpParam": {
+                    "shape": "MockOpParam",
+                    "documentation": "",
+                    "location": "uri",
+                    "locationName": "param",
+                },
+            },
+        },
+    },
+}
+
+MOCK_RULESET = {
+    "version": "1.0",
+    "parameters": {},
+    "rules": [
+        {
+            "conditions": [],
+            "endpoint": {
+                "url": "https://mockservice.us-west-2.amazonaws.com/"
+            },
+            "type": "endpoint",
+        },
+    ],
+}
+
+
+def test_user_agent_has_cbor_feature_id(patched_session, monkeypatch):
+    patch_load_service_model(
+        patched_session, monkeypatch, MOCK_SERVICE_MODEL, MOCK_RULESET
+    )
+    client = patched_session.create_client(
+        'mockservice', region_name='us-west-2'
+    )
+    with ClientHTTPStubber(client) as stub_client:
+        stub_client.add_response()
+        # The mock CBOR operation registers `'PROTOCOL_RPC_V2_CBOR': 'M'`
+        client.mock_operation(MockOpParam='mock-op-param-value')
+    ua_string = get_captured_ua_strings(stub_client)[0]
+    feature_list = parse_registered_feature_ids(ua_string)
+    assert 'M' in feature_list

tests/functional/botocore/test_useragent.py

@@ -22,6 +22,25 @@
 from tests import ClientHTTPStubber
 
 
+def get_captured_ua_strings(stubber):
+    """Get captured request-level user agent strings from stubber.
+
+    :type stubber: tests.BaseHTTPStubber
+    """
+    return [req.headers['User-Agent'].decode() for req in stubber.requests]
+
+
+def parse_registered_feature_ids(ua_string):
+    """Parse registered feature ids in user agent string.
+
+    :type ua_string: str
+    :rtype: list[str]
+    """
+    ua_fields = ua_string.split(' ')
+    feature_field = [field for field in ua_fields if field.startswith('m/')][0]
+    return feature_field[2:].split(',')
+
+
 class UACapHTTPStubber(ClientHTTPStubber):
     """
     Wrapper for ClientHTTPStubber that captures UA header from one request.

tests/functional/test_useragent.py

@@ -1,5 +1,8 @@
 from awscli import __version__ as awscli_version
 from tests import CLIRunner
+from tests.functional.botocore.test_useragent import (
+    parse_registered_feature_ids,
+)
 
 
 def assert_expected_user_agent(result, service, operation):
@@ -32,4 +35,5 @@ def test_user_agent_for_customization():
     result = cli_runner.run([service, operation])
     assert_expected_user_agent(result, service, operation)
     ua_string = result.aws_requests[0].http_requests[0].headers['User-Agent']
-    assert 'm/C' in ua_string
+    feature_list = parse_registered_feature_ids(ua_string)
+    assert 'C' in feature_list