From 8a8467debe3b3143744fdf206325290bc3b7cceb Mon Sep 17 00:00:00 2001
From: KantaHasegawa-26 <161685065+26KantaHasegawa@users.noreply.github.com>
Date: Thu, 14 Mar 2024 00:29:23 +0900
Subject: [PATCH] feat add option outputlocation

---
 src/athena-query.ts | 5 +++++
 src/helper.ts       | 5 +++++
 test/index.test.ts  | 4 ++++
 3 files changed, 14 insertions(+)

diff --git a/src/athena-query.ts b/src/athena-query.ts
index eb28781..55570ac 100644
--- a/src/athena-query.ts
+++ b/src/athena-query.ts
@@ -17,6 +17,11 @@ type Options = {
    * The name of the data catalog used in the query execution.
    */
   catalog?: string;
+
+  /**
+   * The location in Amazon S3 where your query and calculation results are stored.
+   */
+  outputLocation?: string;
 };
 
 export class AthenaQuery {
diff --git a/src/helper.ts b/src/helper.ts
index 1dfda65..f0de841 100644
--- a/src/helper.ts
+++ b/src/helper.ts
@@ -1,6 +1,7 @@
 import type {
   Athena,
   GetQueryResultsCommandOutput,
+  ResultConfiguration,
 } from "@aws-sdk/client-athena";
 
 export type AtheneRecordData = Record<string, string | number | BigInt | null>;
@@ -13,6 +14,7 @@ async function startQueryExecution(params: {
   workgroup?: string;
   db?: string;
   catalog?: string;
+  outputLocation?: string;
 }) {
   const output = await params.athena.startQueryExecution({
     QueryString: params.sql,
@@ -22,6 +24,9 @@ async function startQueryExecution(params: {
       Database: params.db || "default",
       Catalog: params.catalog,
     },
+    ResultConfiguration: {
+      OutputLocation: params.outputLocation,
+    },
   });
 
   if (!output.QueryExecutionId) {
diff --git a/test/index.test.ts b/test/index.test.ts
index ba64239..dea5083 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -258,6 +258,7 @@ test("pass args to sdk", async () => {
     db: "test-db",
     workgroup: "test-workgroup",
     catalog: "test-catalog",
+    outputLocation: "s3//example/path",
   });
   const resultGen = athenaQuery.query("SELECT test FROM test;", {
     executionParameters: ["test", 123, 456n],
@@ -276,6 +277,9 @@ test("pass args to sdk", async () => {
       Catalog: "test-catalog",
       Database: "test-db",
     },
+    ResultConfiguration: {
+      OutputLocation: "s3//example/path",
+    },
   });
 
   expect(