Add token provider infrastructure for token federation (Token Federation 1/3)

lib/DBSQLClient.ts

@@ -19,6 +19,11 @@ 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,
+} from './connection/auth/tokenProvider';
 import IDBSQLLogger, { LogLevel } from './contracts/IDBSQLLogger';
 import DBSQLLogger from './DBSQLLogger';
 import CloseableCollection from './utils/CloseableCollection';
@@ -143,6 +148,12 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
         });
       case 'custom':
         return options.provider;
+      case 'token-provider':
+        return new TokenProviderAuthenticator(options.tokenProvider, this);
+      case 'external-token':
+        return new TokenProviderAuthenticator(new ExternalTokenProvider(options.getToken), this);
+      case 'static-token':
+        return new TokenProviderAuthenticator(StaticTokenProvider.fromJWT(options.staticToken), this);
       // no default
     }
   }

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

lib/connection/auth/tokenProvider/ITokenProvider.ts

@@ -0,0 +1,19 @@
+import Token from './Token';
+
+/**
+ * Interface for token providers that supply access tokens for authentication.
+ * Token providers can be wrapped with caching and federation decorators.
+ */
+export default interface ITokenProvider {
+  /**
+   * Retrieves an access token for authentication.
+   * @returns A Promise that resolves to a Token object containing the access token
+   */
+  getToken(): Promise<Token>;
+
+  /**
+   * Returns the name of this token provider for logging and debugging purposes.
+   * @returns The provider name
+   */
+  getName(): string;
+}

lib/connection/auth/tokenProvider/StaticTokenProvider.ts

@@ -0,0 +1,58 @@
+import ITokenProvider from './ITokenProvider';
+import Token from './Token';
+
+/**
+ * A token provider that returns a static token.
+ * Useful for testing or when the token is obtained through external means.
+ */
+export default class StaticTokenProvider implements ITokenProvider {
+  private readonly token: Token;
+
+  /**
+   * Creates a new StaticTokenProvider.
+   * @param accessToken - The access token string
+   * @param options - Optional token configuration (tokenType, expiresAt, refreshToken, scopes)
+   */
+  constructor(
+    accessToken: string,
+    options?: {
+      tokenType?: string;
+      expiresAt?: Date;
+      refreshToken?: string;
+      scopes?: string[];
+    },
+  ) {
+    this.token = new Token(accessToken, options);
+  }
+
+  /**
+   * Creates a StaticTokenProvider from a JWT string.
+   * The expiration time will be extracted from the JWT payload.
+   * @param jwt - The JWT token string
+   * @param options - Optional token configuration
+   */
+  static fromJWT(
+    jwt: string,
+    options?: {
+      tokenType?: string;
+      refreshToken?: string;
+      scopes?: string[];
+    },
+  ): StaticTokenProvider {
+    const token = Token.fromJWT(jwt, options);
+    return new StaticTokenProvider(token.accessToken, {
+      tokenType: token.tokenType,
+      expiresAt: token.expiresAt,
+      refreshToken: token.refreshToken,
+      scopes: token.scopes,
+    });
+  }
+
+  async getToken(): Promise<Token> {
+    return this.token;
+  }
+
+  getName(): string {
+    return 'StaticTokenProvider';
+  }
+}

lib/connection/auth/tokenProvider/Token.ts

@@ -0,0 +1,151 @@
+import { HeadersInit } from 'node-fetch';
+
+/**
+ * Safety buffer in seconds to consider a token expired before its actual expiration time.
+ * This prevents using tokens that are about to expire during in-flight requests.
+ */
+const EXPIRATION_BUFFER_SECONDS = 30;
+
+/**
+ * Represents an access token with optional metadata and lifecycle management.
+ */
+export default class Token {
+  private readonly _accessToken: string;
+
+  private readonly _tokenType: string;
+
+  private readonly _expiresAt?: Date;
+
+  private readonly _refreshToken?: string;
+
+  private readonly _scopes?: string[];
+
+  constructor(
+    accessToken: string,
+    options?: {
+      tokenType?: string;
+      expiresAt?: Date;
+      refreshToken?: string;
+      scopes?: string[];
+    },
+  ) {
+    this._accessToken = accessToken;
+    this._tokenType = options?.tokenType ?? 'Bearer';
+    this._expiresAt = options?.expiresAt;
+    this._refreshToken = options?.refreshToken;
+    this._scopes = options?.scopes;
+  }
+
+  /**
+   * The access token string.
+   */
+  get accessToken(): string {
+    return this._accessToken;
+  }
+
+  /**
+   * The token type (e.g., "Bearer").
+   */
+  get tokenType(): string {
+    return this._tokenType;
+  }
+
+  /**
+   * The expiration time of the token, if known.
+   */
+  get expiresAt(): Date | undefined {
+    return this._expiresAt;
+  }
+
+  /**
+   * The refresh token, if available.
+   */
+  get refreshToken(): string | undefined {
+    return this._refreshToken;
+  }
+
+  /**
+   * The scopes associated with this token.
+   */
+  get scopes(): string[] | undefined {
+    return this._scopes;
+  }
+
+  /**
+   * Checks if the token has expired, including a safety buffer.
+   * Returns false if expiration time is unknown.
+   */
+  isExpired(): boolean {
+    if (!this._expiresAt) {
+      return false;
+    }
+    const now = new Date();
+    const bufferMs = EXPIRATION_BUFFER_SECONDS * 1000;
+    return this._expiresAt.getTime() - bufferMs <= now.getTime();
+  }
+
+  /**
+   * Sets the Authorization header on the provided headers object.
+   * @param headers - The headers object to modify
+   * @returns The modified headers object with Authorization set
+   */
+  setAuthHeader(headers: HeadersInit): HeadersInit {
+    return {
+      ...headers,
+      Authorization: `${this._tokenType} ${this._accessToken}`,
+    };
+  }
+
+  /**
+   * Creates a Token from a JWT string, extracting the expiration time from the payload.
+   * If the JWT cannot be decoded, the token is created without expiration info.
+   * The server will validate the token anyway, so decoding failures are handled gracefully.
+   * @param jwt - The JWT token string
+   * @param options - Additional token options (tokenType, refreshToken, scopes)
+   * @returns A new Token instance with expiration extracted from the JWT (if available)
+   */
+  static fromJWT(
+    jwt: string,
+    options?: {
+      tokenType?: string;
+      refreshToken?: string;
+      scopes?: string[];
+    },
+  ): Token {
+    let expiresAt: Date | undefined;
+
+    try {
+      const parts = jwt.split('.');
+      if (parts.length >= 2) {
+        const payload = Buffer.from(parts[1], 'base64').toString('utf8');
+        const decoded = JSON.parse(payload);
+        if (typeof decoded.exp === 'number') {
+          expiresAt = new Date(decoded.exp * 1000);
+        }
+      }
+    } catch {
+      // If we can't decode the JWT, we'll proceed without expiration info
+      // The server will validate the token anyway
+    }
+
+    return new Token(jwt, {
+      tokenType: options?.tokenType,
+      expiresAt,
+      refreshToken: options?.refreshToken,
+      scopes: options?.scopes,
+    });
+  }
+
+  /**
+   * Converts the token to a plain object for serialization.
+   */
+  toJSON(): Record<string, unknown> {
+    return {
+      accessToken: this._accessToken,
+      tokenType: this._tokenType,
+      expiresAt: this._expiresAt?.toISOString(),
+      refreshToken: this._refreshToken,
+      scopes: this._scopes,
+    };
+  }
+}

