fix(telemetry): iter-2 review fixes — rebase regressions, type-safe wiring, error telemetry

Rebase the regressed exporter/aggregator/feature-flag-cache on main's
hardened versions and re-apply only the genuinely new functionality
(CONNECTION_CLOSE event, chunk-timing aggregation) on top. Closes the
critical findings from the multi-reviewer audit:

  - SSRF guard, redactSensitive, sanitizeProcessName, hasAuthorization,
    auth-missing warn-once — all restored via main's telemetryUtils.
  - MetricsAggregator memory bounds (maxPendingMetrics with error-preferring
    drop, maxErrorsPerStatement, statementTtlMs eviction) restored.
  - FeatureFlagCache in-flight fetch dedup and TTL clamp [60s, 3600s]
    restored; lib/telemetry/urlUtils.ts deleted.
  - close() now properly awaits aggregator drain — fixes the close()/flush
    race that PR #362 already fixed once.
  - Driver version reads lib/version.ts via buildUserAgentString instead
    of hardcoded '1.0.0'; uuidv4() restored in place of Math.random().
  - TelemetryTerminalError re-exported from lib/index.ts.

Type-safe wiring:

  - Added optional getTelemetryEmitter() / getTelemetryAggregator() to
    IClientContext; removed all 7 `(this.context as any)` casts.
  - Six copy-pasted event listeners in DBSQLClient.initializeTelemetry
    collapsed into one `Object.values(TelemetryEventType)` loop — closes
    the listener-name mismatch that silently dropped error events.
  - mapAuthType now covers all 6 authType values instead of defaulting
    everything to 'pat'.

TelemetryClient now owns the host-scoped resources:

  - TelemetryClientProvider is a process-wide singleton (getInstance()).
  - TelemetryClient owns DatabricksTelemetryExporter, MetricsAggregator,
    CircuitBreakerRegistry, and FeatureFlagCache for its host. Implements
    IClientContext itself so the owned components have a stable context
    that survives any single DBSQLClient's close.
  - DBSQLClient instances on the same host share the breaker counters,
    feature-flag cache, exporter, and HTTP batches. Fixes the per-instance
    breaker-fragmentation noted in iter-2 architecture review.
  - Each DBSQLClient still holds its own TelemetryEventEmitter (respects
    per-client telemetryEnabled); emitters bridge into the shared aggregator.
  - Exporter falls back to context.getAuthProvider() when no explicit auth
    provider is passed, so the shared exporter resolves auth through the
    TelemetryClient's FIFO of registered DBSQLClients.

Error telemetry wired across operation entry points:

  - Re-added emitErrorEvent(error) on DBSQLOperation; uses
    ExceptionClassifier.isTerminal() to classify.
  - fetchChunk, cancel, close, getMetadata wrap their bodies in try/catch
    that calls emitErrorEvent before re-throwing. Verified end-to-end
    against a real Azure Databricks workspace: failed query produces
    STATEMENT_COMPLETE + ERROR (with redacted stack) on the wire.
  - Removed the await getMetadata() call from emitStatementComplete —
    eliminates the extra Thrift RPC on every close (F19) AND prevents
    spurious error telemetry from getMetadata's wrapper firing during
    close-cleanup of an already-failed operation.

Other:

  - Iterating Map.keys() while mutating made safe via snapshot in close().
  - STATEMENT_COMPLETE no longer zeroes accumulated chunk metrics when
    the emit doesn't supply them (matches sibling-field guards).
  - Tests for the rebased modules restored from main; provider tests
    updated for the singleton API; deleted unused TelemetryExporterStub.

484 unit tests passing. Diff vs main: ~+2110/-383, down from the
original PR's +3640/-1173.

Co-authored-by: Isaac

Author: samikshya-db

lib/DBSQLClient.ts

@@ -32,13 +32,10 @@ import IDBSQLLogger, { LogLevel } from './contracts/IDBSQLLogger';
 import DBSQLLogger from './DBSQLLogger';
 import CloseableCollection from './utils/CloseableCollection';
 import IConnectionProvider from './connection/contracts/IConnectionProvider';
