[3/7] Telemetry Client Management: TelemetryClient and Provider (#326)

samikshya-db · claude · web-flow · commit de440ef6586c · 2026-04-22T10:39:18.000+05:30
* Add telemetry infrastructure: CircuitBreaker and FeatureFlagCache This is part 2 of 7 in the telemetry implementation stack. Components: - CircuitBreaker: Per-host endpoint protection with state management - FeatureFlagCache: Per-host feature flag caching with reference counting - CircuitBreakerRegistry: Manages circuit breakers per host Circuit Breaker: - States: CLOSED (normal), OPEN (failing), HALF_OPEN (testing recovery) - Default: 5 failures trigger OPEN, 60s timeout, 2 successes to CLOSE - Per-host isolation prevents cascade failures - All state transitions logged at debug level Feature Flag Cache: - Per-host caching with 15-minute TTL - Reference counting for connection lifecycle management - Automatic cache expiration and refetch - Context removed when refCount reaches zero Testing: - 32 comprehensive unit tests for CircuitBreaker - 29 comprehensive unit tests for FeatureFlagCache - 100% function coverage, >80% line/branch coverage - CircuitBreakerStub for testing other components Dependencies: - Builds on [1/7] Types and Exception Classifier * Add authentication support for REST API calls Implements getAuthHeaders() method for authenticated REST API requests: - Added getAuthHeaders() to IClientContext interface - Implemented in DBSQLClient using authProvider.authenticate() - Updated FeatureFlagCache to fetch from connector-service API with auth - Added driver version support for version-specific feature flags - Replaced placeholder implementation with actual REST API calls Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com> * Fix feature flag and telemetry export endpoints - Change feature flag endpoint to use NODEJS client type - Fix telemetry endpoints to /telemetry-ext and /telemetry-unauth - Update payload to match proto with system_configuration - Add shared buildUrl utility for protocol handling * Match JDBC telemetry payload format - Change payload structure to match JDBC: uploadTime, items, protoLogs - protoLogs contains JSON-stringified TelemetryFrontendLog objects - Remove workspace_id (JDBC doesn't populate it) - Remove debug logs added during testing * Fix lint errors - Fix import order in FeatureFlagCache - Replace require() with import for driverVersion - Fix variable shadowing - Disable prefer-default-export for urlUtils * Add telemetry infrastructure: CircuitBreaker and FeatureFlagCache This is part 2 of 7 in the telemetry implementation stack. Components: - CircuitBreaker: Per-host endpoint protection with state management - FeatureFlagCache: Per-host feature flag caching with reference counting - CircuitBreakerRegistry: Manages circuit breakers per host Circuit Breaker: - States: CLOSED (normal), OPEN (failing), HALF_OPEN (testing recovery) - Default: 5 failures trigger OPEN, 60s timeout, 2 successes to CLOSE - Per-host isolation prevents cascade failures - All state transitions logged at debug level Feature Flag Cache: - Per-host caching with 15-minute TTL - Reference counting for connection lifecycle management - Automatic cache expiration and refetch - Context removed when refCount reaches zero Testing: - 32 comprehensive unit tests for CircuitBreaker - 29 comprehensive unit tests for FeatureFlagCache - 100% function coverage, >80% line/branch coverage - CircuitBreakerStub for testing other components Dependencies: - Builds on [1/7] Types and Exception Classifier * Add telemetry client management: TelemetryClient and Provider This is part 3 of 7 in the telemetry implementation stack. Components: - TelemetryClient: HTTP client for telemetry export per host - TelemetryClientProvider: Manages per-host client lifecycle with reference counting TelemetryClient: - Placeholder HTTP client for telemetry export - Per-host isolation for connection pooling - Lifecycle management (open/close) - Ready for future HTTP implementation TelemetryClientProvider: - Reference counting tracks connections per host - Automatically creates clients on first connection - Closes and removes clients when refCount reaches zero - Thread-safe per-host management Design Pattern: - Follows JDBC driver pattern for resource management - One client per host, shared across connections - Efficient resource utilization - Clean lifecycle management Testing: - 31 comprehensive unit tests for TelemetryClient - 31 comprehensive unit tests for TelemetryClientProvider - 100% function coverage, >80% line/branch coverage - Tests verify reference counting and lifecycle Dependencies: - Builds on [1/7] Types and [2/7] Infrastructure * Add authentication support for REST API calls Implements getAuthHeaders() method for authenticated REST API requests: - Added getAuthHeaders() to IClientContext interface - Implemented in DBSQLClient using authProvider.authenticate() - Updated FeatureFlagCache to fetch from connector-service API with auth - Added driver version support for version-specific feature flags - Replaced placeholder implementation with actual REST API calls Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com> * Fix feature flag and telemetry export endpoints - Change feature flag endpoint to use NODEJS client type - Fix telemetry endpoints to /telemetry-ext and /telemetry-unauth - Update payload to match proto with system_configuration - Add shared buildUrl utility for protocol handling * Match JDBC telemetry payload format - Change payload structure to match JDBC: uploadTime, items, protoLogs - protoLogs contains JSON-stringified TelemetryFrontendLog objects - Remove workspace_id (JDBC doesn't populate it) - Remove debug logs added during testing * Fix lint errors - Fix import order in FeatureFlagCache - Replace require() with import for driverVersion - Fix variable shadowing - Disable prefer-default-export for urlUtils * Add missing getAuthHeaders method to ClientContextStub Fix TypeScript compilation error by implementing getAuthHeaders method required by IClientContext interface. * Add missing getAuthHeaders method to ClientContextStub Fix TypeScript compilation error by implementing getAuthHeaders method required by IClientContext interface. * Fix prettier formatting * Fix prettier formatting * Add DRIVER_NAME constant for nodejs-sql-driver * Add DRIVER_NAME constant for nodejs-sql-driver * Add missing telemetry fields to match JDBC Added osArch, runtimeVendor, localeName, charSetEncoding, and processName fields to DriverConfiguration to match JDBC implementation. * Add missing telemetry fields to match JDBC Added osArch, runtimeVendor, localeName, charSetEncoding, and processName fields to DriverConfiguration to match JDBC implementation. * Fix TypeScript compilation: add missing fields to system_configuration interface * Fix TypeScript compilation: add missing fields to system_configuration interface * Fix telemetry PR review comments from #325 Three fixes addressing review feedback: 1. Fix documentation typo (sreekanth-db comment) - DatabricksTelemetryExporter.ts:94 - Changed "TelemetryFrontendLog" to "DatabricksTelemetryLog" 2. Add proxy support (jadewang-db comment) - DatabricksTelemetryExporter.ts:exportInternal() - Get HTTP agent from connection provider - Pass agent to fetch for proxy support - Follows same pattern as CloudFetchResultHandler and DBSQLSession - Supports http/https/socks proxies with authentication 3. Fix flush timer to prevent rate limiting (sreekanth-db comment) - MetricsAggregator.ts:flush() - Reset timer after manual flushes (batch size, terminal errors) - Ensures consistent 30s spacing between exports - Prevents rapid successive flushes (e.g., batch at 25s, timer at 30s) All changes follow existing driver patterns and maintain backward compatibility. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com> * Add proxy support to feature flag fetching Feature flag fetching was also missing proxy support like telemetry exporter was. Applied the same fix: - Get HTTP agent from connection provider - Pass agent to fetch call for proxy support - Follows same pattern as CloudFetchResultHandler and DBSQLSession - Supports http/https/socks proxies with authentication This completes proxy support for all HTTP operations in the telemetry system (both telemetry export and feature flag fetching). Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com> * Merge main into telemetry-3-client-management; prefer main versions for infrastructure Takes main's versions of CircuitBreaker, FeatureFlagCache, DatabricksTelemetryExporter, MetricsAggregator, TelemetryEventEmitter, and types — bringing in SSRF hardening, overflow protection, async flush/close, errorStack redaction, and IAuthentication-based auth headers. Removes IClientContext.getAuthHeaders() in favour of the direct authProvider injection pattern from main. Co-authored-by: samikshya-chand_data * fix: prettier formatting and add coverage/ to .prettierignore Co-authored-by: samikshya-chand_data * fix: delete host entry before awaiting close to prevent stale-client race If getOrCreateClient ran after refCount hit 0 but before clients.delete(host), it would receive a closing TelemetryClient. Deleting synchronously first means any concurrent caller gets a fresh instance instead. Co-authored-by: samikshya-chand_data * refactor: simplify TelemetryClient.close() and TelemetryClientProvider - close() made synchronous (no I/O, no reason for async); uses finally to guarantee this.closed=true even when logger throws - TelemetryClientProvider.releaseClient() made synchronous for the same reason - Misleading ConcurrentHashMap reference removed from docstring - getRefCount/getActiveClients marked @internal (test-only surface) - Update tests to match: rejects→throws stubs, remove await on sync calls Co-authored-by: samikshya-chand_data * fix: address PR #326 review feedback from vikrant - F5: Add refCount underflow guard and defensive warn-log in releaseClient so a double-release cannot drive the count negative. Also add a test that double-release does not throw or corrupt state. - F6: Normalize host used as map key (lowercase, strip scheme, default port, trailing slash/dot, whitespace). Aliases like https://Example.COM/ now coalesce to a single entry; add a soft-limit warn at 128 entries to surface alias/leak bugs. Tests added for normalization + alias-release. - F9: Replace `catch (error: any) { error.message }` with a narrowed `error instanceof Error ? error.message : String(error)` so non-Error throws (string, null) cannot escape the swallow-all contract. Test added. - F10: Rewrite header and inline comments to describe the current sync close() semantics instead of an async-await model the code does not implement; note the invariant to re-establish when close() becomes async. - F12: Relax exact debug-log string assertions to regex matchers so cosmetic log rewording does not break unrelated tests. Replace the `expect(true).to.be.true` tautology in TelemetryClient.test.ts with a behavioural `expect(() => client.close()).to.not.throw()`. Signed-off-by: samikshya-chand_data <samikshya.chand@databricks.com> Co-authored-by: Isaac Signed-off-by: samikshya-chand_data <samikshya.chand@databricks.com> * style: apply prettier formatting to TelemetryClientProvider.test.ts Co-authored-by: Isaac Signed-off-by: samikshya-chand_data <samikshya.chand@databricks.com> --------- Signed-off-by: samikshya-chand_data <samikshya.chand@databricks.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
diff --git a/.prettierignore b/.prettierignore
@@ -5,6 +5,7 @@ node_modules
 .nyc_output
 coverage_e2e
 coverage_unit
+coverage
 .clinic
 
 dist
diff --git a/lib/telemetry/TelemetryClient.ts b/lib/telemetry/TelemetryClient.ts
@@ -0,0 +1,65 @@
+/**
+ * Copyright (c) 2025 Databricks Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import IClientContext from '../contracts/IClientContext';
+import { LogLevel } from '../contracts/IDBSQLLogger';
+
+/**
+ * Telemetry client for a specific host.
+ * Managed by TelemetryClientProvider with reference counting.
+ * One client instance is shared across all connections to the same host.
+ */
+class TelemetryClient {
+  private closed: boolean = false;
+
+  constructor(private context: IClientContext, private host: string) {
+    const logger = context.getLogger();
+    logger.log(LogLevel.debug, `Created TelemetryClient for host: ${host}`);
+  }
+
+  /**
+   * Gets the host associated with this client.
+   */
+  getHost(): string {
+    return this.host;
+  }
+
+  /**
+   * Checks if the client has been closed.
+   */
+  isClosed(): boolean {
+    return this.closed;
+  }
+
+  /**
+   * Closes the telemetry client and releases resources.
+   * Should only be called by TelemetryClientProvider when reference count reaches zero.
+   */
+  close(): void {
+    if (this.closed) {
+      return;
+    }
+    try {
+      this.context.getLogger().log(LogLevel.debug, `Closing TelemetryClient for host: ${this.host}`);
+    } catch {
+      // swallow
+    } finally {
+      this.closed = true;
+    }
+  }
+}
+
+export default TelemetryClient;
diff --git a/lib/telemetry/TelemetryClientProvider.ts b/lib/telemetry/TelemetryClientProvider.ts
@@ -0,0 +1,173 @@
+/**
+ * Copyright (c) 2025 Databricks Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import IClientContext from '../contracts/IClientContext';
+import { LogLevel } from '../contracts/IDBSQLLogger';
+import TelemetryClient from './TelemetryClient';
+
+/**
+ * Holds a telemetry client and its reference count.
+ * The reference count tracks how many connections are using this client.
+ */
+interface TelemetryClientHolder {
+  client: TelemetryClient;
+  refCount: number;
+}
+
+// Soft cap on distinct host entries. Above this the provider warns once so a
+// misconfigured caller (per-request hosts, unnormalized aliases) is visible in
+// logs rather than silently growing the map.
+const MAX_CLIENTS_SOFT_LIMIT = 128;
+
+/**
+ * Manages one telemetry client per host.
+ * Prevents rate limiting by sharing clients across connections to the same host.
+ * Instance-based (not singleton), stored in DBSQLClient.
+ *
+ * Reference counts are incremented and decremented synchronously, and
+ * `close()` is sync today, so there is no await between map mutation and
+ * client teardown. The map entry is removed before `close()` runs so a
+ * concurrent `getOrCreateClient` call for the same host gets a fresh
+ * instance rather than receiving this closing one. When `close()` becomes
+ * async (e.g. HTTP flush in [5/7]) the flow will need to `await` after the
+ * delete to preserve the same invariant.
+ */
+class TelemetryClientProvider {
+  private clients: Map<string, TelemetryClientHolder>;
+
+  private softLimitWarned = false;
+
+  constructor(private context: IClientContext) {
+    this.clients = new Map();
+    const logger = context.getLogger();
+    logger.log(LogLevel.debug, 'Created TelemetryClientProvider');
+  }
+
+  /**
+   * Canonicalize host so aliases (scheme, default port, trailing slash, case,
+   * trailing dot, surrounding whitespace) map to the same entry. Kept to a
+   * lightweight lexical normalization — `buildTelemetryUrl` still performs
+   * the strict security validation when a request is actually built.
+   */
+  private static normalizeHostKey(host: string): string {
+    return host
+      .trim()
+      .toLowerCase()
+      .replace(/^https?:\/\//, '')
+      .replace(/\/+$/, '')
+      .replace(/\.$/, '')
+      .replace(/:443$/, '');
+  }
+
+  /**
+   * Gets or creates a telemetry client for the specified host.
+   * Increments the reference count for the client.
+   *
+   * @param host The host identifier (e.g., "workspace.cloud.databricks.com")
+   * @returns The telemetry client for the host
+   */
+  getOrCreateClient(host: string): TelemetryClient {
+    const logger = this.context.getLogger();
+    const key = TelemetryClientProvider.normalizeHostKey(host);
+    let holder = this.clients.get(key);
+
+    if (!holder) {
+      const client = new TelemetryClient(this.context, key);
+      holder = {
+        client,
+        refCount: 0,
+      };
+      this.clients.set(key, holder);
+      logger.log(LogLevel.debug, `Created new TelemetryClient for host: ${host}`);
+
+      if (!this.softLimitWarned && this.clients.size > MAX_CLIENTS_SOFT_LIMIT) {
+        this.softLimitWarned = true;
+        logger.log(
+          LogLevel.warn,
+          `TelemetryClientProvider has ${this.clients.size} distinct hosts — possible alias or leak`,
+        );
+      }
+    }
+
+    holder.refCount += 1;
+    logger.log(LogLevel.debug, `TelemetryClient reference count for ${host}: ${holder.refCount}`);
+
+    return holder.client;
+  }
+
+  /**
+   * Releases a telemetry client for the specified host.
+   * Decrements the reference count and closes the client when it reaches zero.
+   *
+   * @param host The host identifier
+   */
+  releaseClient(host: string): void {
+    const logger = this.context.getLogger();
+    const key = TelemetryClientProvider.normalizeHostKey(host);
+    const holder = this.clients.get(key);
+
+    if (!holder) {
+      logger.log(LogLevel.debug, `No TelemetryClient found for host: ${host}`);
+      return;
+    }
+
+    // Guard against double-release: a caller releasing more times than it got
+    // would otherwise drive refCount negative and close a client another
+    // caller is still holding. Warn loudly and refuse to decrement further.
+    if (holder.refCount <= 0) {
+      logger.log(LogLevel.warn, `Unbalanced release for TelemetryClient host: ${host}`);
+      return;
+    }
+
+    holder.refCount -= 1;
+    logger.log(LogLevel.debug, `TelemetryClient reference count for ${host}: ${holder.refCount}`);
+
+    // Close and remove client when reference count reaches zero.
+    // Remove from map before calling close so a concurrent getOrCreateClient
+    // creates a fresh client rather than receiving this closing one.
+    if (holder.refCount <= 0) {
+      this.clients.delete(key);
+      try {
+        holder.client.close();
+        logger.log(LogLevel.debug, `Closed and removed TelemetryClient for host: ${host}`);
+      } catch (error) {
+        const msg = error instanceof Error ? error.message : String(error);
+        logger.log(LogLevel.debug, `Error releasing TelemetryClient: ${msg}`);
+      }
+    }
+  }
+
+  /**
+   * @internal Exposed for testing only.
+   */
+  getRefCount(host: string): number {
+    const holder = this.clients.get(TelemetryClientProvider.normalizeHostKey(host));
+    return holder ? holder.refCount : 0;
+  }
+
+  /**
+   * @internal Exposed for testing only.
+   */
+  getActiveClients(): Map<string, TelemetryClient> {
+    const result = new Map<string, TelemetryClient>();
+    for (const [host, holder] of this.clients.entries()) {
+      result.set(host, holder.client);
+    }
+    return result;
+  }
+}
+
+export default TelemetryClientProvider;
diff --git a/lib/telemetry/urlUtils.ts b/lib/telemetry/urlUtils.ts
@@ -0,0 +1,31 @@
+/**
+ * Copyright (c) 2025 Databricks Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * Build full URL from host and path, handling protocol correctly.
+ * @param host The hostname (with or without protocol)
+ * @param path The path to append (should start with /)
+ * @returns Full URL with protocol
+ */
+// eslint-disable-next-line import/prefer-default-export
+export function buildUrl(host: string, path: string): string {
+  // Check if host already has protocol
+  if (host.startsWith('http://') || host.startsWith('https://')) {
+    return `${host}${path}`;
+  }
+  // Add https:// if no protocol present
+  return `https://${host}${path}`;
+}
diff --git a/tests/unit/telemetry/TelemetryClient.test.ts b/tests/unit/telemetry/TelemetryClient.test.ts
diff --git a/tests/unit/telemetry/TelemetryClientProvider.test.ts b/tests/unit/telemetry/TelemetryClientProvider.test.ts
