[3/7] Telemetry Client Management: TelemetryClient and Provider

.prettierignore

@@ -5,6 +5,7 @@ node_modules
 .nyc_output
 coverage_e2e
 coverage_unit
+coverage
 .clinic
 
 dist

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;

lib/telemetry/TelemetryClientProvider.ts

@@ -0,0 +1,131 @@
+/**
+ * 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;
+}
+
+/**
+ * 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 synchronously so there are no async races
+ * on the count itself. The map entry is deleted before awaiting close() so a
+ * concurrent getOrCreateClient call always gets a fresh instance.
+ */
+class TelemetryClientProvider {
+  private clients: Map<string, TelemetryClientHolder>;
+
+  constructor(private context: IClientContext) {
+    this.clients = new Map();
+    const logger = context.getLogger();
+    logger.log(LogLevel.debug, 'Created TelemetryClientProvider');
+  }
+
+  /**
+   * 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();
+    let holder = this.clients.get(host);
+
+    if (!holder) {
+      // Create new client for this host
+      const client = new TelemetryClient(this.context, host);
+      holder = {
+        client,
+        refCount: 0,
+      };
+      this.clients.set(host, holder);
+      logger.log(LogLevel.debug, `Created new TelemetryClient for host: ${host}`);
+    }
+
+    // Increment reference count
+    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 holder = this.clients.get(host);
+
+    if (!holder) {
+      logger.log(LogLevel.debug, `No TelemetryClient found for host: ${host}`);
+      return;
+    }
+
+    // Decrement reference count
+    holder.refCount -= 1;
+    logger.log(LogLevel.debug, `TelemetryClient reference count for ${host}: ${holder.refCount}`);
+
+    // Close and remove client when reference count reaches zero.
+    // Delete from map before awaiting close so a concurrent getOrCreateClient
+    // creates a fresh client rather than receiving this closing one.
+    if (holder.refCount <= 0) {
+      this.clients.delete(host);
+      try {
+        holder.client.close();
+        logger.log(LogLevel.debug, `Closed and removed TelemetryClient for host: ${host}`);
+      } catch (error: any) {
+        // Swallow all exceptions per requirement
+        logger.log(LogLevel.debug, `Error releasing TelemetryClient: ${error.message}`);
+      }
+    }
+  }
+
+  /**
+   * @internal Exposed for testing only.
+   */
+  getRefCount(host: string): number {
+    const holder = this.clients.get(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;

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}`;
+}

tests/unit/telemetry/TelemetryClient.test.ts

@@ -0,0 +1,155 @@
+/**
+ * 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 { expect } from 'chai';
+import sinon from 'sinon';
+import TelemetryClient from '../../../lib/telemetry/TelemetryClient';
+import ClientContextStub from '../.stubs/ClientContextStub';
+import { LogLevel } from '../../../lib/contracts/IDBSQLLogger';
+
+describe('TelemetryClient', () => {
+  const HOST = 'workspace.cloud.databricks.com';
+
+  describe('Constructor', () => {
+    it('should create client with host', () => {
+      const context = new ClientContextStub();
+      const client = new TelemetryClient(context, HOST);
+
+      expect(client.getHost()).to.equal(HOST);
+      expect(client.isClosed()).to.be.false;
+    });
+
+    it('should log creation at debug level', () => {
+      const context = new ClientContextStub();
+      const logSpy = sinon.spy(context.logger, 'log');
+
+      new TelemetryClient(context, HOST);
+
+      expect(logSpy.calledWith(LogLevel.debug, `Created TelemetryClient for host: ${HOST}`)).to.be.true;
+    });
+  });
+
+  describe('getHost', () => {
+    it('should return the host identifier', () => {
+      const context = new ClientContextStub();
+      const client = new TelemetryClient(context, HOST);
+
+      expect(client.getHost()).to.equal(HOST);
+    });
+  });
+
+  describe('isClosed', () => {
+    it('should return false initially', () => {
+      const context = new ClientContextStub();
+      const client = new TelemetryClient(context, HOST);
+
+      expect(client.isClosed()).to.be.false;
+    });
+
+    it('should return true after close', async () => {
+      const context = new ClientContextStub();
+      const client = new TelemetryClient(context, HOST);
+
+      client.close();
+
+      expect(client.isClosed()).to.be.true;
+    });
+  });
+
+  describe('close', () => {
+    it('should set closed flag', async () => {
+      const context = new ClientContextStub();
+      const client = new TelemetryClient(context, HOST);
+
+      client.close();
+
+      expect(client.isClosed()).to.be.true;
+    });
+
+    it('should log closure at debug level', async () => {
+      const context = new ClientContextStub();
+      const logSpy = sinon.spy(context.logger, 'log');
+      const client = new TelemetryClient(context, HOST);
+
+      client.close();
+
+      expect(logSpy.calledWith(LogLevel.debug, `Closing TelemetryClient for host: ${HOST}`)).to.be.true;
+    });
+
+    it('should be idempotent', async () => {
+      const context = new ClientContextStub();
+      const logSpy = sinon.spy(context.logger, 'log');
+      const client = new TelemetryClient(context, HOST);
+
+      client.close();
+      const firstCallCount = logSpy.callCount;
+
+      client.close();
+
+      // Should not log again on second close
+      expect(logSpy.callCount).to.equal(firstCallCount);
+      expect(client.isClosed()).to.be.true;
+    });
+
+    it('should swallow all exceptions', async () => {
+      const context = new ClientContextStub();
+      const client = new TelemetryClient(context, HOST);
+
+      // Force an error by stubbing the logger
+      const error = new Error('Logger error');
+      sinon.stub(context.logger, 'log').throws(error);
+
+      // Should not throw
+      client.close();
+      // If we get here without throwing, the test passes
+      expect(true).to.be.true;
+    });
+
+    it('should still set closed when logger throws', () => {
+      const context = new ClientContextStub();
+      const client = new TelemetryClient(context, HOST);
+
+      sinon.stub(context.logger, 'log').throws(new Error('Logger error'));
+
+      client.close();
+
+      expect(client.isClosed()).to.be.true;
+    });
+  });
+
+  describe('Context usage', () => {
+    it('should use logger from context', () => {
+      const context = new ClientContextStub();
+      const logSpy = sinon.spy(context.logger, 'log');
+
+      new TelemetryClient(context, HOST);
+
+      expect(logSpy.called).to.be.true;
+    });
+
+    it('should log all messages at debug level only', async () => {
+      const context = new ClientContextStub();
+      const logSpy = sinon.spy(context.logger, 'log');
+      const client = new TelemetryClient(context, HOST);
+
+      client.close();
+
+      logSpy.getCalls().forEach((call) => {
+        expect(call.args[0]).to.equal(LogLevel.debug);
+      });
+    });
+  });
+});