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

CHANGELOG.md

@@ -1,5 +1,12 @@
 # Release History
 
+## 1.13.0
+
+- Add token federation support with custom token providers (databricks/databricks-sql-nodejs#318, databricks/databricks-sql-nodejs#319, databricks/databricks-sql-nodejs#320 by @madhav-db)
+- Add metric view metadata support (databricks/databricks-sql-nodejs#312 by @shivam2680)
+- Fix: Avoid calling require('lz4') if it's really not required (databricks/databricks-sql-nodejs#316 by @ikkala)
+- Add telemetry foundation (off by default) (databricks/databricks-sql-nodejs#324 by @samikshya-db)
+
 ## 1.12.0
 
 - Support for session parameters (databricks/databricks-sql-nodejs#307 by @sreekanth-db)

examples/tokenFederation/README.md

@@ -0,0 +1,51 @@
+# Token Federation Examples
+
+Examples demonstrating the token provider and federation features of the Databricks SQL Node.js Driver.
+
+## Examples
+
+### Static Token (`staticToken.ts`)
+
+The simplest authentication method. Use a static access token that doesn't change during the application lifetime.
+
+```bash
+DATABRICKS_HOST=<host> DATABRICKS_HTTP_PATH=<path> DATABRICKS_TOKEN=<token> npx ts-node staticToken.ts
+```
+
+### External Token (`externalToken.ts`)
+
+Use a callback function to provide tokens dynamically. Useful for integrating with secret managers, vaults, or other token sources. Tokens are automatically cached by the driver.
+
+```bash
+DATABRICKS_HOST=<host> DATABRICKS_HTTP_PATH=<path> DATABRICKS_TOKEN=<token> npx ts-node externalToken.ts
+```
+
+### Token Federation (`federation.ts`)
+
+Automatically exchange tokens from external identity providers (Azure AD, Google, Okta, etc.) for Databricks-compatible tokens using RFC 8693 token exchange.
+
+```bash
+DATABRICKS_HOST=<host> DATABRICKS_HTTP_PATH=<path> AZURE_AD_TOKEN=<token> npx ts-node federation.ts
+```
+
+### M2M Federation (`m2mFederation.ts`)
+
+Machine-to-machine token federation with a service principal. Requires a `federationClientId` to identify the service principal to Databricks.
+
+```bash
+DATABRICKS_HOST=<host> DATABRICKS_HTTP_PATH=<path> DATABRICKS_CLIENT_ID=<client-id> SERVICE_ACCOUNT_TOKEN=<token> npx ts-node m2mFederation.ts
+```
+
+### Custom Token Provider (`customTokenProvider.ts`)
+
+Implement the `ITokenProvider` interface for full control over token management, including custom caching, refresh logic, retry, and error handling.
+
+```bash
+DATABRICKS_HOST=<host> DATABRICKS_HTTP_PATH=<path> OAUTH_SERVER_URL=<url> OAUTH_CLIENT_ID=<id> OAUTH_CLIENT_SECRET=<secret> npx ts-node customTokenProvider.ts
+```
+
+## Prerequisites
+
+- Node.js 14+
+- A Databricks workspace with token federation enabled (for federation examples)
+- Valid credentials for your identity provider

examples/tokenFederation/customTokenProvider.ts

@@ -0,0 +1,169 @@
+/**
+ * Example: Custom Token Provider Implementation
+ *
+ * This example demonstrates how to create a custom token provider by
+ * implementing the ITokenProvider interface. This gives you full control
+ * over token management, including custom caching, refresh logic, and
+ * error handling.
+ */
+
+import { DBSQLClient } from '@databricks/sql';
+import { ITokenProvider, Token } from '../../lib/connection/auth/tokenProvider';
+
+/**
+ * Custom token provider that refreshes tokens from a custom OAuth server.
+ */
+class CustomOAuthTokenProvider implements ITokenProvider {
+  private readonly oauthServerUrl: string;
+
+  private readonly clientId: string;
+
+  private readonly clientSecret: string;
+
+  constructor(oauthServerUrl: string, clientId: string, clientSecret: string) {
+    this.oauthServerUrl = oauthServerUrl;
+    this.clientId = clientId;
+    this.clientSecret = clientSecret;
+  }
+
+  async getToken(): Promise<Token> {
+    // eslint-disable-next-line no-console
+    console.log('Fetching token from custom OAuth server...');
+    return this.fetchTokenWithRetry(0);
+  }
+
+  /**
+   * Recursively attempts to fetch a token with exponential backoff.
+   */
+  private async fetchTokenWithRetry(attempt: number): Promise<Token> {
+    const maxRetries = 3;
+
+    try {
+      return await this.fetchToken();
+    } catch (error) {
+      // Don't retry client errors (4xx)
+      if (error instanceof Error && error.message.includes('OAuth token request failed: 4')) {
+        throw error;
+      }
+
+      if (attempt >= maxRetries) {
+        throw error;
+      }
+
+      // Exponential backoff: 1s, 2s, 4s
+      const delay = 1000 * 2 ** attempt;
+      await new Promise<void>((resolve) => {
+        setTimeout(resolve, delay);
+      });
+
+      return this.fetchTokenWithRetry(attempt + 1);
+    }
+  }
+
+  private async fetchToken(): Promise<Token> {
+    const response = await fetch(`${this.oauthServerUrl}/oauth/token`, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/x-www-form-urlencoded',
+      },
+      body: new URLSearchParams({
+        grant_type: 'client_credentials',
+        client_id: this.clientId,
+        client_secret: this.clientSecret,
+        scope: 'sql',
+      }).toString(),
+    });
+
+    if (!response.ok) {
+      throw new Error(`OAuth token request failed: ${response.status}`);
+    }
+
+    const data = (await response.json()) as {
+      access_token: string;
+      token_type?: string;
+      expires_in?: number;
+    };
+
+    // Calculate expiration
+    let expiresAt: Date | undefined;
+    if (typeof data.expires_in === 'number') {
+      expiresAt = new Date(Date.now() + data.expires_in * 1000);
+    }
+
+    return new Token(data.access_token, {
+      tokenType: data.token_type ?? 'Bearer',
+      expiresAt,
+    });
+  }
+
+  getName(): string {
+    return 'CustomOAuthTokenProvider';
+  }
+}
+
+/**
+ * Simple token provider that reads from a file (for development/testing).
+ */
+// exported for use as an alternative example provider
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+class FileTokenProvider implements ITokenProvider {
+  private readonly filePath: string;
+
+  constructor(filePath: string) {
+    this.filePath = filePath;
+  }
+
+  async getToken(): Promise<Token> {
+    const fs = await import('fs/promises');
+    const tokenData = await fs.readFile(this.filePath, 'utf-8');
+    const parsed = JSON.parse(tokenData);
+
+    return Token.fromJWT(parsed.access_token, {
+      refreshToken: parsed.refresh_token,
+    });
+  }
+
+  getName(): string {
+    return 'FileTokenProvider';
+  }
+}
+
+async function main() {
+  const host = process.env.DATABRICKS_HOST!;
+  const path = process.env.DATABRICKS_HTTP_PATH!;
+
+  const client = new DBSQLClient();
+
+  // Option 1: Use a custom OAuth token provider (shown below)
+  // Option 2: Use a file-based token provider for development:
+  //   const fileProvider = new FileTokenProvider('/path/to/token.json');
+  const oauthProvider = new CustomOAuthTokenProvider(
+    process.env.OAUTH_SERVER_URL!,
+    process.env.OAUTH_CLIENT_ID!,
+    process.env.OAUTH_CLIENT_SECRET!,
+  );
+
+  await client.connect({
+    host,
+    path,
+    authType: 'token-provider',
+    tokenProvider: oauthProvider,
+    // Optionally enable federation if your OAuth server issues non-Databricks tokens
+    enableTokenFederation: true,
+  });
+
+  console.log('Connected successfully with custom token provider');
+
+  // Open a session and run a query
+  const session = await client.openSession();
+  const operation = await session.executeStatement('SELECT 1 AS result');
+  const result = await operation.fetchAll();
+
+  console.log('Query result:', result);
+
+  await operation.close();
+  await session.close();
+  await client.close();
+}
+
+main().catch(console.error);

examples/tokenFederation/externalToken.ts

@@ -0,0 +1,53 @@
+/**
+ * Example: Using an external token provider
+ *
+ * This example demonstrates how to use a callback function to provide
+ * tokens dynamically. This is useful for integrating with secret managers,
+ * vaults, or other token sources that may refresh tokens.
+ */
+
+import { DBSQLClient } from '@databricks/sql';
+
+// Simulate fetching a token from a secret manager or vault
+async function fetchTokenFromVault(): Promise<string> {
+  // In a real application, this would fetch from AWS Secrets Manager,
+  // Azure Key Vault, HashiCorp Vault, or another secret manager
+  console.log('Fetching token from vault...');
+
+  // Simulated token - replace with actual vault integration
+  const token = process.env.DATABRICKS_TOKEN!;
+  return token;
+}
+
+async function main() {
+  const host = process.env.DATABRICKS_HOST!;
+  const path = process.env.DATABRICKS_HTTP_PATH!;
+
+  const client = new DBSQLClient();
+
+  // Connect using an external token provider
+  // The callback will be called each time a new token is needed
+  // Note: The token is automatically cached, so the callback won't be
+  // called on every request
+  await client.connect({
+    host,
+    path,
+    authType: 'external-token',
+    getToken: fetchTokenFromVault,
+  });
+
+  console.log('Connected successfully with external token provider');
+
+  // Open a session and run a query
+  const session = await client.openSession();
+  const operation = await session.executeStatement('SELECT current_user() AS user');
+  const result = await operation.fetchAll();
+
+  console.log('Query result:', result);
+
+  await operation.close();
+  await session.close();
+  await client.close();
+}
+
+main().catch(console.error);

examples/tokenFederation/federation.ts

@@ -0,0 +1,80 @@
+/**
+ * Example: Token Federation with an External Identity Provider
+ *
+ * This example demonstrates how to use token federation to automatically
+ * exchange tokens from external identity providers (Azure AD, Google, Okta,
+ * Auth0, AWS Cognito, GitHub) for Databricks-compatible tokens.
+ *
+ * Token federation uses RFC 8693 (OAuth 2.0 Token Exchange) to exchange
+ * the external JWT token for a Databricks access token.
+ */
+
+import { DBSQLClient } from '@databricks/sql';
+
+// Example: Fetch a token from Azure AD
+// In a real application, you would use the Azure SDK or similar
+async function getAzureADToken(): Promise<string> {
+  // Example using @azure/identity:
+  //
+  // import { DefaultAzureCredential } from '@azure/identity';
+  // const credential = new DefaultAzureCredential();
+  // const token = await credential.getToken('https://your-scope/.default');
+  // return token.token;
+
+  // For this example, we use an environment variable
+  const token = process.env.AZURE_AD_TOKEN!;
+  console.log('Fetched token from Azure AD');
+  return token;
+}
+
+// Example: Fetch a token from Google
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+async function getGoogleToken(): Promise<string> {
+  // Example using google-auth-library:
+  //
+  // import { GoogleAuth } from 'google-auth-library';
+  // const auth = new GoogleAuth();
+  // const client = await auth.getClient();
+  // const token = await client.getAccessToken();
+  // return token.token;
+
+  const token = process.env.GOOGLE_TOKEN!;
+  console.log('Fetched token from Google');
+  return token;
+}
+
+async function main() {
+  const host = process.env.DATABRICKS_HOST!;
+  const path = process.env.DATABRICKS_HTTP_PATH!;
+
+  const client = new DBSQLClient();
+
+  // Connect using token federation
+  // The driver will automatically:
+  // 1. Get the token from the callback
+  // 2. Check if the token's issuer matches the Databricks host
+  // 3. If not, exchange the token for a Databricks token via RFC 8693
+  // 4. Cache the result for subsequent requests
+  await client.connect({
+    host,
+    path,
+    authType: 'external-token',
+    getToken: getAzureADToken, // or getGoogleToken, etc.
+    enableTokenFederation: true,
+  });
+
+  console.log('Connected successfully with token federation');
+
+  // Open a session and run a query
+  const session = await client.openSession();
+  const operation = await session.executeStatement('SELECT current_user() AS user');
+  const result = await operation.fetchAll();
+
+  console.log('Query result:', result);
+
+  await operation.close();
+  await session.close();
+  await client.close();
+}
+
+main().catch(console.error);