-import FeatureFlagCache from './telemetry/FeatureFlagCache';
+import TelemetryClient from './telemetry/TelemetryClient';
 import TelemetryClientProvider from './telemetry/TelemetryClientProvider';
 import TelemetryEventEmitter from './telemetry/TelemetryEventEmitter';
-import MetricsAggregator from './telemetry/MetricsAggregator';
-import DatabricksTelemetryExporter from './telemetry/DatabricksTelemetryExporter';
-import { CircuitBreakerRegistry } from './telemetry/CircuitBreaker';
-import { DriverConfiguration, DRIVER_NAME } from './telemetry/types';
+import { DriverConfiguration, DRIVER_NAME, TelemetryEventType } from './telemetry/types';
 import driverVersion from './version';
 
 function prependSlash(str: string): string {
@@ -84,23 +81,21 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
 
   private readonly sessions = new CloseableCollection<DBSQLSession>();
 
-  // Telemetry components (instance-based, NOT singletons)
+  // Telemetry components — `telemetryClient` is the shared per-host owner
+  // (process-wide via TelemetryClientProvider). The exporter, aggregator,
+  // circuit-breaker registry and feature-flag cache live on it. Each
+  // DBSQLClient still owns its own `telemetryEmitter` so it respects its
+  // own `telemetryEnabled` flag.
   private host?: string;
 
   private httpPath?: string;
 
   private authType?: string;
 
-  private featureFlagCache?: FeatureFlagCache;
-
-  private telemetryClientProvider?: TelemetryClientProvider;
+  private telemetryClient?: TelemetryClient;
 
   private telemetryEmitter?: TelemetryEventEmitter;
 
-  private telemetryAggregator?: MetricsAggregator;
-
-  private circuitBreakerRegistry?: CircuitBreakerRegistry;
-
   private static getDefaultLogger(): IDBSQLLogger {
     if (!this.defaultLogger) {
       this.defaultLogger = new DBSQLLogger();
@@ -298,19 +293,23 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
    * Distinguishes between U2M and M2M OAuth flows.
    */
   private mapAuthType(options: ConnectionOptions): string {
-    if (options.authType === 'databricks-oauth') {
-      // Check if M2M (has client secret) or U2M (no client secret)
-      return options.oauthClientSecret === undefined
-        ? 'external-browser' // U2M OAuth (User-to-Machine)
-        : 'oauth-m2m'; // M2M OAuth (Machine-to-Machine)
-    }
-
-    if (options.authType === 'custom') {
-      return 'custom'; // Custom auth provider
+    switch (options.authType) {
+      case 'databricks-oauth':
+        return options.oauthClientSecret === undefined ? 'external-browser' : 'oauth-m2m';
+      case 'custom':
+        return 'custom';
+      case 'token-provider':
+        return 'token-provider';
+      case 'external-token':
+        return 'external-token';
+      case 'static-token':
+        return 'static-token';
+      case 'access-token':
+      case undefined:
+        return 'pat';
+      default:
+        return 'unknown';
     }
-
-    // 'access-token' or undefined
-    return 'pat'; // Personal Access Token
   }
 
   /**
@@ -373,80 +372,34 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
     }
 
     try {
-      // Create circuit breaker registry (shared by feature flags and telemetry)
-      this.circuitBreakerRegistry = new CircuitBreakerRegistry(this);
-
-      // Create feature flag cache instance with circuit breaker protection
-      this.featureFlagCache = new FeatureFlagCache(this, this.circuitBreakerRegistry);
-      this.featureFlagCache.getOrCreateContext(this.host);
+      // Acquire (or create) the per-host TelemetryClient from the
+      // process-wide provider. The shared client owns the circuit-breaker
+      // registry, feature-flag cache, exporter, and aggregator. Multiple
+      // DBSQLClient instances on the same host share these resources so
+      // breaker counters and HTTP batches don't fragment per-instance.
+      this.telemetryClient = TelemetryClientProvider.getInstance().getOrCreateClient(this, this.host);
 
-      // Check if telemetry enabled via feature flag
-      const enabled = await this.featureFlagCache.isTelemetryEnabled(this.host);
+      // Use the shared feature-flag cache (registered in the previous step).
+      const enabled = await this.telemetryClient.getFeatureFlagCache().isTelemetryEnabled(this.host);
 
       if (!enabled) {
+        // Release our refcount immediately; we won't be emitting.
+        await TelemetryClientProvider.getInstance().releaseClient(this, this.host);
+        this.telemetryClient = undefined;
         this.logger.log(LogLevel.debug, 'Telemetry: disabled');
         return;
       }
 
-      // Create telemetry components (all instance-based)
-      this.telemetryClientProvider = new TelemetryClientProvider(this);
+      // Each DBSQLClient still owns its own emitter so it respects its own
+      // `telemetryEnabled` flag and feature-flag result. All emitters bridge
+      // into the SHARED aggregator on the TelemetryClient.
       this.telemetryEmitter = new TelemetryEventEmitter(this);
-
-      // Get or create telemetry client for this host (increments refCount)
-      this.telemetryClientProvider.getOrCreateClient(this.host);
-
-      // Create telemetry exporter with shared circuit breaker registry
-      const exporter = new DatabricksTelemetryExporter(this, this.host, this.circuitBreakerRegistry);
-      this.telemetryAggregator = new MetricsAggregator(this, exporter);
-
-      // Wire up event listeners
-      this.telemetryEmitter.on('connection.open', (event) => {
-        try {
-          this.telemetryAggregator?.processEvent(event);
-        } catch (error: any) {
-          this.logger.log(LogLevel.debug, `Error processing connection.open event: ${error.message}`);
-        }
-      });
-
-      this.telemetryEmitter.on('connection.close', (event) => {
-        try {
-          this.telemetryAggregator?.processEvent(event);
-        } catch (error: any) {
-          this.logger.log(LogLevel.debug, `Error processing connection.close event: ${error.message}`);
-        }
-      });
-
-      this.telemetryEmitter.on('statement.start', (event) => {
-        try {
-          this.telemetryAggregator?.processEvent(event);
-        } catch (error: any) {
-          this.logger.log(LogLevel.debug, `Error processing statement.start event: ${error.message}`);
-        }
-      });
-
-      this.telemetryEmitter.on('statement.complete', (event) => {
-        try {
-          this.telemetryAggregator?.processEvent(event);
-        } catch (error: any) {
-          this.logger.log(LogLevel.debug, `Error processing statement.complete event: ${error.message}`);
-        }
-      });
-
-      this.telemetryEmitter.on('cloudfetch.chunk', (event) => {
-        try {
-          this.telemetryAggregator?.processEvent(event);
-        } catch (error: any) {
-          this.logger.log(LogLevel.debug, `Error processing cloudfetch.chunk event: ${error.message}`);
-        }
-      });
-
-      this.telemetryEmitter.on('error', (event) => {
-        try {
-          this.telemetryAggregator?.processEvent(event);
-        } catch (error: any) {
-          this.logger.log(LogLevel.debug, `Error processing error event: ${error.message}`);
-        }
-      });
+      const sharedAggregator = this.telemetryClient.getAggregator();
+      for (const eventType of Object.values(TelemetryEventType)) {
+        this.telemetryEmitter.on(eventType, (event) => {
+          sharedAggregator.processEvent(event);
+        });
+      }
 
       this.logger.log(LogLevel.debug, 'Telemetry: enabled');
     } catch (error: any) {
@@ -613,27 +566,17 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
   public async close(): Promise<void> {
     await this.sessions.closeAll();
 
-    // Cleanup telemetry
-    if (this.host) {
+    // Cleanup telemetry. Releasing our refcount on the shared TelemetryClient
+    // is awaited because the underlying close() drains the final HTTP POST —
+    // a caller doing `await client.close(); process.exit(0)` would otherwise
+    // truncate the in-flight request when this is the last refcount holder.
+    if (this.host && this.telemetryClient) {
       try {
-        // Step 1: Close aggregator (stops timer, completes statements, final flush)
-        if (this.telemetryAggregator) {
-          this.telemetryAggregator.close();
-        }
-
-        // Step 2: Release telemetry client (decrements ref count, closes if last)
-        if (this.telemetryClientProvider) {
-          await this.telemetryClientProvider.releaseClient(this.host);
-        }
-
-        // Step 3: Release feature flag context (decrements ref count)
-        if (this.featureFlagCache) {
-          this.featureFlagCache.releaseContext(this.host);
-        }
+        await TelemetryClientProvider.getInstance().releaseClient(this, this.host);
       } catch (error: any) {
-        // Swallow all telemetry cleanup errors
         this.logger.log(LogLevel.debug, `Telemetry cleanup error: ${error.message}`);
       }
+      this.telemetryClient = undefined;
     }
 
     this.client = undefined;
@@ -687,4 +630,14 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
   public getAuthProvider(): IAuthentication | undefined {
     return this.authProvider;
   }
+
+  /** @internal */
+  public getTelemetryEmitter(): TelemetryEventEmitter | undefined {
+    return this.telemetryEmitter;
+  }
+
+  /** @internal */
+  public getTelemetryAggregator() {
+    return this.telemetryClient?.getAggregator();
+  }
 }

lib/DBSQLOperation.ts

@@ -34,8 +34,8 @@ import { definedOrError } from './utils';
 import { OperationChunksIterator, OperationRowsIterator } from './utils/OperationIterator';
 import HiveDriverError from './errors/HiveDriverError';
 import IClientContext from './contracts/IClientContext';
-import ExceptionClassifier from './telemetry/ExceptionClassifier';
 import { mapOperationTypeToTelemetryType, mapResultFormatToTelemetryType } from './telemetry/telemetryTypeMappers';
+import ExceptionClassifier from './telemetry/ExceptionClassifier';
 
 interface DBSQLOperationConstructorOptions {
   handle: TOperationHandle;
@@ -181,6 +181,15 @@ export default class DBSQLOperation implements IOperation {
    * const result = await queryOperation.fetchChunk({maxRows: 1000});
    */
   public async fetchChunk(options?: FetchOptions): Promise<Array<object>> {
+    try {
+      return await this.fetchChunkInternal(options);
+    } catch (err: any) {
+      this.emitErrorEvent(err);
+      throw err;
+    }
+  }
+
+  private async fetchChunkInternal(options?: FetchOptions): Promise<Array<object>> {
     await this.failIfClosed();
 
     if (!this.operationHandle.hasResultSet) {
@@ -257,6 +266,15 @@ export default class DBSQLOperation implements IOperation {
    * @throws {StatusError}
    */
   public async cancel(): Promise<Status> {
+    try {
+      return await this.cancelInternal();
+    } catch (err: any) {
+      this.emitErrorEvent(err);
+      throw err;
+    }
+  }
+
+  private async cancelInternal(): Promise<Status> {
     if (this.closed || this.cancelled) {
       return Status.success();
     }
@@ -281,6 +299,15 @@ export default class DBSQLOperation implements IOperation {
    * @throws {StatusError}
    */
   public async close(): Promise<Status> {
+    try {
+      return await this.closeInternal();
+    } catch (err: any) {
+      this.emitErrorEvent(err);
+      throw err;
+    }
+  }
+
+  private async closeInternal(): Promise<Status> {
     if (this.closed || this.cancelled) {
       return Status.success();
     }
@@ -341,9 +368,14 @@ export default class DBSQLOperation implements IOperation {
   }
 
   public async getMetadata(): Promise<TGetResultSetMetadataResp> {
-    await this.failIfClosed();
-    await this.waitUntilReady();
-    return this.fetchMetadata();
+    try {
+      await this.failIfClosed();
+      await this.waitUntilReady();
+      return await this.fetchMetadata();
+    } catch (err: any) {
+      this.emitErrorEvent(err);
+      throw err;
+    }
   }
 
   private async failIfClosed(): Promise<void> {
@@ -509,7 +541,7 @@ export default class DBSQLOperation implements IOperation {
    */
   private emitStatementStart(): void {
     try {
-      const { telemetryEmitter } = this.context as any;
+      const telemetryEmitter = this.context.getTelemetryEmitter?.();
       if (!telemetryEmitter) {
         return;
       }
@@ -530,23 +562,19 @@ export default class DBSQLOperation implements IOperation {
    */
   private async emitStatementComplete(): Promise<void> {
     try {
-      const { telemetryEmitter } = this.context as any;
-      const { telemetryAggregator } = this.context as any;
+      const telemetryEmitter = this.context.getTelemetryEmitter?.();
+      const telemetryAggregator = this.context.getTelemetryAggregator?.();
       if (!telemetryEmitter || !telemetryAggregator) {
         return;
       }
 
-      // Fetch metadata if not already fetched to get result format
-      let resultFormat: string | undefined;
-      try {
-        if (!this.metadata && !this.cancelled) {
-          await this.getMetadata();
-        }
-        resultFormat = mapResultFormatToTelemetryType(this.metadata?.resultFormat);
-      } catch (error) {
-        // If metadata fetch fails, continue without it
-        resultFormat = undefined;
-      }
+      // Use whatever metadata was already fetched by the result-handling
+      // path. Do NOT trigger a `getMetadata()` here — that issues a Thrift
+      // RPC on every close (doubles close latency for short DDL/DML) AND
+      // throws if the operation is already in an error/closed state, which
+      // would then fire spurious error telemetry from `getMetadata`'s error
+      // wrapper.
+      const resultFormat = mapResultFormatToTelemetryType(this.metadata?.resultFormat);
 
       const latencyMs = Date.now() - this.startTime;
 
@@ -566,25 +594,24 @@ export default class DBSQLOperation implements IOperation {
   }
 
   /**
-   * Emit error telemetry event with terminal classification.
-   * CRITICAL: All exceptions swallowed and logged at LogLevel.debug ONLY.
+   * Emit a telemetry error event for an exception thrown by an operation.
+   * Terminal errors (per `ExceptionClassifier`) trigger an immediate flush
+   * in the aggregator; retryable errors are buffered until the statement
+   * completes. All exceptions from this method itself are swallowed at
+   * debug level — telemetry must never break the driver.
    */
   private emitErrorEvent(error: Error): void {
     try {
-      const { telemetryEmitter } = this.context as any;
-      if (!telemetryEmitter) {
-        return;
-      }
-
-      // Classify the exception
-      const isTerminal = ExceptionClassifier.isTerminal(error);
+      const telemetryEmitter = this.context.getTelemetryEmitter?.();
+      if (!telemetryEmitter) return;
 
       telemetryEmitter.emitError({
         statementId: this.id,
         sessionId: this.sessionId,
         errorName: error.name || 'Error',
         errorMessage: error.message || 'Unknown error',
-        isTerminal,
+        errorStack: error.stack,
+        isTerminal: ExceptionClassifier.isTerminal(error),
       });
     } catch (emitError: any) {
       this.context.getLogger().log(LogLevel.debug, `Error emitting error event: ${emitError.message}`);

lib/DBSQLSession.ts

@@ -599,7 +599,7 @@ export default class DBSQLSession implements IDBSQLSession {
 
     // Emit connection close telemetry
     const closeLatency = Date.now() - this.openTime;
-    const { telemetryEmitter } = this.context as any;
+    const telemetryEmitter = this.context.getTelemetryEmitter?.();
     if (telemetryEmitter) {
       telemetryEmitter.emitConnectionClose({
         sessionId: this.id,