From 17a6f4f1fbb9762f3dd195903618b4e495d07c41 Mon Sep 17 00:00:00 2001
From: Jon Carstens <jjcarstens@me.com>
Date: Tue, 11 Nov 2025 17:14:00 -0700
Subject: [PATCH] Use returned `DataManifestLocation` if available

Fixes #42

In some situations, the manifest file may be a slightly different path than
what is set in `output_locations` (reasons not fully known). This would
cause failures fetching the locations of results since the manifest would
not exist due to the hardcoded path.

This changes to use the `DataManifestLocation` key if it is returned
in the results rather than assuming the path
---
 lib/req_athena.ex    | 24 ++++++++++++++----------
 lib/req_athena/s3.ex |  6 ++----
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/lib/req_athena.ex b/lib/req_athena.ex
index cdcd16b..0878d6c 100644
--- a/lib/req_athena.ex
+++ b/lib/req_athena.ex
@@ -351,12 +351,12 @@ defmodule ReqAthena do
   end
 
   defp get_json_result(request, response) do
-    output_location = Request.get_private(request, :athena_output_location)
+    manifest_location = Request.get_private(request, :athena_manifest_location)
 
     aws_credentials = aws_credentials_from_request(request)
     req_s3 = ReqAthena.S3.new(aws_credentials)
 
-    locations = ReqAthena.S3.get_locations(req_s3, output_location)
+    locations = ReqAthena.S3.get_locations(req_s3, manifest_location)
 
     # OPTIMIZE: use tasks to retrieve locations.
     results =
@@ -374,7 +374,7 @@ defmodule ReqAthena do
   end
 
   defp get_explorer_result(request, response) do
-    output_location = Request.get_private(request, :athena_output_location)
+    manifest_location = Request.get_private(request, :athena_manifest_location)
 
     aws_credentials = aws_credentials_from_request(request)
 
@@ -384,7 +384,7 @@ defmodule ReqAthena do
 
     decode_body = Req.Request.get_option(request, :decode_body, true)
 
-    result = fetcher_and_builder.(output_location, aws_credentials, decode_body)
+    result = fetcher_and_builder.(manifest_location, aws_credentials, decode_body)
 
     Request.halt(request, %{response | body: result})
   end
@@ -397,9 +397,9 @@ defmodule ReqAthena do
   end
 
   @doc false
-  def fetch_and_build_dataframe(output_location, aws_credentials, decode_body) do
+  def fetch_and_build_dataframe(manifest_location, aws_credentials, decode_body) do
     req_s3 = ReqAthena.S3.new(aws_credentials)
-    locations = ReqAthena.S3.get_locations(req_s3, output_location)
+    locations = ReqAthena.S3.get_locations(req_s3, manifest_location)
 
     if decode_body do
       build_lazy_frame(locations, aws_credentials)
@@ -456,13 +456,17 @@ defmodule ReqAthena do
         Request.halt(request, Req.post!(request))
 
       "SUCCEEDED" ->
+        output_location = body["QueryExecution"]["ResultConfiguration"]["OutputLocation"]
+
+        manifest_location =
+          body["QueryExecution"]["Statistics"]["DataManifestLocation"] ||
+            output_location <> "-manifest.csv"
+
         request =
           request
           |> prepare_action("GetQueryResults")
-          |> Request.put_private(
-            :athena_output_location,
-            body["QueryExecution"]["ResultConfiguration"]["OutputLocation"]
-          )
+          |> Request.put_private(:athena_output_location, output_location)
+          |> Request.put_private(:athena_manifest_location, manifest_location)
           |> Request.put_private(:athena_query_execution_id, query_execution_id)
 
         Request.halt(request, Req.post!(request))
diff --git a/lib/req_athena/s3.ex b/lib/req_athena/s3.ex
index daa76f7..eca8178 100644
--- a/lib/req_athena/s3.ex
+++ b/lib/req_athena/s3.ex
@@ -4,11 +4,9 @@ defmodule ReqAthena.S3 do
     options |> Req.new() |> ReqS3.attach(aws_sigv4: aws_credentials)
   end
 
-  def get_locations(req_s3, output_location) do
-    manifest_csv_location = output_location <> "-manifest.csv"
-
+  def get_locations(req_s3, manifest_location) do
     req_s3
-    |> get_body(manifest_csv_location)
+    |> get_body(manifest_location)
     |> String.trim()
     |> String.split("\n")
   end