From 10f7ecace495b618fcc0c9a9b0bc969098bfb5b6 Mon Sep 17 00:00:00 2001
From: Philip Sampaio <philip.sampaio@gmail.com>
Date: Wed, 11 Sep 2024 12:21:37 -0300
Subject: [PATCH] Apply suggestions from code review

Co-authored-by: Wojtek Mach <wojtekmach@users.noreply.github.com>
---
 lib/req_athena.ex        | 23 +++++++++++++----------
 lib/req_athena/s3.ex     |  2 +-
 test/req_athena_test.exs |  2 +-
 3 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/lib/req_athena.ex b/lib/req_athena.ex
index 7d4bd3d..065db04 100644
--- a/lib/req_athena.ex
+++ b/lib/req_athena.ex
@@ -45,23 +45,26 @@ defmodule ReqAthena do
 
     * `:database` - Required. The AWS Athena database name.
 
-    * `:output_location` - Conditional. The S3 URL location to output AWS Athena query results.
+    * `:output_location` - Optional. The S3 URL location to output AWS Athena query results.
       Results will be saved as Parquet and loaded with Explorer only if this option is given.
 
     * `:workgroup` - Conditional. The AWS Athena workgroup.
 
     * `:cache_query` - Optional. Forces a non-cached result from AWS Athena.
 
-    * `:format` - Optional. It changes the output format. By default this is
-      `:none`, which means that we return the decoded result from the Athena API.
-      The supported formats are: `:csv`, `:explorer,`, and `:json`.
+    * `:format` - Optional. The output format. Can be one of:
+
+        * `:none` (default) - return decoded API response from Athena.
+        
+        * `:csv` - return contents of the CSV file.
+        
+        * `:json` - return contents of the JSON file.
+        
+          Note: Req by default automatically decodes JSON response body ([`decode_body`](Req.Steps.decode_body1/) step)
+          and to prevent it from doing so, set `decode_body: false`.
+          
+        * `:explorer` - return contents in parquet format, lazy loaded into Explorer data frame.
 
-      For `:csv`, the contents of the CSV file are the output instead of the API return.
-      When `:json` is used, the contents of the JSON files are going to be the output.
-      Notice that the body is decoded by default and to prevent that, you need to use
-      the `:decode_body` option, so you get the "raw" data.
-      The `:explorer` format will perform the query unloading it to Parquet files, and
-      then will lazy load these parquet files into an Explorer dataframe.
 
       There are some limitations when using the `:json` and `:explorer` format.
       First, you need to install Explorer in order to use the `:explorer` format.
diff --git a/lib/req_athena/s3.ex b/lib/req_athena/s3.ex
index c8ffd09..daa76f7 100644
--- a/lib/req_athena/s3.ex
+++ b/lib/req_athena/s3.ex
@@ -14,7 +14,7 @@ defmodule ReqAthena.S3 do
   end
 
   def get_body(req_s3, location) do
-    response = Req.get!(req_s3, url: location)
+    %{status: 200} = response = Req.get!(req_s3, url: location)
 
     response.body
   end
diff --git a/test/req_athena_test.exs b/test/req_athena_test.exs
index 2759429..1e93b32 100644
--- a/test/req_athena_test.exs
+++ b/test/req_athena_test.exs
@@ -433,7 +433,7 @@ defmodule ReqAthenaTest do
              |> Req.Request.put_private(:athena_dataframe_builder, fn output_location,
                                                                       credentials,
                                                                       decode_body ->
-               assert String.starts_with?(output_location, "s3://")
+               assert "s3://" <> _ = output_location
 
                assert decode_body