Merge pull request #5 from classmethod/use-delegated-yield

Author: yamatatsu

README.md

@@ -34,8 +34,8 @@ import AthenaQuery from "athena-query";
 const athena = new Athena({});
 const athenaQuery = new AthenaQuery(athena);
 
-for await (const items of athenaQuery.query("SELECT * FROM waf_logs;")) {
-  console.log(items); // You can get all items with pagination by query.
+for await (const item of athenaQuery.query("SELECT * FROM waf_logs;")) {
+  console.log(item); // You can get all items across pagination.
 }
 ```
 

lib/athena-query.ts

@@ -16,7 +16,11 @@ export class AthenaQuery {
   async *query(
     sql: string,
     options?: { executionParameters?: string[]; maxResults: number }
-  ): AsyncGenerator<Record<string, string | number | BigInt | null>[]> {
+  ): AsyncGenerator<
+    Record<string, string | number | BigInt | null>,
+    void,
+    undefined
+  > {
     const QueryExecutionId = await helpers.startQueryExecution({
       athena: this.athena,
       sql,
@@ -39,7 +43,7 @@ export class AthenaQuery {
         QueryExecutionId,
       });
 
-      yield queryResults.items;
+      yield* queryResults.items;
 
       nextToken = queryResults.nextToken;
     } while (nextToken);

test/index.test.ts

@@ -62,19 +62,17 @@ test("parse to json following ColumnInfo", async () => {
 
   const res1 = await resultGen.next();
 
-  expect(res1.value).toEqual([
-    {
-      name: "test-name-1",
-      disabled: true,
-      timestamp: 1669718600001n,
-      score1: 101,
-      score2: 102,
-      score3: 103,
-      score4: 104,
-      rate1: 1.01,
-      rate2: 1.02,
-    },
-  ]);
+  expect(res1.value).toEqual({
+    name: "test-name-1",
+    disabled: true,
+    timestamp: 1669718600001n,
+    score1: 101,
+    score2: 102,
+    score3: 103,
+    score4: 104,
+    rate1: 1.01,
+    rate2: 1.02,
+  });
 });
 
 test("wait query completed", async () => {
@@ -91,10 +89,7 @@ test("wait query completed", async () => {
         ResultSetMetadata: {
           ColumnInfo: [{ Name: "name", Type: "varchar" }],
         },
-        Rows: [
-          { Data: [{ VarCharValue: "test-name-1" }] },
-          { Data: [{ VarCharValue: "test-name-2" }] },
-        ],
+        Rows: [{ Data: [{ VarCharValue: "test-name-1" }] }],
       },
     });
 
@@ -103,10 +98,7 @@ test("wait query completed", async () => {
   const res1 = await resultGen.next();
 
   expect(res1.done).toBe(false);
-  expect(res1.value).toEqual([
-    { name: "test-name-1" },
-    { name: "test-name-2" },
-  ]);
+  expect(res1.value).toEqual({ name: "test-name-1" });
 });
 
 test("get items with generator", async () => {
@@ -141,18 +133,19 @@ test("get items with generator", async () => {
 
   const res1 = await queryResultGen.next();
   expect(res1.done).toBe(false);
-  expect(res1.value).toEqual([
-    { name: "test-name-1" },
-    { name: "test-name-2" },
-  ]);
+  expect(res1.value).toEqual({ name: "test-name-1" });
 
   const res2 = await queryResultGen.next();
   expect(res2.done).toBe(false);
-  expect(res2.value).toEqual([{ name: "test-name-3" }]);
+  expect(res2.value).toEqual({ name: "test-name-2" });
 
   const res3 = await queryResultGen.next();
-  expect(res3.done).toBe(true);
-  expect(res3.value).toBe(undefined);
+  expect(res3.done).toBe(false);
+  expect(res3.value).toEqual({ name: "test-name-3" });
+
+  const res4 = await queryResultGen.next();
+  expect(res4.done).toBe(true);
+  expect(res4.value).toBe(undefined);
 });
 
 test("get all item with generator", async () => {
@@ -197,8 +190,8 @@ test("get all item with generator", async () => {
 
   const allItems = [];
 
-  for await (const items of athenaQuery.query("")) {
-    allItems.push(...items);
+  for await (const item of athenaQuery.query("")) {
+    allItems.push(item);
   }
 
   expect(allItems).toEqual([