lib/connection/auth/tokenProvider/TokenProviderAuthenticator.ts

@@ -0,0 +1,44 @@
+import { HeadersInit } from 'node-fetch';
+import IAuthentication from '../../contracts/IAuthentication';
+import ITokenProvider from './ITokenProvider';
+import IClientContext from '../../../contracts/IClientContext';
+import { LogLevel } from '../../../contracts/IDBSQLLogger';
+
+/**
+ * Adapts an ITokenProvider to the IAuthentication interface used by the driver.
+ * This allows token providers to be used with the existing authentication system.
+ */
+export default class TokenProviderAuthenticator implements IAuthentication {
+  private readonly tokenProvider: ITokenProvider;
+
+  private readonly context: IClientContext;
+
+  private readonly headers: HeadersInit;
+
+  /**
+   * Creates a new TokenProviderAuthenticator.
+   * @param tokenProvider - The token provider to use for authentication
+   * @param context - The client context for logging
+   * @param headers - Additional headers to include with each request
+   */
+  constructor(tokenProvider: ITokenProvider, context: IClientContext, headers?: HeadersInit) {
+    this.tokenProvider = tokenProvider;
+    this.context = context;
+    this.headers = headers ?? {};
+  }
+
+  async authenticate(): Promise<HeadersInit> {
+    const logger = this.context.getLogger();
+    const providerName = this.tokenProvider.getName();
+
+    logger.log(LogLevel.debug, `TokenProviderAuthenticator: getting token from ${providerName}`);
+
+    const token = await this.tokenProvider.getToken();
+
+    if (token.isExpired()) {
+      logger.log(LogLevel.warn, `TokenProviderAuthenticator: token from ${providerName} is expired`);
+    }
+
+    return token.setAuthHeader(this.headers);
+  }
+}

lib/connection/auth/tokenProvider/index.ts

@@ -0,0 +1,5 @@
+export { default as ITokenProvider } from './ITokenProvider';
+export { default as Token } from './Token';
+export { default as StaticTokenProvider } from './StaticTokenProvider';
+export { default as ExternalTokenProvider, TokenCallback } from './ExternalTokenProvider';
+export { default as TokenProviderAuthenticator } from './TokenProviderAuthenticator';

lib/contracts/IDBSQLClient.ts

@@ -3,6 +3,8 @@ import IDBSQLSession from './IDBSQLSession';
 import IAuthentication from '../connection/contracts/IAuthentication';
 import { ProxyOptions } from '../connection/contracts/IConnectionOptions';
 import OAuthPersistence from '../connection/auth/DatabricksOAuth/OAuthPersistence';
+import ITokenProvider from '../connection/auth/tokenProvider/ITokenProvider';
+import { TokenCallback } from '../connection/auth/tokenProvider/ExternalTokenProvider';
 
 export interface ClientOptions {
   logger?: IDBSQLLogger;
@@ -24,6 +26,18 @@ type AuthOptions =
   | {
       authType: 'custom';
       provider: IAuthentication;
+    }
+  | {
+      authType: 'token-provider';
+      tokenProvider: ITokenProvider;
+    }
+  | {
+      authType: 'external-token';
+      getToken: TokenCallback;
+    }
+  | {
+      authType: 'static-token';
+      staticToken: string;
     };
 
 export type ConnectionOptions = {