(Token Federation 2/3)

lib/DBSQLClient.ts

@@ -19,6 +19,14 @@ import HiveDriverError from './errors/HiveDriverError';
 import { buildUserAgentString, definedOrError } from './utils';
 import PlainHttpAuthentication from './connection/auth/PlainHttpAuthentication';
 import DatabricksOAuth, { OAuthFlow } from './connection/auth/DatabricksOAuth';
+import {
+  TokenProviderAuthenticator,
+  StaticTokenProvider,
+  ExternalTokenProvider,
+  CachedTokenProvider,
+  FederationProvider,
+  ITokenProvider,
+} from './connection/auth/tokenProvider';
 import IDBSQLLogger, { LogLevel } from './contracts/IDBSQLLogger';
 import DBSQLLogger from './DBSQLLogger';
 import CloseableCollection from './utils/CloseableCollection';
@@ -143,10 +151,63 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
         });
       case 'custom':
         return options.provider;
+      case 'token-provider':
+        return new TokenProviderAuthenticator(
+          this.wrapTokenProvider(
+            options.tokenProvider,
+            options.host,
+            options.enableTokenFederation,
+            options.federationClientId,
+          ),
+          this,
+        );
+      case 'external-token':
+        return new TokenProviderAuthenticator(
+          this.wrapTokenProvider(
+            new ExternalTokenProvider(options.getToken),
+            options.host,
+            options.enableTokenFederation,
+            options.federationClientId,
+          ),
+          this,
+        );
+      case 'static-token':
+        return new TokenProviderAuthenticator(
+          this.wrapTokenProvider(
+            StaticTokenProvider.fromJWT(options.staticToken),
+            options.host,
+            options.enableTokenFederation,
+            options.federationClientId,
+          ),
+          this,
+        );
       // no default
     }
   }
 
+  /**
+   * Wraps a token provider with caching and optional federation.
+   * Caching is always enabled by default. Federation is opt-in.
+   */
+  private wrapTokenProvider(
+    provider: ITokenProvider,
+    host: string,
+    enableFederation?: boolean,
+    federationClientId?: string,
+  ): ITokenProvider {
+    // Always wrap with caching first
+    let wrapped: ITokenProvider = new CachedTokenProvider(provider);
+
+    // Optionally wrap with federation
+    if (enableFederation) {
+      wrapped = new FederationProvider(wrapped, host, {
+        clientId: federationClientId,
+      });
+    }
+
+    return wrapped;
+  }
+
   private createConnectionProvider(options: ConnectionOptions): IConnectionProvider {
     return new HttpConnection(this.getConnectionOptions(options), this);
   }

lib/connection/auth/tokenProvider/CachedTokenProvider.ts

@@ -0,0 +1,98 @@
+import ITokenProvider from './ITokenProvider';
+import Token from './Token';
+
+/**
+ * Default refresh threshold in milliseconds (5 minutes).
+ * Tokens will be refreshed when they are within this threshold of expiring.
+ */
+const DEFAULT_REFRESH_THRESHOLD_MS = 5 * 60 * 1000;
+
+/**
+ * A token provider that wraps another provider with automatic caching.
+ * Tokens are cached and reused until they are close to expiring.
+ */
+export default class CachedTokenProvider implements ITokenProvider {
+  private readonly baseProvider: ITokenProvider;
+
+  private readonly refreshThresholdMs: number;
+
+  private cache: Token | null = null;
+
+  private refreshPromise: Promise<Token> | null = null;
+
+  /**
+   * Creates a new CachedTokenProvider.
+   * @param baseProvider - The underlying token provider to cache
+   * @param options - Optional configuration
+   * @param options.refreshThresholdMs - Refresh tokens this many ms before expiry (default: 5 minutes)
+   */
+  constructor(
+    baseProvider: ITokenProvider,
+    options?: {
+      refreshThresholdMs?: number;
+    },
+  ) {
+    this.baseProvider = baseProvider;
+    this.refreshThresholdMs = options?.refreshThresholdMs ?? DEFAULT_REFRESH_THRESHOLD_MS;
+  }
+
+  async getToken(): Promise<Token> {
+    // Return cached token if it's still valid
+    if (this.cache && !this.shouldRefresh(this.cache)) {
+      return this.cache;
+    }
+
+    // If already refreshing, wait for that to complete
+    if (this.refreshPromise) {
+      return this.refreshPromise;
+    }
+
+    // Start refresh
+    this.refreshPromise = this.refreshToken();
+
+    try {
+      const token = await this.refreshPromise;
+      return token;
+    } finally {
+      this.refreshPromise = null;
+    }
+  }
+
+  getName(): string {
+    return `cached[${this.baseProvider.getName()}]`;
+  }
+
+  /**
+   * Clears the cached token, forcing a refresh on the next getToken() call.
+   */
+  clearCache(): void {
+    this.cache = null;
+  }
+
+  /**
+   * Determines if the token should be refreshed.
+   * @param token - The token to check
+   * @returns true if the token should be refreshed
+   */
+  private shouldRefresh(token: Token): boolean {
+    // If no expiration is known, don't refresh proactively
+    if (!token.expiresAt) {
+      return false;
+    }
+
+    const now = Date.now();
+    const expiresAtMs = token.expiresAt.getTime();
+    const refreshAtMs = expiresAtMs - this.refreshThresholdMs;
+
+    return now >= refreshAtMs;
+  }
+
+  /**
+   * Fetches a new token from the base provider and caches it.
+   */
+  private async refreshToken(): Promise<Token> {
+    const token = await this.baseProvider.getToken();
+    this.cache = token;
+    return token;
+  }
+}

lib/connection/auth/tokenProvider/ExternalTokenProvider.ts

@@ -0,0 +1,52 @@
+import ITokenProvider from './ITokenProvider';
+import Token from './Token';
+
+/**
+ * Type for the callback function that retrieves tokens from external sources.
+ */
+export type TokenCallback = () => Promise<string>;
+
+/**
+ * A token provider that delegates token retrieval to an external callback function.
+ * Useful for integrating with secret managers, vaults, or other token sources.
+ */
+export default class ExternalTokenProvider implements ITokenProvider {
+  private readonly getTokenCallback: TokenCallback;
+
+  private readonly parseJWT: boolean;
+
+  private readonly providerName: string;
+
+  /**
+   * Creates a new ExternalTokenProvider.
+   * @param getToken - Callback function that returns the access token string
+   * @param options - Optional configuration
+   * @param options.parseJWT - If true, attempt to extract expiration from JWT payload (default: true)
+   * @param options.name - Custom name for this provider (default: "ExternalTokenProvider")
+   */
+  constructor(
+    getToken: TokenCallback,
+    options?: {
+      parseJWT?: boolean;
+      name?: string;
+    },
+  ) {
+    this.getTokenCallback = getToken;
+    this.parseJWT = options?.parseJWT ?? true;
+    this.providerName = options?.name ?? 'ExternalTokenProvider';
+  }
+
+  async getToken(): Promise<Token> {
+    const accessToken = await this.getTokenCallback();
+
+    if (this.parseJWT) {
+      return Token.fromJWT(accessToken);
+    }
+
+    return new Token(accessToken);
+  }
+
+  getName(): string {
+    return this.providerName;
+  }
+}