From 6b4fd5477bc58449e584fa721f9a5f74121b1c78 Mon Sep 17 00:00:00 2001
From: Curtis Vogt <curtis.vogt@gmail.com>
Date: Fri, 18 Jun 2021 10:49:46 -0500
Subject: [PATCH] Introduce AWS.Response

---
 src/AWS.jl                 |  25 ++++++-
 src/AWSMetadata.jl         |   1 +
 src/utilities/request.jl   |  65 +++--------------
 src/utilities/response.jl  |  75 +++++++++++++++++++
 src/utilities/utilities.jl |   3 -
 test/AWS.jl                | 146 +++++++++++--------------------------
 6 files changed, 148 insertions(+), 167 deletions(-)
 create mode 100644 src/utilities/response.jl

diff --git a/src/AWS.jl b/src/AWS.jl
index c5c60beb30..2c638f6a3a 100644
--- a/src/AWS.jl
+++ b/src/AWS.jl
@@ -32,6 +32,7 @@ include("AWSConfig.jl")
 include("AWSMetadata.jl")
 
 include(joinpath("utilities", "request.jl"))
+include(joinpath("utilities", "response.jl"))
 include(joinpath("utilities", "sign.jl"))
 
 
@@ -203,6 +204,9 @@ function (service::RestXMLService)(
     aws_config::AbstractAWSConfig=global_aws_config(),
 )
     return_headers = _pop!(args, "return_headers", false)
