Improve awsQueryCompatible support for JSON and CBOR (#3128)

Author: mullermp

gems/aws-sdk-core/CHANGELOG.md

@@ -1,6 +1,8 @@
 Unreleased Changes
 ------------------
 
+* Feature - Support functionality for services that migrate from AWS Query to AWS JSON or CBOR.
+
 * Issue - Fix RPCv2 protocol to always send an Accept header for CBOR.
 
 3.210.0 (2024-10-18)

gems/aws-sdk-core/lib/aws-sdk-core/client_stubs.rb

@@ -307,14 +307,14 @@ def data_to_http_resp(operation_name, data)
 
     def protocol_helper
       case config.api.metadata['protocol']
-      when 'json'        then Stubbing::Protocols::Json
-      when 'rest-json'   then Stubbing::Protocols::RestJson
-      when 'rest-xml'    then Stubbing::Protocols::RestXml
-      when 'query'       then Stubbing::Protocols::Query
-      when 'ec2'         then Stubbing::Protocols::EC2
+      when 'json'               then Stubbing::Protocols::Json
+      when 'rest-json'          then Stubbing::Protocols::RestJson
+      when 'rest-xml'           then Stubbing::Protocols::RestXml
+      when 'query'              then Stubbing::Protocols::Query
+      when 'ec2'                then Stubbing::Protocols::EC2
       when 'smithy-rpc-v2-cbor' then Stubbing::Protocols::RpcV2
-      when 'api-gateway' then Stubbing::Protocols::ApiGateway
-      else raise "unsupported protocol"
+      when 'api-gateway'        then Stubbing::Protocols::ApiGateway
+      else raise 'unsupported protocol'
       end.new
     end
   end

gems/aws-sdk-core/lib/aws-sdk-core/json/error_handler.rb

@@ -26,7 +26,8 @@ def extract_error(body, context)
       def error_code(json, context)
         code =
           if aws_query_error?(context)
-            error = context.http_response.headers['x-amzn-query-error'].split(';')[0]
+            query_header = context.http_response.headers['x-amzn-query-error']
+            error, _type = query_header.split(';') # type not supported
             remove_prefix(error, context)
           else
             json['__type']

gems/aws-sdk-core/lib/aws-sdk-core/json/handler.rb

@@ -21,6 +21,7 @@ def build_request(context)
         context.http_request.http_method = 'POST'
         context.http_request.headers['Content-Type'] = content_type(context)
         context.http_request.headers['X-Amz-Target'] = target(context)
+        context.http_request.headers['X-Amzn-Query-Mode'] = 'true' if query_compatible?(context)
         context.http_request.body = build_body(context)
       end
 

gems/aws-sdk-core/lib/aws-sdk-core/rpc_v2/error_handler.rb

@@ -39,7 +39,8 @@ def extract_error(body, context)
       def error_code(data, context)
         code =
           if aws_query_error?(context)
-            error = context.http_response.headers['x-amzn-query-error'].split(';')[0]
+            query_header = context.http_response.headers['x-amzn-query-error']
+            error, _type = query_header.split(';') # type not supported
             remove_prefix(error, context)
           else
             data['__type']

gems/aws-sdk-core/lib/aws-sdk-core/rpc_v2/handler.rb

@@ -20,7 +20,8 @@ def with_metric(&block)
       end
 
       def build_request(context)
-        context.http_request.headers['smithy-protocol'] = 'rpc-v2-cbor'
+        context.http_request.headers['Smithy-Protocol'] = 'rpc-v2-cbor'
+        context.http_request.headers['X-Amzn-Query-Mode'] = 'true' if query_compatible?(context)
         context.http_request.http_method = 'POST'
         context.http_request.body = build_body(context)
         build_url(context)

gems/aws-sdk-core/lib/aws-sdk-core/stubbing.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module Aws
-  # setup autoloading for Stubbing module
+  # @api private
   module Stubbing
     autoload :EmptyStub, 'aws-sdk-core/stubbing/empty_stub'
     autoload :DataApplicator, 'aws-sdk-core/stubbing/data_applicator'
@@ -19,4 +19,4 @@ module Protocols
       autoload :ApiGateway, 'aws-sdk-core/stubbing/protocols/api_gateway'
     end
   end
-end
+end

gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/ec2.rb

@@ -3,6 +3,7 @@
 module Aws
   module Stubbing
     module Protocols
+      # @api private
       class EC2
 
         def stub_data(api, operation, data)
@@ -16,17 +17,17 @@ def stub_data(api, operation, data)
         end
 
         def stub_error(error_code)
-          http_resp = Seahorse::Client::Http::Response.new
-          http_resp.status_code = 400
-          http_resp.body = <<-XML.strip
-<ErrorResponse>
-  <Error>
-    <Code>#{error_code}</Code>
-    <Message>stubbed-response-error-message</Message>
-  </Error>
-</ErrorResponse>
+          resp = Seahorse::Client::Http::Response.new
+          resp.status_code = 400
+          resp.body = <<~XML.strip
+            <ErrorResponse>
+              <Error>
+                <Code>#{error_code}</Code>
+                <Message>stubbed-response-error-message</Message>
+              </Error>
+            </ErrorResponse>
           XML
-          http_resp
+          resp
         end
 
         private
@@ -37,7 +38,7 @@ def build_body(api, operation, data)
           xml.shift
           xml.pop
           xmlns = "http://ec2.amazonaws.com/doc/#{api.version}/".inspect
-          xml.unshift("  <requestId>stubbed-request-id</requestId>")
+          xml.unshift('  <requestId>stubbed-request-id</requestId>')
           xml.unshift("<#{operation.name}Response xmlns=#{xmlns}>\n")
           xml.push("</#{operation.name}Response>\n")
           xml.join

gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/json.rb

@@ -3,27 +3,28 @@
 module Aws
   module Stubbing
     module Protocols
+      # @api private
       class Json
 
         def stub_data(api, operation, data)
           resp = Seahorse::Client::Http::Response.new
           resp.status_code = 200
-          resp.headers["Content-Type"] = content_type(api)
-          resp.headers["x-amzn-RequestId"] = "stubbed-request-id"
+          resp.headers['Content-Type'] = content_type(api)
+          resp.headers['x-amzn-RequestId'] = 'stubbed-request-id'
           resp.body = build_body(operation, data)
           resp
         end
 
         def stub_error(error_code)
-          http_resp = Seahorse::Client::Http::Response.new
-          http_resp.status_code = 400
-          http_resp.body = <<-JSON.strip
-{
-  "code": #{error_code.inspect},
-  "message": "stubbed-response-error-message"
-}
+          resp = Seahorse::Client::Http::Response.new
+          resp.status_code = 400
+          resp.body = <<~JSON.strip
+            {
+              "code": #{error_code.inspect},
+              "message": "stubbed-response-error-message"
+            }
           JSON
-          http_resp
+          resp
         end
 
         private

gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/query.rb

@@ -3,6 +3,7 @@
 module Aws
   module Stubbing
     module Protocols
+      # @api private
       class Query
 
         def stub_data(api, operation, data)
@@ -13,20 +14,20 @@ def stub_data(api, operation, data)
         end
 
         def stub_error(error_code)
-          http_resp = Seahorse::Client::Http::Response.new
-          http_resp.status_code = 400
-          http_resp.body = XmlError.new(error_code).to_xml
-          http_resp
+          resp = Seahorse::Client::Http::Response.new
+          resp.status_code = 400
+          resp.body = XmlError.new(error_code).to_xml
+          resp
         end
 
         private
 
         def build_body(api, operation, data)
           xml = []
           builder = Aws::Xml::DocBuilder.new(target: xml, indent: '  ')
-          builder.node(operation.name + 'Response', xmlns: xmlns(api)) do
+          builder.node("#{operation.name}Response", xmlns: xmlns(api)) do
             if (rules = operation.output)
-              rules.location_name = operation.name + 'Result'
+              rules.location_name = "#{operation.name}Result"
               Xml::Builder.new(rules, target: xml, pad:'  ').to_xml(data)
             end
             builder.node('ResponseMetadata') do