fix(telemetry): address review highs — refcount leak, FIFO fallthrough, public types, README, test gaps

- DBSQLClient.initializeTelemetry: release the per-host refcount on the
  catch path so a throw between getOrCreateClient and the success log
  doesn't leak a TelemetryClient (with its flush timer / exporter / FFCache)
  for the lifetime of the process on long-running supervisors.
- TelemetryClient.getClient/getDriver: walk the FIFO with try/catch
  fallthrough to mirror getConnectionProvider, so a closed-but-not-yet-
  released head context doesn't take the whole shared pool down.
  getAuthProvider: return the first defined entry from the FIFO.
- lib/index.ts: re-export TelemetryEventType, DEFAULT_TELEMETRY_CONFIG,
  and the consumer-facing telemetry payload types so SDK users don't
  need to deep-import for type-checked event/metric handling.
- README: add a Telemetry section covering what's collected, the three
  opt-out paths (env var, programmatic, server-side feature flag), the
  tunable knobs, and the await-close requirement for short-lived
  processes.
- MetricsAggregator.test.ts: cover chunk-timing aggregation
  (initial=first-positive, slowest=max, sum=running) and CONNECTION_CLOSE
  → DELETE_SESSION emission. Both were acknowledged coverage gaps.
- TelemetryEventEmitter.test.ts: cover emitConnectionClose — emission
  shape, disabled-flag suppression, and listener-exception swallow.

Co-authored-by: Isaac

Author: samikshya-db

README.md