+    return_stream = _pop!(args, "return_stream", false)
+    return_raw = _pop!(args, "return_raw", false)
+    response_stream = _pop!(args, "response_stream", nothing)
 
     request = Request(;
         _extract_common_kw_args(service, args)...,
@@ -226,7 +230,8 @@ function (service::RestXMLService)(
 
     request.url = generate_service_url(aws_config, request.service, request.resource)
 
-    return submit_request(aws_config, request; return_headers=return_headers)
+    response = submit_request(aws_config, request)
+    return legacy(response; return_headers, return_stream, return_raw, response_stream)
 end
 
 
@@ -254,6 +259,9 @@ function (service::QueryService)(
 )
     POST_RESOURCE = "/"
     return_headers = _pop!(args, "return_headers", false)
+    return_stream = _pop!(args, "return_stream", false)
+    return_raw = _pop!(args, "return_raw", false)
+    response_stream = _pop!(args, "response_stream", nothing)
 
     request = Request(;
         _extract_common_kw_args(service, args)...,
@@ -268,7 +276,8 @@ function (service::QueryService)(
     args["Version"] = service.api_version
     request.content = HTTP.escapeuri(_flatten_query(service.name, args))
 
-    return submit_request(aws_config, request; return_headers=return_headers)
+    response = submit_request(aws_config, request)
+    return legacy(response; return_headers, return_stream, return_raw, response_stream)
 end
 
 """
@@ -295,6 +304,9 @@ function (service::JSONService)(
 )
     POST_RESOURCE = "/"
     return_headers = _pop!(args, "return_headers", false)
+    return_stream = _pop!(args, "return_stream", false)
+    return_raw = _pop!(args, "return_raw", false)
+    response_stream = _pop!(args, "response_stream", nothing)
 
     request = Request(;
         _extract_common_kw_args(service,args)...,
@@ -307,7 +319,8 @@ function (service::JSONService)(
     request.headers["Content-Type"] = "application/x-amz-json-$(service.json_version)"
     request.headers["X-Amz-Target"] = "$(service.target).$(operation)"
 
-    return submit_request(aws_config, request; return_headers=return_headers)
+    response = submit_request(aws_config, request)
+    return legacy(response; return_headers, return_stream, return_raw, response_stream)
 end
 
 """
@@ -334,6 +347,9 @@ function (service::RestJSONService)(
     aws_config::AbstractAWSConfig=global_aws_config(),
 )
     return_headers = _pop!(args, "return_headers", false)
+    return_stream = _pop!(args, "return_stream", false)
+    return_raw = _pop!(args, "return_raw", false)
+    response_stream = _pop!(args, "response_stream", nothing)
 
     request = Request(;
         _extract_common_kw_args(service, args)...,
@@ -350,7 +366,8 @@ function (service::RestJSONService)(
     request.headers["Content-Type"] = "application/json"
     request.content = json(args)
 
-    return submit_request(aws_config, request; return_headers=return_headers)
+    response = submit_request(aws_config, request)
+    return legacy(response; return_headers, return_stream, return_raw, response_stream)
 end
 
 end  # module AWS
diff --git a/src/AWSMetadata.jl b/src/AWSMetadata.jl
index 25d1255a58..fc27f60603 100644
--- a/src/AWSMetadata.jl
+++ b/src/AWSMetadata.jl
@@ -1,3 +1,4 @@
+
 module AWSMetadata
 
 using Base64
diff --git a/src/utilities/request.jl b/src/utilities/request.jl
index e2b4815a2a..52b7e550d2 100644
--- a/src/utilities/request.jl
+++ b/src/utilities/request.jl
@@ -8,16 +8,13 @@ Base.@kwdef mutable struct Request
     resource::String=""
     url::String=""
 
-    return_stream::Bool=false
-    response_stream::Union{<:IO, Nothing}=nothing
     http_options::AbstractDict{Symbol,<:Any}=LittleDict{Symbol,String}()
-    return_raw::Bool=false
     response_dict_type::Type{<:AbstractDict}=LittleDict
 end
 
 
 """
-    submit_request(aws::AbstractAWSConfig, request::Request; return_headers::Bool=false)
+    submit_request(aws::AbstractAWSConfig, request::Request)
 
 Submit the request to AWS.
 
@@ -25,13 +22,10 @@ Submit the request to AWS.
 - `aws::AbstractAWSConfig`: AWSConfig containing credentials and other information for fulfilling the request, default value is the global configuration
 - `request::Request`: All the information about making a request to AWS
 
-# Keywords
-- `return_headers::Bool=false`: True if you want the headers from the response returned back
-
 # Returns
-- `Tuple or Dict`: Tuple if returning_headers, otherwise just return a Dict of the response body
+- `Response`: A struct containing the response
 """
-function submit_request(aws::AbstractAWSConfig, request::Request; return_headers::Bool=false)
+function submit_request(aws::AbstractAWSConfig, request::Request)
     response = nothing
     TOO_MANY_REQUESTS = 429
     EXPIRED_ERROR_CODES = ["ExpiredToken", "ExpiredTokenException", "RequestExpired"]
@@ -91,48 +85,7 @@ function submit_request(aws::AbstractAWSConfig, request::Request; return_headers
         end
     end
 
-    response_dict_type = request.response_dict_type
-
-    # For HEAD request, return headers...
-    if request.request_method == "HEAD"
-        return response_dict_type(response.headers)
-    end
-
-    # Return response stream if requested...
-    if request.return_stream
-        return request.response_stream
-    end
-
-    # Return raw data if requested...
-    if request.return_raw
-        return (return_headers ? (response.body, response.headers) : response.body)
-    end
-
-    # Parse response data according to mimetype...
-    mime = HTTP.header(response, "Content-Type", "")
-
-    if isempty(mime)
-        if length(response.body) > 5 && response.body[1:5] == b"<?xml"
-            mime = "text/xml"
-        end
-    end
-
-    body = String(copy(response.body))
-
-    if occursin(r"/xml", mime)
-        xml_dict_type = response_dict_type{Union{Symbol, String}, Any}
-        body = parse_xml(body)
-        root = XMLDict.root(body.x)
-
-        return (return_headers ? (xml_dict(root, xml_dict_type), response_dict_type(response.headers)) : xml_dict(root, xml_dict_type))
-    elseif occursin(r"/x-amz-json-1.[01]$", mime) || endswith(mime, "json")
-        info = isempty(response.body) ? nothing : JSON.parse(body, dicttype=response_dict_type)
-        return (return_headers ? (info, response_dict_type(response.headers)) : info)
-    elseif startswith(mime, "text/")
-        return (return_headers ? (body, response_dict_type(response.headers)) : body)
-    else
-        return (return_headers ? (response.body, response.headers) : response.body)
-    end
+    return response
 end
 
 
@@ -140,20 +93,20 @@ function _http_request(request::Request)
     @repeat 4 try
         http_stack = HTTP.stack(redirect=false, retry=false, aws_authorization=false)
 
-        if request.return_stream && request.response_stream === nothing
-            request.response_stream = Base.BufferStream()
-        end
+        response_stream = Base.BufferStream()
 
-        return HTTP.request(
+        r = HTTP.request(
             http_stack,
             request.request_method,
             HTTP.URI(request.url),
             HTTP.mkheaders(request.headers),
             request.content;
             require_ssl_verification=false,
-            response_stream=request.response_stream,
+            response_stream=response_stream,
             request.http_options...
         )
+
+        return Response(r, response_stream)
     catch e
         # Base.IOError is needed because HTTP.jl can often have errors that aren't
         # caught and wrapped in an HTTP.IOError
diff --git a/src/utilities/response.jl b/src/utilities/response.jl
new file mode 100644
index 0000000000..bcb1d25705
--- /dev/null
+++ b/src/utilities/response.jl
@@ -0,0 +1,75 @@
+# TODO: Include only the fields we care about
+struct Response
+    response::HTTP.Response
+    body::BufferStream
+end
+
+function Base.getproperty(r::Response, f::Symbol)
+    if f === :headers
+        r.response.headers
+    else
+        getfield(r, f)
+    end
+end
+
+function mime_type(r::Response)
+    # Parse response data according to mimetype...
+    mime = HTTP.header(r.response, "Content-Type", "")
+
+    # https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/MIME_types/Common_types
+
+    # Reading consumes the buffer so this is unsafe currently
+    # if isempty(mime) && read(r.body, 5) == b"<?xml"
+    #     "xml"
+
+    # e.g. "application/xml" or "text/xml"
+    if occursin(r"/xml", mime)
+        MIME(Symbol("application/xml"))
+    elseif occursin(r"/x-amz-json-1.[01]$", mime) || endswith(mime, "json")
+        MIME(Symbol("application/json"))
+    elseif startswith(mime, "text/")
+        MIME(Symbol("text/plain"))
+    else
+        MIME(Symbol("text/plain"))
+    end
+end
+
+
+# TODO: Interface isn't perfect. For example "text/plain" completely ignores `T`
+
+Base.parse(::Type{T}, r::Response) where T = parse(T, r, mime_type(r))
+
+function Base.parse(::Type{T}, r::Response, ::MIME"application/xml") where T <: AbstractDict
+    xml = parse_xml(String(r.body))
+    root = XMLDict.root(xml.x)  # TODO: Why x?
+    return xml_dict(root, T))
+end
+
+function Base.parse(::Type{T}, r::Response, ::MIME"application/json") where T <: AbstractDict
+    return JSON.parse(r.body, dicttype=T)
+end
+
+function Base.parse(::Type{T}, r::Response, ::MIME"text/plain") where T <: AbstractDict
+    return String(r.body)
+end
+
+function legacy(r::Response; response_dict_type::Type, return_headers=false, return_stream=false, return_raw=false, response_stream=nothing)
+    # For HEAD request, return headers...
+    if r.response.request.method == "HEAD"
+        return response_dict_type(r.response.headers)
+    end
+
+    # Return response stream if requested...
+    if return_stream
+        return write(response_stream, read(r.body))
+    end
+
+    # Return raw data if requested...
+    if return_raw
+        content = read(r.body)
+        return return_headers ? (content, r.response.headers) : content
+    end
+
+    content = parse(response_dict_type, r)
+    return return_headers ? (content, r.response.headers) : content
+end
diff --git a/src/utilities/utilities.jl b/src/utilities/utilities.jl
index a3d5dedad9..e0414e7bbd 100644
--- a/src/utilities/utilities.jl
+++ b/src/utilities/utilities.jl
@@ -61,9 +61,6 @@ function _extract_common_kw_args(service, args)
     return (
         service=service.name,
         api_version=service.api_version,
-        return_stream=_pop!(args, "return_stream", false),
-        return_raw=_pop!(args, "return_raw", false),
-        response_stream=_pop!(args, "response_stream", nothing),
         headers=LittleDict{String, String}(_pop!(args, "headers", [])),
         http_options=_pop!(args, "http_options", LittleDict{Symbol, String}()),
         response_dict_type=_pop!(args, "response_dict_type", LittleDict),
diff --git a/test/AWS.jl b/test/AWS.jl
index da12d7a07a..4e968fdbb4 100644
--- a/test/AWS.jl
+++ b/test/AWS.jl
@@ -112,8 +112,6 @@ end
     end
 
     @testset "HEAD request method" begin
-        expected_result_type = LittleDict
-
         request = Request(
             service="s3",
             api_version="api_version",
@@ -123,10 +121,10 @@ end
 
         apply(Patches._aws_http_request_patch) do
             Patches._response!()
-            result = AWS.submit_request(aws, request)
+            response = AWS.submit_request(aws, request)
 
-            @test result == expected_result_type(Patches.headers)
-            @test result isa expected_result_type
+            @test response.headers == Patches.headers
+            @test response isa Vector
         end
     end
 
@@ -135,16 +133,14 @@ end
             service="s3",
             api_version="api_version",
             request_method="GET",
-            return_stream=true,
-            response_stream=Base.BufferStream(),
             url="https://s3.us-east-1.amazonaws.com/sample-bucket"
         )
 
         apply(Patches._aws_http_request_patch) do
             Patches._response!()
-            result = AWS.submit_request(aws, request)
+            response = AWS.submit_request(aws, request)
 
-            @test result == request.response_stream
+            @test response.body isa Base.BufferStream
         end
     end
 
@@ -157,23 +153,12 @@ end
             return_raw=true
         )
 
-        @testset "body" begin
-            apply(Patches._aws_http_request_patch) do
-                Patches._response!()
-                result = AWS.submit_request(aws, request)
-
-                @test String(result) == Patches.body
-            end
-        end
-
-        @testset "body and headers" begin
-            apply(Patches._aws_http_request_patch) do
-                Patches._response!()
-                body, headers = AWS.submit_request(aws, request; return_headers=true)
+        apply(Patches._aws_http_request_patch) do
+            Patches._response!()
+            response = AWS.submit_request(aws, request)
 
-                @test String(body) == Patches.body
-                @test headers == Patches.headers
-            end
+            @test String(response.body) == Patches.body
+            @test response.headers == Patches.headers
         end
     end
 
@@ -189,49 +174,30 @@ end
             apply(Patches._aws_http_request_patch) do
                 @testset "default" begin
                     expected_body = ""
-                    expected_headers = Pair["Content-Type"=>"",]
+                    expected_headers = Pair["Content-Type" => ""]
 
                     Patches._response!(headers=expected_headers, body=expected_body)
 
-                    @testset "body" begin
-                       result = AWS.submit_request(aws, request)
+                    response = AWS.submit_request(aws, request)
 
-                       @test String(result) == expected_body
-                    end
-
-                    @testset "body and headers" begin
-                        body, headers = AWS.submit_request(aws, request; return_headers=true)
-
-                        @test String(body) == expected_body
-                        @test headers == expected_headers
-                    end
+                    @test String(response.body) == expected_body
+                    @test response.headers == expected_headers
                 end
 
                 @testset "text/xml" begin
-                    expected_headers = Pair["Content-Type"=>"text/xml",]
+                    expected_headers = Pair["Content-Type" => "text/xml"]
                     expected_body_type = LittleDict{Union{Symbol, String}, Any}
                     expected_body = _expected_body(Patches.body, expected_body_type)
 
-                    expected_header_type = LittleDict{SubString{String}, SubString{String}}
-
                     Patches._response!(headers=expected_headers)
 
-                    @testset "body" begin
-                        result = AWS.submit_request(aws, request)
-
-                        @test result isa expected_body_type
-                        @test result == expected_body
-                    end
-
-                    @testset "body and headers" begin
-                        body, headers = AWS.submit_request(aws, request; return_headers=true)
+                    response = AWS.submit_request(aws, request)
+                    body = parse(expected_body_type, response)
 
-                        @test body == expected_body
-                        @test body isa expected_body_type
-
-                        @test headers == LittleDict(expected_headers)
-                        @test headers isa expected_header_type
-                    end
+                    @test body isa expected_body_type
+                    @test body == expected_body
+                    @test response.headers == expected_headers
+                    @test response.headers isa Vector
                 end
             end
         end
@@ -246,32 +212,23 @@ end
 
             expected_body_type = LittleDict{Union{Symbol, String}, Any}
             expected_body = _expected_body(Patches.body, expected_body_type)
-            expected_headers = Pair["Content-Type"=>"application/xml",]
+            expected_headers = Pair["Content-Type" => "application/xml"]
 
             apply(Patches._aws_http_request_patch) do
                 Patches._response!(headers=expected_headers,)
 
-                @testset "body" begin
-                    result = AWS.submit_request(aws, request)
+                response = AWS.submit_request(aws, request)
+                body = parse(expected_body_type, response)
 
-                    @test result == expected_body
-                    @test result isa expected_body_type
-                end
-
-                @testset "body and headers" begin
-                    body, headers = AWS.submit_request(aws, request; return_headers=true)
-
-                    @test body == expected_body
-                    @test body isa expected_body_type
-
-                    @test headers == LittleDict(expected_headers)
-                    @test headers isa LittleDict{SubString{String}, SubString{String}}
-               end
+                @test body isa expected_body_type
+                @test body == expected_body
+                @test response.headers isa Vector
+                @test response.headers == expected_headers
             end
         end
 
         @testset "JSON" begin
-            json_headers = ["Content-Type"=>"application/json"]
+            json_headers = ["Content-Type" => "application/json"]
             body = Dict{String,Any}(
                 "Marker" => nothing,
                 "VaultList" => Any[Dict{String,Any}(
@@ -289,24 +246,15 @@ end
             expected_body = JSON.parse(json_body, dicttype=LittleDict)
 
             apply(Patches._aws_http_request_patch) do
-                Patches._response!(body=json_body, headers=json_headers,)
-
-                @testset "body" begin
-                    result = AWS.submit_request(aws, request)
-
-                    @test result isa expected_body_type
-                    @test result == expected_body
-                end
-
-                @testset "body and headers" begin
-                    body, headers = AWS.submit_request(aws, request; return_headers=true)
+                Patches._response!(body=json_body, headers=json_headers)
 
-                    @test body == expected_body
-                    @test body isa expected_body_type
+                response = AWS.submit_request(aws, request)
+                body = parse(expected_body_type, response)
 
-                    @test headers == LittleDict(json_headers)
-                    @test headers isa LittleDict{SubString{String}, SubString{String}}
-                end
+                @test body isa expected_body_type
+                @test body == expected_body
+                @test response.headers isa Vector
+                @test response.headers == expected_headers
             end
         end
 
@@ -321,26 +269,16 @@ end
             apply(Patches._aws_http_request_patch) do
                 expected_headers = ["Content-Type" => "text/html"]
                 expected_body = Patches.body
-                expected_header_type = LittleDict{SubString{String}, SubString{String}}
 
                 Patches._response!(headers=expected_headers)
 
-                @testset "body" begin
-                    result = AWS.submit_request(aws, request)
-
-                    @test result isa String
-                    @test result == expected_body
-                end
+                response = AWS.submit_request(aws, request)
+                body = String(response.body)
 
-                @testset "body and headers" begin
-                    body, headers = AWS.submit_request(aws, request; return_headers=true)
-
-                    @test body == expected_body
-                    @test body isa String
-
-                    @test headers == LittleDict(expected_headers)
-                    @test headers isa expected_header_type
-                end
+                @test body isa String
+                @test body == expected_body
+                @test response.headers isa Vector
+                @test response.headers == expected_headers
             end
         end
     end