Skip to content

Commit 870bcb3

Browse files
madhav-dbclaude
andauthored
Add token provider infrastructure for token federation (Token Federation 1/3) (#318)
* Add token provider infrastructure for token federation This PR introduces the foundational token provider system that enables custom token sources for authentication. This is the first of three PRs implementing token federation support. New components: - ITokenProvider: Core interface for token providers - Token: Token class with JWT parsing and expiration handling - StaticTokenProvider: Provides a constant token - ExternalTokenProvider: Delegates to a callback function - TokenProviderAuthenticator: Adapts token providers to IAuthentication New auth types in ConnectionOptions: - 'token-provider': Use a custom ITokenProvider - 'external-token': Use a callback function - 'static-token': Use a static token string * Fix TokenProviderAuthenticator test - remove log assertions LoggerStub doesn't have a logs property, so removed tests that checked for debug and warning log messages. The important behavior (token provider authentication) is still tested. * Fix prettier formatting in TokenProviderAuthenticator * Fix Copilot issues: update fromJWT docs and remove TokenCallback duplication - Updated Token.fromJWT() documentation to reflect that it handles decoding failures gracefully instead of throwing errors - Removed duplicate TokenCallback type definition from IDBSQLClient.ts - Now imports TokenCallback from ExternalTokenProvider.ts to maintain a single source of truth * address comments * Retry token fetch when expired before throwing error TokenProviderAuthenticator now requests a fresh token from the provider when the initial token is expired, only throwing if the retry also returns an expired token. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 666584e commit 870bcb3

File tree

12 files changed

+861
-0
lines changed

12 files changed

+861
-0
lines changed

lib/DBSQLClient.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@ import HiveDriverError from './errors/HiveDriverError';
1919
import { buildUserAgentString, definedOrError } from './utils';
2020
import PlainHttpAuthentication from './connection/auth/PlainHttpAuthentication';
2121
import DatabricksOAuth, { OAuthFlow } from './connection/auth/DatabricksOAuth';
22+
import {
23+
TokenProviderAuthenticator,
24+
StaticTokenProvider,
25+
ExternalTokenProvider,
26+
} from './connection/auth/tokenProvider';
2227
import IDBSQLLogger, { LogLevel } from './contracts/IDBSQLLogger';
2328
import DBSQLLogger from './DBSQLLogger';
2429
import CloseableCollection from './utils/CloseableCollection';
@@ -143,6 +148,12 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
143148
});
144149
case 'custom':
145150
return options.provider;
151+
case 'token-provider':
152+
return new TokenProviderAuthenticator(options.tokenProvider, this);
153+
case 'external-token':
154+
return new TokenProviderAuthenticator(new ExternalTokenProvider(options.getToken), this);
155+
case 'static-token':
156+
return new TokenProviderAuthenticator(StaticTokenProvider.fromJWT(options.staticToken), this);
146157
// no default
147158
}
148159
}
Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
import ITokenProvider from './ITokenProvider';
2+
import Token from './Token';
3+
4+
/**
5+
* Type for the callback function that retrieves tokens from external sources.
6+
*/
7+
export type TokenCallback = () => Promise<string>;
8+
9+
/**
10+
* A token provider that delegates token retrieval to an external callback function.
11+
* Useful for integrating with secret managers, vaults, or other token sources.
12+
*/
13+
export default class ExternalTokenProvider implements ITokenProvider {
14+
private readonly getTokenCallback: TokenCallback;
15+
16+
private readonly parseJWT: boolean;
17+
18+
private readonly providerName: string;
19+
20+
/**
21+
* Creates a new ExternalTokenProvider.
22+
* @param getToken - Callback function that returns the access token string
23+
* @param options - Optional configuration
24+
* @param options.parseJWT - If true, attempt to extract expiration from JWT payload (default: true)
25+
* @param options.name - Custom name for this provider (default: "ExternalTokenProvider")
26+
*/
27+
constructor(
28+
getToken: TokenCallback,
29+
options?: {
30+
parseJWT?: boolean;
31+
name?: string;
32+
},
33+
) {
34+
this.getTokenCallback = getToken;
35+
this.parseJWT = options?.parseJWT ?? true;
36+
this.providerName = options?.name ?? 'ExternalTokenProvider';
37+
}
38+
39+
async getToken(): Promise<Token> {
40+
const accessToken = await this.getTokenCallback();
41+
42+
if (this.parseJWT) {
43+
return Token.fromJWT(accessToken);
44+
}
45+
46+
return new Token(accessToken);
47+
}
48+
49+
getName(): string {
50+
return this.providerName;
51+
}
52+
}
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import Token from './Token';
2+
3+
/**
4+
* Interface for token providers that supply access tokens for authentication.
5+
* Token providers can be wrapped with caching and federation decorators.
6+
*/
7+
export default interface ITokenProvider {
8+
/**
9+
* Retrieves an access token for authentication.
10+
* @returns A Promise that resolves to a Token object containing the access token
11+
*/
12+
getToken(): Promise<Token>;
13+
14+
/**
15+
* Returns the name of this token provider for logging and debugging purposes.
16+
* @returns The provider name
17+
*/
18+
getName(): string;
19+
}
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
import ITokenProvider from './ITokenProvider';
2+
import Token, { TokenOptions, TokenFromJWTOptions } from './Token';
3+
4+
/**
5+
* A token provider that returns a static token.
6+
* Useful for testing or when the token is obtained through external means.
7+
*/
8+
export default class StaticTokenProvider implements ITokenProvider {
9+
private readonly token: Token;
10+
11+
/**
12+
* Creates a new StaticTokenProvider.
13+
* @param accessToken - The access token string
14+
* @param options - Optional token configuration (tokenType, expiresAt, refreshToken, scopes)
15+
*/
16+
constructor(accessToken: string, options?: TokenOptions) {
17+
this.token = new Token(accessToken, options);
18+
}
19+
20+
/**
21+
* Creates a StaticTokenProvider from a JWT string.
22+
* The expiration time will be extracted from the JWT payload.
23+
* @param jwt - The JWT token string
24+
* @param options - Optional token configuration
25+
*/
26+
static fromJWT(jwt: string, options?: TokenFromJWTOptions): StaticTokenProvider {
27+
const token = Token.fromJWT(jwt, options);
28+
return new StaticTokenProvider(token.accessToken, {
29+
tokenType: token.tokenType,
30+
expiresAt: token.expiresAt,
31+
refreshToken: token.refreshToken,
32+
scopes: token.scopes,
33+
});
34+
}
35+
36+
async getToken(): Promise<Token> {
37+
return this.token;
38+
}
39+
40+
getName(): string {
41+
return 'StaticTokenProvider';
42+
}
43+
}
Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
import { HeadersInit } from 'node-fetch';
2+
3+
/**
4+
* Safety buffer in seconds to consider a token expired before its actual expiration time.
5+
* This prevents using tokens that are about to expire during in-flight requests.
6+
*/
7+
const EXPIRATION_BUFFER_SECONDS = 30;
8+
9+
/**
10+
* Options for creating a Token instance.
11+
*/
12+
export interface TokenOptions {
13+
/** The token type (e.g., "Bearer"). Defaults to "Bearer". */
14+
tokenType?: string;
15+
/** The expiration time of the token. */
16+
expiresAt?: Date;
17+
/** The refresh token, if available. */
18+
refreshToken?: string;
19+
/** The scopes associated with this token. */
20+
scopes?: string[];
21+
}
22+
23+
/**
24+
* Options for creating a Token from a JWT string.
25+
* Does not include expiresAt since it is extracted from the JWT payload.
26+
*/
27+
export type TokenFromJWTOptions = Omit<TokenOptions, 'expiresAt'>;
28+
29+
/**
30+
* Represents an access token with optional metadata and lifecycle management.
31+
*/
32+
export default class Token {
33+
private readonly _accessToken: string;
34+
35+
private readonly _tokenType: string;
36+
37+
private readonly _expiresAt?: Date;
38+
39+
private readonly _refreshToken?: string;
40+
41+
private readonly _scopes?: string[];
42+
43+
constructor(accessToken: string, options?: TokenOptions) {
44+
this._accessToken = accessToken;
45+
this._tokenType = options?.tokenType ?? 'Bearer';
46+
this._expiresAt = options?.expiresAt;
47+
this._refreshToken = options?.refreshToken;
48+
this._scopes = options?.scopes;
49+
}
50+
51+
/**
52+
* The access token string.
53+
*/
54+
get accessToken(): string {
55+
return this._accessToken;
56+
}
57+
58+
/**
59+
* The token type (e.g., "Bearer").
60+
*/
61+
get tokenType(): string {
62+
return this._tokenType;
63+
}
64+
65+
/**
66+
* The expiration time of the token, if known.
67+
*/
68+
get expiresAt(): Date | undefined {
69+
return this._expiresAt;
70+
}
71+
72+
/**
73+
* The refresh token, if available.
74+
*/
75+
get refreshToken(): string | undefined {
76+
return this._refreshToken;
77+
}
78+
79+
/**
80+
* The scopes associated with this token.
81+
*/
82+
get scopes(): string[] | undefined {
83+
return this._scopes;
84+
}
85+
86+
/**
87+
* Checks if the token has expired, including a safety buffer.
88+
* Returns false if expiration time is unknown.
89+
*/
90+
isExpired(): boolean {
91+
if (!this._expiresAt) {
92+
return false;
93+
}
94+
const now = new Date();
95+
const bufferMs = EXPIRATION_BUFFER_SECONDS * 1000;
96+
return this._expiresAt.getTime() - bufferMs <= now.getTime();
97+
}
98+
99+
/**
100+
* Sets the Authorization header on the provided headers object.
101+
* @param headers - The headers object to modify
102+
* @returns The modified headers object with Authorization set
103+
*/
104+
setAuthHeader(headers: HeadersInit): HeadersInit {
105+
return {
106+
...headers,
107+
Authorization: `${this._tokenType} ${this._accessToken}`,
108+
};
109+
}
110+
111+
/**
112+
* Creates a Token from a JWT string, extracting the expiration time from the payload.
113+
* If the JWT cannot be decoded, the token is created without expiration info.
114+
* The server will validate the token anyway, so decoding failures are handled gracefully.
115+
* @param jwt - The JWT token string
116+
* @param options - Additional token options (tokenType, refreshToken, scopes).
117+
* Note: expiresAt is not accepted here as it is extracted from the JWT payload.
118+
* @returns A new Token instance with expiration extracted from the JWT (if available)
119+
*/
120+
static fromJWT(jwt: string, options?: TokenFromJWTOptions): Token {
121+
let expiresAt: Date | undefined;
122+
123+
try {
124+
const parts = jwt.split('.');
125+
if (parts.length >= 2) {
126+
const payload = Buffer.from(parts[1], 'base64').toString('utf8');
127+
const decoded = JSON.parse(payload);
128+
if (typeof decoded.exp === 'number') {
129+
expiresAt = new Date(decoded.exp * 1000);
130+
}
131+
}
132+
} catch {
133+
// If we can't decode the JWT, we'll proceed without expiration info
134+
// The server will validate the token anyway
135+
}
136+
137+
return new Token(jwt, {
138+
tokenType: options?.tokenType,
139+
expiresAt,
140+
refreshToken: options?.refreshToken,
141+
scopes: options?.scopes,
142+
});
143+
}
144+
145+
/**
146+
* Converts the token to a plain object for serialization.
147+
*/
148+
toJSON(): Record<string, unknown> {
149+
return {
150+
accessToken: this._accessToken,
151+
tokenType: this._tokenType,
152+
expiresAt: this._expiresAt?.toISOString(),
153+
refreshToken: this._refreshToken,
154+
scopes: this._scopes,
155+
};
156+
}
157+
}
Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
import { HeadersInit } from 'node-fetch';
2+
import IAuthentication from '../../contracts/IAuthentication';
3+
import ITokenProvider from './ITokenProvider';
4+
import IClientContext from '../../../contracts/IClientContext';
5+
import { LogLevel } from '../../../contracts/IDBSQLLogger';
6+
7+
/**
8+
* Adapts an ITokenProvider to the IAuthentication interface used by the driver.
9+
* This allows token providers to be used with the existing authentication system.
10+
*/
11+
export default class TokenProviderAuthenticator implements IAuthentication {
12+
private readonly tokenProvider: ITokenProvider;
13+
14+
private readonly context: IClientContext;
15+
16+
private readonly headers: HeadersInit;
17+
18+
/**
19+
* Creates a new TokenProviderAuthenticator.
20+
* @param tokenProvider - The token provider to use for authentication
21+
* @param context - The client context for logging
22+
* @param headers - Additional headers to include with each request
23+
*/
24+
constructor(tokenProvider: ITokenProvider, context: IClientContext, headers?: HeadersInit) {
25+
this.tokenProvider = tokenProvider;
26+
this.context = context;
27+
this.headers = headers ?? {};
28+
}
29+
30+
async authenticate(): Promise<HeadersInit> {
31+
const logger = this.context.getLogger();
32+
const providerName = this.tokenProvider.getName();
33+
34+
logger.log(LogLevel.debug, `TokenProviderAuthenticator: getting token from ${providerName}`);
35+
36+
let token = await this.tokenProvider.getToken();
37+
38+
if (token.isExpired()) {
39+
logger.log(
40+
LogLevel.warn,
41+
`TokenProviderAuthenticator: token from ${providerName} is expired, requesting a new token`,
42+
);
43+
44+
token = await this.tokenProvider.getToken();
45+
46+
if (token.isExpired()) {
47+
const message = `TokenProviderAuthenticator: token from ${providerName} is still expired after refresh`;
48+
logger.log(LogLevel.error, message);
49+
throw new Error(message);
50+
}
51+
}
52+
53+
return token.setAuthHeader(this.headers);
54+
}
55+
}
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
export { default as ITokenProvider } from './ITokenProvider';
2+
export { default as Token } from './Token';
3+
export { default as StaticTokenProvider } from './StaticTokenProvider';
4+
export { default as ExternalTokenProvider, TokenCallback } from './ExternalTokenProvider';
5+
export { default as TokenProviderAuthenticator } from './TokenProviderAuthenticator';

lib/contracts/IDBSQLClient.ts

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ import IDBSQLSession from './IDBSQLSession';
33
import IAuthentication from '../connection/contracts/IAuthentication';
44
import { ProxyOptions } from '../connection/contracts/IConnectionOptions';
55
import OAuthPersistence from '../connection/auth/DatabricksOAuth/OAuthPersistence';
6+
import ITokenProvider from '../connection/auth/tokenProvider/ITokenProvider';
7+
import { TokenCallback } from '../connection/auth/tokenProvider/ExternalTokenProvider';
68

79
export interface ClientOptions {
810
logger?: IDBSQLLogger;
@@ -24,6 +26,18 @@ type AuthOptions =
2426
| {
2527
authType: 'custom';
2628
provider: IAuthentication;
29+
}
30+
| {
31+
authType: 'token-provider';
32+
tokenProvider: ITokenProvider;
33+
}
34+
| {
35+
authType: 'external-token';
36+
getToken: TokenCallback;
37+
}
38+
| {
39+
authType: 'static-token';
40+
staticToken: string;
2741
};
2842

2943
export type ConnectionOptions = {

0 commit comments

Comments
 (0)