Improve awsQueryCompatible support for JSON and CBOR (#3128)

aws · Oct 18, 2024 · 7fd998d · 7fd998d
1 parent d555e5a
commit 7fd998d
Show file tree

Hide file tree

Showing 30 changed files with 4,253 additions and 320 deletions.
diff --git a/gems/aws-sdk-core/CHANGELOG.md b/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)

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/client_stubs.rb b/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

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/json/error_handler.rb b/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']

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/json/handler.rb b/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
 

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/rpc_v2/error_handler.rb b/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']

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/rpc_v2/handler.rb b/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)

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/stubbing.rb b/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
diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/ec2.rb b/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

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/json.rb b/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

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/query.rb b/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

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/rest.rb b/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/rest.rb
@@ -5,6 +5,7 @@
 module Aws
   module Stubbing
     module Protocols
+      # @api private
       class Rest
 
         include Seahorse::Model::Shapes
@@ -22,7 +23,7 @@ def stub_data(api, operation, data)
         def new_http_response
           resp = Seahorse::Client::Http::Response.new
           resp.status_code = 200
-          resp.headers["x-amzn-RequestId"] = "stubbed-request-id"
+          resp.headers['x-amzn-RequestId'] = 'stubbed-request-id'
           resp
         end
 

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/rest_json.rb b/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/rest_json.rb
@@ -3,6 +3,7 @@
 module Aws
   module Stubbing
     module Protocols
+      # @api private
       class RestJson < Rest
 
         def body_for(_a, _b, rules, data)
@@ -14,15 +15,15 @@ def body_for(_a, _b, rules, data)
         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
 
       end

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/rest_xml.rb b/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/rest_xml.rb
@@ -3,25 +3,26 @@
 module Aws
   module Stubbing
     module Protocols
+      # @api private
       class RestXml < Rest
 
         def body_for(api, operation, rules, data)
           if eventstream?(rules)
             encode_eventstream_response(rules, data, Xml::Builder)
           else
             xml = []
-            rules.location_name = operation.name + 'Result'
+            rules.location_name = "#{operation.name}Result"
             rules['xmlNamespace'] = { 'uri' => api.metadata['xmlNamespace'] }
             Xml::Builder.new(rules, target:xml).to_xml(data)
             xml.join
           end
         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
 
         def xmlns(api)

diff --git a/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/rpc_v2.rb b/gems/aws-sdk-core/lib/aws-sdk-core/stubbing/protocols/rpc_v2.rb
@@ -3,35 +3,33 @@
 module Aws
   module Stubbing
     module Protocols
+      # @api private
       class RpcV2
 
-        def stub_data(api, operation, data)
+        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['Smithy-Protocol'] = 'rpc-v2-cbor'
+          resp.headers['Content-Type'] = 'application/cbor'
           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"
-}
-          JSON
-          http_resp
+          resp = Seahorse::Client::Http::Response.new
+          resp.status_code = 400
+          resp.body = Aws::RpcV2.encode(
+            {
+              'code' => error_code,
+              'message' => 'stubbed-response-error-message'
+            }
+          )
+          resp
         end
 
         private
 
-        def content_type(api)
-          'application/cbor'
-        end
-
         def build_body(operation, data)
           Aws::RpcV2::Builder.new(operation.output).serialize(data)
         end

diff --git a/gems/aws-sdk-core/spec/api_helper.rb b/gems/aws-sdk-core/spec/api_helper.rb
@@ -45,6 +45,14 @@ def sample_ec2 # ec2 has its own protocol
       end
     end
 
+    def sample_rpcv2_cbor # cloudwatch logs changed to cbor
+      @sample_rpcv2_cbor ||= begin
+        api = File.expand_path('../fixtures/apis/logs.json', __FILE__)
+        api = JSON.load(File.read(api))
+        sample_service(api: api)
+      end
+    end
+
     def sample_shapes
       {
         'StructureShape' => {