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 51fb815..d7ed99c 100644 --- a/lib/req_athena/s3.ex +++ b/lib/req_athena/s3.ex @@ -4,12 +4,10 @@ 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 |> Req.merge(decode_body: false) - |> get_body(manifest_csv_location) + |> get_body(manifest_location) |> String.trim() |> String.split("\n") end