add exp-backoff, handle nexttoken, remove destroy

Author: cyyeh

packages/extension-driver-redshift/package.json

@@ -23,7 +23,8 @@
   },
   "license": "Apache-2.0",
   "dependencies": {
-    "@aws-sdk/client-redshift-data": "^3.405.0"
+    "@aws-sdk/client-redshift-data": "^3.405.0",
+    "exponential-backoff": "^3.1.1"
   },
   "peerDependencies": {
     "@vulcan-sql/core": "~0.9.1-0"

packages/extension-driver-redshift/src/lib/redshiftDataSource.ts

@@ -22,6 +22,7 @@ import {
   GetStatementResultCommand,
   SqlParameter,
 } from '@aws-sdk/client-redshift-data';
+import { backOff } from 'exponential-backoff';
 
 export type RedshiftOptions = RedshiftDataClientConfig & Omit<ExecuteStatementCommandInput, "Sql" | "Parameters">;
 
@@ -89,7 +90,6 @@ export class RedShiftDataSource extends DataSource<any, RedshiftOptions> {
       this.logger.debug(
         `Errors occurred, release connection from ${profileName}`
       );
-      redshiftClient.destroy();
       throw e;
     }
   }
@@ -114,7 +114,9 @@ export class RedShiftDataSource extends DataSource<any, RedshiftOptions> {
       const statementCommandResult = await redshiftClient.send(executeStatementCommand);
       return await this.getResultFromExecuteStatement(statementCommandResult, redshiftClient);
     } catch (e) {
-      redshiftClient.destroy();
+      this.logger.debug(
+        `Errors occurred, release connection from ${profileName}`
+      );
       throw e;
     }
   }
@@ -132,7 +134,7 @@ export class RedShiftDataSource extends DataSource<any, RedshiftOptions> {
     // https://github.com/aws/aws-sdk-js-v3/blob/29056f4ca545f7e5cf951b915bb52178305fc305/clients/client-redshift-data/src/models/models_0.ts#L604
     while (!describeStatementResponse || describeStatementResponse.Status !== 'FINISHED') {
       const describeStatementCommand = new DescribeStatementCommand(describeStatementRequestInput);
-      describeStatementResponse = await redshiftClient.send(describeStatementCommand);
+      describeStatementResponse = await backOff(() =>redshiftClient.send(describeStatementCommand));
 
       if (
         describeStatementResponse.Status === 'ABORTED' || 
@@ -142,15 +144,26 @@ export class RedShiftDataSource extends DataSource<any, RedshiftOptions> {
       }
     }
 
-    const getStatementResultCommandParams: GetStatementResultCommandInput = {
+    let getStatementResultCommandParams: GetStatementResultCommandInput = {
       "Id": describeStatementResponse.Id
     };
-    const getStatementResultCommand = new GetStatementResultCommand(getStatementResultCommandParams);
-    const getStatementResultResponse = await redshiftClient.send(getStatementResultCommand);
+    let getStatementResultCommand = new GetStatementResultCommand(getStatementResultCommandParams);
+    let getStatementResultResponse = await redshiftClient.send(getStatementResultCommand);
+    const records = getStatementResultResponse.Records! || [];
+    const columns = getStatementResultResponse.ColumnMetadata || [];
+
+    while (getStatementResultResponse.NextToken) {
+      getStatementResultCommandParams = {
+        "Id": describeStatementResponse.Id,
+        "NextToken": getStatementResultResponse.NextToken,
+      };
+      getStatementResultCommand = new GetStatementResultCommand(getStatementResultCommandParams);
+      getStatementResultResponse = await redshiftClient.send(getStatementResultCommand);
+      records.push(...(getStatementResultResponse.Records! || []));
+    }
 
     return {
       getColumns: () => {
-        const columns = getStatementResultResponse.ColumnMetadata || [];
         return columns.map((column) => ({
           name: column.name || '',
           type: mapFromRedShiftTypeId(column.typeName?.toLowerCase() || ''),
@@ -159,8 +172,6 @@ export class RedShiftDataSource extends DataSource<any, RedshiftOptions> {
       getData: () => new Readable({
         objectMode: true,
         read() {
-          const records = getStatementResultResponse.Records! || [];
-          const columns = getStatementResultResponse.ColumnMetadata || [];
           for (const record of records) {
             const row: RedShiftDataRow = {};
             for (const [i, recordField] of record.entries()) {

packages/extension-driver-redshift/yarn.lock

@@ -768,6 +768,11 @@ bowser@^2.11.0:
   resolved "https://registry.yarnpkg.com/bowser/-/bowser-2.11.0.tgz#5ca3c35757a7aa5771500c70a73a9f91ef420a8f"
   integrity sha512-AlcaJBi/pqqJBIQ8U9Mcpc9i8Aqxn88Skv5d+xBX006BY5u8N3mGLHa5Lgppa7L/HfwgwLgZ6NYs+Ag6uUmJRA==
 
+exponential-backoff@^3.1.1:
+  version "3.1.1"
+  resolved "https://registry.yarnpkg.com/exponential-backoff/-/exponential-backoff-3.1.1.tgz#64ac7526fe341ab18a39016cd22c787d01e00bf6"
+  integrity sha512-dX7e/LHVJ6W3DE1MHWi9S1EYzDESENfLrYohG2G++ovZrYOkm4Knwa0mc1cn84xJOR4KEU0WSchhLbd0UklbHw==
+
 fast-xml-parser@4.2.5:
   version "4.2.5"
   resolved "https://registry.yarnpkg.com/fast-xml-parser/-/fast-xml-parser-4.2.5.tgz#a6747a09296a6cb34f2ae634019bf1738f3b421f"