@@ -51,6 +51,70 @@ client
   });
 ```
 
+## Telemetry
+
+Starting with version 1.13, the driver collects telemetry — connection,
+statement, and CloudFetch chunk metrics, plus error events with redacted
+stack traces — to help Databricks improve driver performance and
+reliability. **Telemetry is enabled by default and gated by a server-side
+feature flag**: events are emitted only when the workspace's feature flag
+is on. No SQL text, parameter values, or row data are ever included.
+
+### What's collected
+
+- Connection lifecycle (`CREATE_SESSION`, `DELETE_SESSION`) with latency.
+- Statement lifecycle (`STATEMENT_START`, `STATEMENT_COMPLETE`) with
+  execution latency, operation type, and result format.
+- CloudFetch chunk timings and byte counts.
+- Error events with redacted stack traces (Bearer/JWT tokens, OAuth
+  secrets, home-directory paths, and Databricks PATs are stripped before
+  emission).
+
+See `TelemetryEvent` and `TelemetryMetric` in the package exports for the
+exact payload shapes.
+
+### Opting out
+
+Three independent ways to disable telemetry, in order of precedence:
+
+1. **Environment variable** — set `DATABRICKS_TELEMETRY_DISABLED` to one
+   of `1`, `true`, `yes`, or `on` (case-insensitive). Other values
+   (empty, `0`, `false`, `off`, `no`) are ignored, leaving the runtime
+   config in charge.
+2. **Programmatic** — pass `telemetryEnabled: false` to `connect()`:
+   ```javascript
+   await client.connect({
+     host,
+     path,
+     token,
+     telemetryEnabled: false,
+   });
+   ```
+3. **Server-side** — Databricks-managed feature flag; if disabled for
+   your workspace, the driver does not emit telemetry regardless of
+   client config.
+
+### Tuning
+
+If you keep telemetry on, the following knobs are available on
+`ConnectionOptions` (see JSDoc on `IDBSQLClient.ts` for defaults and
+units):
+
+- `telemetryAuthenticatedExport` — set to `false` to ship reduced
+  payloads (no statement/session correlation IDs, generic User-Agent)
+  via the unauthenticated endpoint.
+- `telemetryBatchSize`, `telemetryFlushIntervalMs`, `telemetryMaxRetries`
+  — batching and retry tuning.
+- `telemetryCircuitBreakerThreshold`, `telemetryCircuitBreakerTimeout` —
+  circuit-breaker tuning for the export endpoint.
+- `telemetryCloseTimeoutMs` — bound on `await client.close()` waiting for
+  the final flush.
+
+> **Note for short-lived processes**: always `await client.close()`
+> before `process.exit(0)` so the final batch is flushed. Without an
+> explicit close, the periodic flush timer is `unref()`'d to avoid
+> holding the event loop open, so any unflushed events are dropped.
+
 ## Run Tests
 
 ### Unit tests

lib/DBSQLClient.ts

@@ -404,8 +404,23 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
 
       this.logger.log(LogLevel.debug, 'Telemetry: enabled');
     } catch (error: any) {
-      // Swallow all telemetry initialization errors
-      this.logger.log(LogLevel.debug, `Telemetry initialization error: ${error.message}`);
+      // Swallow all telemetry initialization errors. If we acquired a refcount
+      // before the throw, release it — otherwise the per-host TelemetryClient
+      // (and its flush timer / exporter / FFCache) leaks for the lifetime of
+      // the process on long-running supervisors that retry-connect.
+      if (this.telemetryClient) {
+        try {
+          await TelemetryClientProvider.getInstance().releaseClient(this, this.host);
+        } catch (releaseError: any) {
+          this.logger.log(
+            LogLevel.debug,
+            `Telemetry release-after-init-failure error: ${releaseError?.message ?? releaseError}`,
+          );
+        }
+        this.telemetryClient = undefined;
+        this.telemetryEmitter = undefined;
+      }
+      this.logger.log(LogLevel.debug, `Telemetry initialization error: ${error?.message ?? error}`);
     }
   }
 

lib/index.ts

@@ -26,6 +26,18 @@ export type { default as ITokenProvider } from './connection/auth/tokenProvider/
 export { CircuitBreakerOpenError, CIRCUIT_BREAKER_OPEN_CODE } from './telemetry/CircuitBreaker';
 export { TelemetryTerminalError } from './telemetry/DatabricksTelemetryExporter';
 
+// Telemetry event/metric/config shapes for consumers that want to inspect
+// telemetry payloads or pre-validate config. The emitter, aggregator, and
+// per-host client are deliberately not re-exported — they are internal.
+export { TelemetryEventType, DEFAULT_TELEMETRY_CONFIG } from './telemetry/types';
+export type {
+  TelemetryEvent,
+  TelemetryMetric,
+  TelemetryConfiguration,
+  StatementMetrics,
+  DriverConfiguration,
+} from './telemetry/types';
+
 export const auth = {
   PlainHttpAuthentication,
   // Token provider classes for custom authentication

lib/telemetry/TelemetryClient.ts

@@ -175,21 +175,43 @@ class TelemetryClient implements IClientContext {
   }
 
   async getClient(): Promise<IThriftClient> {
-    if (this.contexts.length === 0) {
-      throw new Error(`TelemetryClient: no client available for host ${this.host}`);
+    let lastErr: unknown;
+    for (const ctx of this.contexts) {
+      try {
+        // eslint-disable-next-line no-await-in-loop
+        return await ctx.getClient();
+      } catch (err) {
+        lastErr = err;
+      }
     }
-    return this.contexts[0].getClient();
+    throw lastErr instanceof Error ? lastErr : new Error(`TelemetryClient: no client available for host ${this.host}`);
   }
 
   async getDriver(): Promise<IDriver> {
-    if (this.contexts.length === 0) {
-      throw new Error(`TelemetryClient: no driver available for host ${this.host}`);
+    let lastErr: unknown;
+    for (const ctx of this.contexts) {
+      try {
+        // eslint-disable-next-line no-await-in-loop
+        return await ctx.getDriver();
+      } catch (err) {
+        lastErr = err;
+      }
     }
-    return this.contexts[0].getDriver();
+    throw lastErr instanceof Error ? lastErr : new Error(`TelemetryClient: no driver available for host ${this.host}`);
   }
 
   getAuthProvider(): IAuthentication | undefined {
-    return this.authProviders[0];
+    // Walk the FIFO and return the first usable provider. A registered head
+    // whose underlying DBSQLClient has revoked credentials will surface as
+    // an `authenticate()` failure inside the exporter retry loop, but we
+    // can avoid that round-trip when the array contains a clearly-defined
+    // entry. Symmetric with `getConnectionProvider` fallthrough.
+    for (const provider of this.authProviders) {
+      if (provider !== undefined) {
+        return provider;
+      }
+    }
+    return undefined;
   }
 
   getTelemetryEmitter(): undefined {

tests/unit/telemetry/MetricsAggregator.test.ts

@@ -121,6 +121,79 @@ describe('MetricsAggregator', () => {
 
       expect(() => aggregator.completeStatement('unknown-stmt')).to.not.throw();
     });
+
+    it('aggregates chunk timing — initial is first-seen, slowest is max, sum accumulates', () => {
+      const context = new ClientContextStub();
+      const exporter = makeExporterStub();
+      const aggregator = new MetricsAggregator(context, exporter as any);
+
+      aggregator.processEvent(statementEvent(TelemetryEventType.STATEMENT_START));
+      aggregator.processEvent(statementEvent(TelemetryEventType.CLOUDFETCH_CHUNK, { latencyMs: 100, bytes: 10 }));
+      aggregator.processEvent(statementEvent(TelemetryEventType.CLOUDFETCH_CHUNK, { latencyMs: 250, bytes: 10 }));
+      aggregator.processEvent(statementEvent(TelemetryEventType.CLOUDFETCH_CHUNK, { latencyMs: 75, bytes: 10 }));
+
+      aggregator.completeStatement('stmt-1');
+      aggregator.flush();
+
+      const stmtMetric = exporter.export.firstCall.args[0][0];
+      expect(stmtMetric.chunkInitialLatencyMs).to.equal(100);
+      expect(stmtMetric.chunkSlowestLatencyMs).to.equal(250);
+      expect(stmtMetric.chunkSumLatencyMs).to.equal(425);
+    });
+
+    it('chunks with non-positive latency do not contribute to timing fields', () => {
+      const context = new ClientContextStub();
+      const exporter = makeExporterStub();
+      const aggregator = new MetricsAggregator(context, exporter as any);
+
+      aggregator.processEvent(statementEvent(TelemetryEventType.STATEMENT_START));
+      // latency=0 (cached/prefetched page) — must be ignored entirely.
+      aggregator.processEvent(statementEvent(TelemetryEventType.CLOUDFETCH_CHUNK, { latencyMs: 0, bytes: 5 }));
+      // latency undefined — emitter didn't set it, must be ignored.
+      aggregator.processEvent(statementEvent(TelemetryEventType.CLOUDFETCH_CHUNK, { bytes: 5 }));
+      // First *positive* latency wins for `initial`, even though earlier chunks already arrived.
+      aggregator.processEvent(statementEvent(TelemetryEventType.CLOUDFETCH_CHUNK, { latencyMs: 60, bytes: 5 }));
+
+      aggregator.completeStatement('stmt-1');
+      aggregator.flush();
+
+      const stmtMetric = exporter.export.firstCall.args[0][0];
+      expect(stmtMetric.chunkInitialLatencyMs).to.equal(60);
+      expect(stmtMetric.chunkSlowestLatencyMs).to.equal(60);
+      expect(stmtMetric.chunkSumLatencyMs).to.equal(60);
+      expect(stmtMetric.chunkCount).to.equal(3); // chunkCount counts all chunks regardless of latency
+    });
+  });
+
+  describe('processEvent() - CONNECTION_CLOSE', () => {
+    it('emits a DELETE_SESSION connection metric immediately', () => {
+      const context = new ClientContextStub();
+      const exporter = makeExporterStub();
+      const aggregator = new MetricsAggregator(context, exporter as any);
+
+      aggregator.processEvent(connectionEvent({ eventType: TelemetryEventType.CONNECTION_CLOSE, latencyMs: 42 }));
+      aggregator.flush();
+
+      expect(exporter.export.calledOnce).to.be.true;
+      const metric = exporter.export.firstCall.args[0][0];
+      expect(metric.metricType).to.equal('connection');
+      expect(metric.operationType).to.equal('DELETE_SESSION');
+      expect(metric.latencyMs).to.equal(42);
+    });
+
+    it('CONNECTION_OPEN and CONNECTION_CLOSE produce distinct operation types in the same batch', () => {
+      const context = new ClientContextStub();
+      const exporter = makeExporterStub();
+      const aggregator = new MetricsAggregator(context, exporter as any);
+
+      aggregator.processEvent(connectionEvent({ eventType: TelemetryEventType.CONNECTION_OPEN, latencyMs: 100 }));
+      aggregator.processEvent(connectionEvent({ eventType: TelemetryEventType.CONNECTION_CLOSE, latencyMs: 5 }));
+      aggregator.flush();
+
+      const batch = exporter.export.firstCall.args[0];
+      expect(batch).to.have.lengthOf(2);
+      expect(batch.map((m: any) => m.operationType)).to.deep.equal(['CREATE_SESSION', 'DELETE_SESSION']);
+    });
   });
 
   describe('processEvent() - error events', () => {

tests/unit/telemetry/TelemetryEventEmitter.test.ts

@@ -222,6 +222,60 @@ describe('TelemetryEventEmitter', () => {
     });
   });
 
+  describe('emitConnectionClose', () => {
+    it('emits connection.close with sessionId and latencyMs', (done) => {
+      emitter.on(TelemetryEventType.CONNECTION_CLOSE, (event: TelemetryEvent) => {
+        expect(event.eventType).to.equal(TelemetryEventType.CONNECTION_CLOSE);
+        expect(event.sessionId).to.equal('session-close-1');
+        expect(event.latencyMs).to.equal(7);
+        expect(event.timestamp).to.be.a('number');
+        done();
+      });
+
+      emitter.emitConnectionClose({ sessionId: 'session-close-1', latencyMs: 7 });
+    });
+
+    it('does not emit when telemetry is disabled', () => {
+      const disabledContext = {
+        getLogger: () => logger,
+        getConfig: () => ({
+          telemetryEnabled: false,
+          directResultsDefaultMaxRows: 10000,
+          fetchChunkDefaultMaxRows: 100000,
+          socketTimeout: 900000,
+          retryMaxAttempts: 30,
+          retriesTimeout: 900000,
+          retryDelayMin: 1000,
+          retryDelayMax: 30000,
+          useCloudFetch: true,
+          cloudFetchConcurrentDownloads: 10,
+          cloudFetchSpeedThresholdMBps: 0,
+          useLZ4Compression: true,
+        }),
+      } as any;
+
+      const disabledEmitter = new TelemetryEventEmitter(disabledContext);
+      let emitted = false;
+      disabledEmitter.on(TelemetryEventType.CONNECTION_CLOSE, () => {
+        emitted = true;
+      });
+
+      disabledEmitter.emitConnectionClose({ sessionId: 's', latencyMs: 1 });
+      expect(emitted).to.be.false;
+    });
+
+    it('swallows exceptions from listeners and logs at debug level', () => {
+      emitter.on(TelemetryEventType.CONNECTION_CLOSE, () => {
+        throw new Error('listener boom');
+      });
+
+      expect(() => emitter.emitConnectionClose({ sessionId: 's', latencyMs: 1 })).to.not.throw();
+      const logStub = logger.log as sinon.SinonStub;
+      const debugCalls = logStub.getCalls().filter((c) => c.args[0] === LogLevel.debug);
+      expect(debugCalls.some((c) => /connection close/i.test(c.args[1]))).to.be.true;
+    });
+  });
+
   describe('emitStatementStart', () => {
     it('should emit statement.start event with correct data', (done) => {
       emitter.on(TelemetryEventType.STATEMENT_START, (event: TelemetryEvent) => {