fix: recover Thrift CloudFetch downloads from rejected URLs, 4xx and stalls

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Thrift CloudFetch downloads now recover from more failures. When cloud storage rejects a presigned URL with 401, 403 or 404 (an expired or invalid signature), the driver fetches a fresh link for the same rows from the server and retries with it; each refresh counts as an attempt against `retryMaxAttempts`. Other 4xx responses are retried with the same URL. `socketTimeout` is now also the inactivity timeout of each download attempt, so an attempt that stops receiving data is retried instead of waiting until `retriesTimeout`; one `retriesTimeout` deadline still covers the whole download, including link refreshes. Previously these failures failed the query, and a stalled attempt waited out the whole deadline.
 - Thrift CloudFetch downloads are now bounded end to end by `retriesTimeout` (default 15 minutes), including every attempt and retry backoff. Previously a stalled download, where the response headers or body never arrived, blocked the fetch forever. A download still running when the deadline expires is aborted and fails with `RetryError` (`TIMEOUT_EXCEEDED`, message `CloudFetch result download failed: Retry timeout exceeded`). (#520)
 - Thrift connections made through `proxy` now apply `checkServerCertificate`, `customCaCert` and `clientCert`/`clientKey`. Previously the proxied path ignored them and verified against Node's process-wide trust store, so a certificate outside the configured trust set could be accepted. **Compatibility:** proxied connections now verify unless `checkServerCertificate: false` is set; as on direct connections, `NODE_TLS_REJECT_UNAUTHORIZED=0` no longer turns verification off. For an `https` proxy these options also apply to the proxy's own certificate. (#485)
 - Honor the client `useCloudFetch: false` default on the kernel backend so complete results can be fetched without cloud-storage downloads. (#574)

lib/DBSQLClient.ts

@@ -742,6 +742,10 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
     if (options.retryDelayMax !== undefined) {
       this.config.retryDelayMax = options.retryDelayMax;
     }
+    // Also bounds each CloudFetch download attempt on the Thrift backend
+    if (options.socketTimeout !== undefined) {
+      this.config.socketTimeout = options.socketTimeout;
+    }
 
     // Override telemetry config if provided in options. Per-key narrowed copy
     // preserves the structural type system: `ConnectionOptions` and

lib/connection/connections/HttpConnection.ts

@@ -4,7 +4,7 @@ import http from 'http';
 import { HeadersInit } from 'node-fetch';
 import { ProxyAgent } from 'proxy-agent';
 
-import IConnectionProvider, { HttpTransactionDetails } from '../contracts/IConnectionProvider';
+import IConnectionProvider, { HttpRetryPolicyOptions, HttpTransactionDetails } from '../contracts/IConnectionProvider';
 import IConnectionOptions, { ProxyOptions } from '../contracts/IConnectionOptions';
 import IClientContext from '../../contracts/IClientContext';
 import { LogLevel } from '../../contracts/IDBSQLLogger';
@@ -227,7 +227,7 @@ export default class HttpConnection implements IConnectionProvider {
     return this.connection;
   }
 
-  public async getRetryPolicy(): Promise<IRetryPolicy<HttpTransactionDetails>> {
-    return new HttpRetryPolicy(this.context);
+  public async getRetryPolicy(options?: HttpRetryPolicyOptions): Promise<IRetryPolicy<HttpTransactionDetails>> {
+    return new HttpRetryPolicy(this.context, options);
   }
 }

lib/connection/connections/HttpRetryPolicy.ts

@@ -1,5 +1,5 @@
 import IRetryPolicy, { ShouldRetryResult, RetryableOperation } from '../contracts/IRetryPolicy';
-import { HttpTransactionDetails } from '../contracts/IConnectionProvider';
+import { HttpRetryPolicyOptions, HttpTransactionDetails } from '../contracts/IConnectionProvider';
 import IClientContext from '../../contracts/IClientContext';
 import RetryError, { RetryErrorCode } from '../../errors/RetryError';
 
@@ -53,10 +53,13 @@ export default class HttpRetryPolicy implements IRetryPolicy<HttpTransactionDeta
 
   private attempt: number;
 
-  constructor(context: IClientContext) {
+  private readonly options: HttpRetryPolicyOptions;
+
+  constructor(context: IClientContext, options: HttpRetryPolicyOptions = {}) {
     this.context = context;
     this.startTime = Date.now();
     this.attempt = 0;
+    this.options = options;
   }
 
   public async shouldRetry(details: HttpTransactionDetails): Promise<ShouldRetryResult> {
@@ -193,14 +196,20 @@ export default class HttpRetryPolicy implements IRetryPolicy<HttpTransactionDeta
 
   protected isRetryable({ response }: HttpTransactionDetails): boolean {
     const statusCode = response.status;
+    const { retryClientErrorsExcept } = this.options;
 
     const result =
       // Retry on all codes below 100
       statusCode < 100 ||
       // ...and on `429 Too Many Requests`
       statusCode === 429 ||
       // ...and on all `5xx` codes except for `501 Not Implemented`
-      (statusCode >= 500 && statusCode !== 501);
+      (statusCode >= 500 && statusCode !== 501) ||
+      // ...and, when enabled, on the other `4xx` codes
+      (retryClientErrorsExcept !== undefined &&
+        statusCode >= 400 &&
+        statusCode < 500 &&
+        !retryClientErrorsExcept.has(statusCode));
 
     return result;
   }

lib/connection/contracts/IConnectionProvider.ts

@@ -7,6 +7,14 @@ export interface HttpTransactionDetails {
   response: Response;
 }
 
+export interface HttpRetryPolicyOptions {
+  /**
+   * Also retry 4xx responses, except these statuses. For idempotent GETs of presigned
+   * URLs (CloudFetch), where a 4xx from storage does not mean the request was wrong.
+   */
+  retryClientErrorsExcept?: ReadonlySet<number>;
+}
+
 export default interface IConnectionProvider {
   getThriftConnection(): Promise<any>;
 
@@ -18,5 +26,5 @@ export default interface IConnectionProvider {
 
   setHeaders(headers: HeadersInit): void;
 
-  getRetryPolicy(): Promise<IRetryPolicy<HttpTransactionDetails>>;
+  getRetryPolicy(options?: HttpRetryPolicyOptions): Promise<IRetryPolicy<HttpTransactionDetails>>;
 }

lib/contracts/IDBSQLClient.ts

@@ -70,6 +70,11 @@ export type ConnectionOptions = {
   port?: number;
   path: string;
   userAgentEntry?: string;
+  /**
+   * Request timeout in milliseconds. On the Thrift backend it is also the inactivity timeout
+   * of each CloudFetch download attempt: an attempt that receives nothing for this long is
+   * aborted and retried. Default 900000 (15 minutes).
+   */
   socketTimeout?: number;
   proxy?: ProxyOptions;
   enableMetricViewMetadata?: boolean;

lib/result/CloudFetchResultHandler.ts

@@ -1,4 +1,5 @@
 import fetch, { RequestInfo, RequestInit, Request, Response } from 'node-fetch';
+import Int64 from 'node-int64';
 import { TGetResultSetMetadataResp, TRowSet, TSparkArrowResultLink } from '../../thrift/TCLIService_types';
 import HiveDriverError from '../errors/HiveDriverError';
 import RetryError, { RetryErrorCode } from '../errors/RetryError';
@@ -20,6 +21,19 @@ type SettledDownload = { ok: true; batch: ArrowBatch } | { ok: false; error: unk
 // Largest delay `setTimeout` accepts; longer values fire immediately, so they are clamped.
 const MAX_TIMER_DELAY_MS = 2 ** 31 - 1;
 
+// Statuses that mean the presigned URL itself was rejected (expired or invalid signature; GCS
+// answers 404 for an expired one). Retrying the same URL can't succeed, so the link is refreshed.
+// Other 4xx responses are retried with the same URL.
+const LINK_REFRESH_STATUSES: ReadonlySet<number> = new Set([401, 403, 404]);
+
+interface LinkRefresher {
+  refreshLinks(startRowOffset: Int64, maxRows: Int64): Promise<Array<TSparkArrowResultLink>>;
+}
+
+function canRefreshLinks(source: unknown): source is LinkRefresher {
+  return typeof (source as Partial<LinkRefresher> | undefined)?.refreshLinks === 'function';
+}
+
 export default class CloudFetchResultHandler implements IResultsProvider<ArrowBatch> {
   private readonly context: IClientContext;
 
@@ -128,7 +142,22 @@ export default class CloudFetchResultHandler implements IResultsProvider<ArrowBa
     }
 
     const startTime = Date.now();
-    const { response, body } = await this.fetch(link.fileLink, { headers: link.httpHeaders });
+    const { retryMaxAttempts } = this.context.getConfig();
+    // One `retriesTimeout` deadline covers the download and all of its link refreshes
+    const deadlineAt = this.downloadDeadline();
+    let currentLink = link;
+    let { response, body } = await this.fetch(currentLink.fileLink, { headers: currentLink.httpHeaders }, deadlineAt);
+    // Each link refresh counts as one more attempt against `retryMaxAttempts`
+    for (let attempt = 1; LINK_REFRESH_STATUSES.has(response.status) && attempt < retryMaxAttempts; attempt += 1) {
+      // eslint-disable-next-line no-await-in-loop
+      const freshLink = await this.refreshLink(currentLink, response.status);
+      if (!freshLink) {
+        break;
+      }
+      currentLink = freshLink;
+      // eslint-disable-next-line no-await-in-loop
+      ({ response, body } = await this.fetch(currentLink.fileLink, { headers: currentLink.httpHeaders }, deadlineAt));
+    }
     if (!response.ok) {
       throw new Error(`CloudFetch HTTP error ${response.status} ${response.statusText}`);
     }
@@ -149,23 +178,51 @@ export default class CloudFetchResultHandler implements IResultsProvider<ArrowBa
     };
   }
 
-  private async fetch(url: RequestInfo, init?: RequestInit): Promise<{ response: Response; body?: ArrayBuffer }> {
+  /** Asks the server for a fresh link to the rows of `link`; `undefined` if there is none. */
+  private async refreshLink(link: TSparkArrowResultLink, status: number): Promise<TSparkArrowResultLink | undefined> {
+    if (!canRefreshLinks(this.source)) {
+      return undefined;
+    }
+    const offset = link.startRowOffset.toNumber();
+    this.context
+      .getLogger()
+      .log(LogLevel.debug, `CloudFetch download got HTTP ${status}; refreshing the link at row offset ${offset}`);
+    const links = await this.source.refreshLinks(link.startRowOffset, link.rowCount);
+    return links.find((candidate) => candidate.startRowOffset.toNumber() === offset);
+  }
+
+  // End (epoch ms) of a download's `retriesTimeout` deadline; `undefined` when it is off
+  private downloadDeadline(): number | undefined {
+    const { retriesTimeout } = this.context.getConfig();
+    return retriesTimeout > 0 ? Date.now() + retriesTimeout : undefined;
+  }
+
+  // `deadlineAt` is shared by a download and its link refreshes
+  private async fetch(
+    url: RequestInfo,
+    init?: RequestInit,
+    deadlineAt: number | undefined = this.downloadDeadline(),
+  ): Promise<{ response: Response; body?: ArrayBuffer }> {
     const connectionProvider = await this.context.getConnectionProvider();
     let requestUrl = url;
     if (typeof requestUrl !== 'string') {
       requestUrl = 'href' in requestUrl ? requestUrl.href : requestUrl.url;
     }
     const agent = await connectionProvider.getAgent(requestUrl);
-    const retryPolicy = await connectionProvider.getRetryPolicy();
-    const { retriesTimeout } = this.context.getConfig();
+    const retryPolicy = await connectionProvider.getRetryPolicy({ retryClientErrorsExcept: LINK_REFRESH_STATUSES });
+    const { socketTimeout } = this.context.getConfig();
+    // `socketTimeout` is an inactivity timeout for each attempt: an attempt that receives
+    // nothing for that long is aborted and retried like a network timeout.
+    const idleTimeoutMs = socketTimeout > 0 ? Math.min(socketTimeout, MAX_TIMER_DELAY_MS) : undefined;
 
     const requestConfig: RequestInit = { agent, ...init };
     // The retry policy checks `retriesTimeout` only between attempts, so a stalled attempt
     // (headers or body never arrive) would block forever. Enforce it as a deadline over the
     // whole download - every attempt plus the backoff between them - and abort the in-flight
     // request and any pending backoff sleep when it expires (#520). Non-positive values keep
     // the deadline off.
-    const deadlineMs = retriesTimeout > 0 ? Math.min(retriesTimeout, MAX_TIMER_DELAY_MS) : undefined;
+    const deadlineMs =
+      deadlineAt === undefined ? undefined : Math.min(Math.max(deadlineAt - Date.now(), 0), MAX_TIMER_DELAY_MS);
     // Aborted only once, by the deadline, so a single controller serves every attempt and backoff.
     const controller = deadlineMs === undefined ? undefined : new AbortController();
     let timeoutError: RetryError | undefined;
@@ -187,14 +244,33 @@ export default class CloudFetchResultHandler implements IResultsProvider<ArrowBa
           throw timeoutError;
         }
         downloaded = undefined;
+        // Aborted by the deadline or when the attempt goes idle
+        const attempt = new AbortController();
+        const abortAttempt = () => attempt.abort();
+        controller?.signal.addEventListener('abort', abortAttempt, { once: true });
+        let attemptTimedOut = false;
+        let idleTimer: NodeJS.Timeout | undefined;
+        const restartIdleTimer = () => {
+          if (idleTimeoutMs !== undefined) {
+            clearTimeout(idleTimer);
+            idleTimer = setTimeout(() => {
+              attemptTimedOut = true;
+              attempt.abort();
+            }, idleTimeoutMs);
+          }
+        };
+        restartIdleTimer();
         const request = new Request(url, {
           ...requestConfig,
-          signal: controller?.signal as unknown as RequestInit['signal'],
+          signal: attempt.signal as unknown as RequestInit['signal'],
         });
         try {
           const response = await fetch(request);
+          restartIdleTimer();
           if (response.ok) {
-            downloaded = await response.arrayBuffer();
+            const bodyRead = response.arrayBuffer();
+            response.body.on('data', restartIdleTimer);
+            downloaded = await bodyRead;
           } else {
             lastFailure = { request, response };
           }
@@ -203,8 +279,13 @@ export default class CloudFetchResultHandler implements IResultsProvider<ArrowBa
           if (timeoutError) {
             throw timeoutError;
           }
-          lastFailure = error;
-          throw error;
+          lastFailure = attemptTimedOut
+            ? Object.assign(new Error('CloudFetch download attempt timed out'), { code: 'ETIMEDOUT' })
+            : error;
+          throw lastFailure;
+        } finally {
+          clearTimeout(idleTimer);
+          controller?.signal.removeEventListener('abort', abortAttempt);
         }
       }, controller?.signal);
 

lib/result/RowSetProvider.ts

@@ -5,6 +5,7 @@ import {
   TFetchResultsResp,
   TOperationHandle,
   TRowSet,
+  TSparkArrowResultLink,
 } from '../../thrift/TCLIService_types';
 import Status from '../dto/Status';
 import IClientContext from '../contracts/IClientContext';
@@ -98,6 +99,21 @@ export default class RowSetProvider implements IResultsProvider<TRowSet | undefi
 
   private hasMoreRowsFlag?: boolean = undefined;
 
+  // Row offset after the last result link returned so far
+  private nextLinkRowOffset: number = 0;
+
+  // A link refresh may move the server's cursor, so after one the next page is requested by offset
+  private fetchByOffset: boolean = false;
+
+  // Serializes FetchResults calls, so a link refresh never races a page fetch on the cursor
+  private lastFetch: Promise<unknown> = Promise.resolve();
+
+  private serializeFetch<T>(fetchResults: () => Promise<T>): Promise<T> {
+    const result = this.lastFetch.then(fetchResults, fetchResults);
+    this.lastFetch = result.catch(() => undefined);
+    return result;
+  }
+
   private get hasMoreRows(): boolean {
     // `hasMoreRowsFlag` is populated only after fetching the first row set.
     // Prior to that, we use a `operationHandle.hasResultSet` flag which
@@ -129,9 +145,38 @@ export default class RowSetProvider implements IResultsProvider<TRowSet | undefi
     Status.assert(response.status);
     this.fetchOrientation = TFetchOrientation.FETCH_NEXT;
     this.hasMoreRowsFlag = checkIfOperationHasMoreRows(response);
+    response.results?.resultLinks?.forEach((link) => {
+      this.nextLinkRowOffset = Math.max(
+        this.nextLinkRowOffset,
+        link.startRowOffset.toNumber() + link.rowCount.toNumber(),
+      );
+    });
     return response.results;
   }
 
+  /**
+   * Fetches fresh result links starting at `startRowOffset`, e.g. to replace a presigned
+   * URL that cloud storage rejected. Returns no links when the operation is already closed.
+   */
+  public async refreshLinks(startRowOffset: Int64, maxRows: Int64): Promise<Array<TSparkArrowResultLink>> {
+    if (this.returnOnlyPrefetchedResults) {
+      return [];
+    }
+    const driver = await this.context.getDriver();
+    const response = await this.serializeFetch(async () => {
+      this.fetchByOffset = true;
+      return driver.fetchResults({
+        operationHandle: this.operationHandle,
+        orientation: TFetchOrientation.FETCH_NEXT,
+        maxRows,
+        fetchType: FetchType.Data,
+        startRowOffset,
+      });
+    });
+    Status.assert(response.status);
+    return response.results?.resultLinks ?? [];
+  }
+
   public async fetchNext({ limit }: ResultsProviderFetchNextOptions) {
     const prefetchedResponse = this.prefetchedResults.shift();
     if (prefetchedResponse) {
@@ -150,12 +195,15 @@ export default class RowSetProvider implements IResultsProvider<TRowSet | undefi
 
     const driver = await this.context.getDriver();
     const startTime = Date.now();
-    const response = await driver.fetchResults({
-      operationHandle: this.operationHandle,
-      orientation: this.fetchOrientation,
-      maxRows: new Int64(limit),
-      fetchType: FetchType.Data,
-    });
+    const response = await this.serializeFetch(() =>
+      driver.fetchResults({
+        operationHandle: this.operationHandle,
+        orientation: this.fetchOrientation,
+        maxRows: new Int64(limit),
+        fetchType: FetchType.Data,
+        ...(this.fetchByOffset ? { startRowOffset: new Int64(this.nextLinkRowOffset) } : {}),
+      }),
+    );
     const latencyMs = Date.now() - startTime;
 
     this.emitChunkEvent(latencyMs, response);

tests/unit/.stubs/ConnectionProviderStub.ts

@@ -1,6 +1,9 @@
 import http from 'http';
 import { HeadersInit } from 'node-fetch';
-import IConnectionProvider, { HttpTransactionDetails } from '../../../lib/connection/contracts/IConnectionProvider';
+import IConnectionProvider, {
+  HttpRetryPolicyOptions,
+  HttpTransactionDetails,
+} from '../../../lib/connection/contracts/IConnectionProvider';
 import IRetryPolicy from '../../../lib/connection/contracts/IRetryPolicy';
 import NullRetryPolicy from '../../../lib/connection/connections/NullRetryPolicy';
 
@@ -19,7 +22,8 @@ export default class ConnectionProviderStub implements IConnectionProvider {
     this.headers = headers;
   }
 
-  public async getRetryPolicy(): Promise<IRetryPolicy<HttpTransactionDetails>> {
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  public async getRetryPolicy(options?: HttpRetryPolicyOptions): Promise<IRetryPolicy<HttpTransactionDetails>> {
     return new NullRetryPolicy();
   }
 }