feat add option outputlocation

classmethod · Mar 13, 2024 · 8a8467d · 8a8467d
1 parent f159b8f
commit 8a8467d
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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(