Add query compatibility mode header (#9083)

.changes/next-release/enhancement-protocol-90080.json

@@ -0,0 +1,5 @@
+{
+  "type": "enhancement",
+  "category": "protocol",
+  "description": "Added support for header enabling service migration off the AWS Query protocol."
+}

awscli/botocore/handlers.py

@@ -1151,6 +1151,12 @@ def _update_status_code(response, **kwargs):
         http_response.status_code = parsed_status_code
 
 
+def add_query_compatibility_header(model, params, **kwargs):
+    if not model.service_model.is_query_compatible:
+        return
+    params['headers']['x-amzn-query-mode'] = 'true'
+
+
 # This is a list of (event_name, handler).
 # When a Session is created, everything in this list will be
 # automatically registered with that Session.
@@ -1201,7 +1207,7 @@ def _update_status_code(response, **kwargs):
     ('docs.response-example.s3.*.complete-section', document_expires_shape),
     ('docs.response-params.s3.*.complete-section', document_expires_shape),
     ('before-endpoint-resolution.s3', customize_endpoint_resolver_builtins),
-
+    ('before-call', add_query_compatibility_header),
     ('before-call.s3', add_expect_header),
     ('before-call.glacier', add_glacier_version),
     ('before-call.api-gateway', add_accept_header),

awscli/botocore/model.py

@@ -446,6 +446,10 @@ def signature_version(self):
     def signature_version(self, value):
         self._signature_version = value
 
+    @CachedProperty
+    def is_query_compatible(self):
+        return 'awsQueryCompatible' in self.metadata
+
     def __repr__(self):
         return '%s(%s)' % (self.__class__.__name__, self.service_name)
 

tests/functional/botocore/test_sqs.py

@@ -0,0 +1,32 @@
+# Copyright 2023 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 BaseSessionTest, ClientHTTPStubber
+
+
+class BaseSQSOperationTest(BaseSessionTest):
+    def setUp(self):
+        super().setUp()
+        self.region = "us-west-2"
+        self.client = self.session.create_client("sqs", self.region)
+        self.http_stubber = ClientHTTPStubber(self.client)
+
+
+class SQSQueryCompatibleTest(BaseSQSOperationTest):
+    def test_query_compatibility_mode_header_sent(self):
+        with self.http_stubber as stub:
+            stub.add_response()
+            self.client.delete_queue(QueueUrl="not-a-real-queue-botocore")
+            request = self.http_stubber.requests[0]
+            assert 'x-amzn-query-mode' in request.headers
+            assert request.headers['x-amzn-query-mode'] == b'true'

tests/unit/botocore/test_handlers.py

@@ -1698,4 +1698,21 @@ def test_document_response_params_without_expires(document_expires_mocks):
     )
     mocks['section'].get_section.assert_not_called()
     mocks['param_section'].add_new_section.assert_not_called()
-    mocks['doc_section'].write.assert_not_called()
+    mocks['doc_section'].write.assert_not_called()
+
+
+def test_add_query_compatibility_header():
+    service_model = ServiceModel({'metadata': {'awsQueryCompatible': {}}})
+    operation_model = OperationModel(mock.Mock(), service_model)
+    request_dict = {'headers': {}}
+    handlers.add_query_compatibility_header(operation_model, request_dict)
+    assert 'x-amzn-query-mode' in request_dict['headers']
+    assert request_dict['headers']['x-amzn-query-mode'] == 'true'
+
+
+def test_does_not_add_query_compatibility_header():
+    service_model = ServiceModel({'metadata': {}})
+    operation_model = OperationModel(mock.Mock(), service_model)
+    request_dict = {'headers': {}}
+    handlers.add_query_compatibility_header(operation_model, request_dict)
+    assert 'x-amzn-query-mode' not in request_dict['headers']