Query compatibility mode header (#3295)

.changes/next-release/enhancement-QueryCompatibilityModeHeader-5257.json

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

botocore/handlers.py

@@ -1287,6 +1287,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.
@@ -1348,6 +1354,7 @@ def _update_status_code(response, **kwargs):
     ('docs.response-params.s3.*.complete-section', document_expires_shape),
     ('before-endpoint-resolution.s3', customize_endpoint_resolver_builtins),
     ('before-call', add_recursion_detection_header),
+    ('before-call', add_query_compatibility_header),
     ('before-call.s3', add_expect_header),
     ('before-call.glacier', add_glacier_version),
     ('before-call.apigateway', add_accept_header),

botocore/model.py

@@ -478,6 +478,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 f'{self.__class__.__name__}({self.service_name})'
 

tests/functional/test_sqs.py

@@ -46,3 +46,11 @@ def test_query_compatible_error_parsing(self):
                 self.client.delete_queue(
                     QueueUrl="not-a-real-queue-botocore",
                 )
+
+    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/test_handlers.py

@@ -1944,3 +1944,20 @@ 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()
+
+
+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']