Skip to content

Commit ba4d0b4

Browse files
madhav-dbmadhav-db
committed
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
1 parent 72d1e80 commit ba4d0b4

12 files changed

Lines changed: 847 additions & 0 deletions

File tree

‎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
}

‎lib/connection/auth/tokenProvider/ExternalTokenProvider.ts‎

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

‎lib/connection/auth/tokenProvider/ITokenProvider.ts‎

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

‎lib/connection/auth/tokenProvider/StaticTokenProvider.ts‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
import ITokenProvider from './ITokenProvider';
2+
import Token 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(
17+
accessToken: string,
18+
options?: {
19+
tokenType?: string;
20+
expiresAt?: Date;
21+
refreshToken?: string;
22+
scopes?: string[];
23+
},
24+
) {
25+
this.token = new Token(accessToken, options);
26+
}
27+
28+
/**
29+
* Creates a StaticTokenProvider from a JWT string.
30+
* The expiration time will be extracted from the JWT payload.
31+
* @param jwt - The JWT token string
32+
* @param options - Optional token configuration
33+
*/
34+
static fromJWT(
35+
jwt: string,
36+
options?: {
37+
tokenType?: string;
38+
refreshToken?: string;
39+
scopes?: string[];
40+
},
41+
): StaticTokenProvider {
42+
const token = Token.fromJWT(jwt, options);
43+
return new StaticTokenProvider(token.accessToken, {
44+
tokenType: token.tokenType,
45+
expiresAt: token.expiresAt,
46+
refreshToken: token.refreshToken,
47+
scopes: token.scopes,
48+
});
49+
}
50+
51+
async getToken(): Promise<Token> {
52+
return this.token;
53+
}
54+
55+
getName(): string {
56+
return 'StaticTokenProvider';
57+
}
58+
}

‎lib/connection/auth/tokenProvider/Token.ts‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
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+
* Represents an access token with optional metadata and lifecycle management.
11+
*/
12+
export default class Token {
13+
private readonly _accessToken: string;
14+
15+
private readonly _tokenType: string;
16+
17+
private readonly _expiresAt?: Date;
18+
19+
private readonly _refreshToken?: string;
20+
21+
private readonly _scopes?: string[];
22+
23+
constructor(
24+
accessToken: string,
25+
options?: {
26+
tokenType?: string;
27+
expiresAt?: Date;
28+
refreshToken?: string;
29+
scopes?: string[];
30+
},
31+
) {
32+
this._accessToken = accessToken;
33+
this._tokenType = options?.tokenType ?? 'Bearer';
34+
this._expiresAt = options?.expiresAt;
35+
this._refreshToken = options?.refreshToken;
36+
this._scopes = options?.scopes;
37+
}
38+
39+
/**
40+
* The access token string.
41+
*/
42+
get accessToken(): string {
43+
return this._accessToken;
44+
}
45+
46+
/**
47+
* The token type (e.g., "Bearer").
48+
*/
49+
get tokenType(): string {
50+
return this._tokenType;
51+
}
52+
53+
/**
54+
* The expiration time of the token, if known.
55+
*/
56+
get expiresAt(): Date | undefined {
57+
return this._expiresAt;
58+
}
59+
60+
/**
61+
* The refresh token, if available.
62+
*/
63+
get refreshToken(): string | undefined {
64+
return this._refreshToken;
65+
}
66+
67+
/**
68+
* The scopes associated with this token.
69+
*/
70+
get scopes(): string[] | undefined {
71+
return this._scopes;
72+
}
73+
74+
/**
75+
* Checks if the token has expired, including a safety buffer.
76+
* Returns false if expiration time is unknown.
77+
*/
78+
isExpired(): boolean {
79+
if (!this._expiresAt) {
80+
return false;
81+
}
82+
const now = new Date();
83+
const bufferMs = EXPIRATION_BUFFER_SECONDS * 1000;
84+
return this._expiresAt.getTime() - bufferMs <= now.getTime();
85+
}
86+
87+
/**
88+
* Sets the Authorization header on the provided headers object.
89+
* @param headers - The headers object to modify
90+
* @returns The modified headers object with Authorization set
91+
*/
92+
setAuthHeader(headers: HeadersInit): HeadersInit {
93+
return {
94+
...headers,
95+
Authorization: `${this._tokenType} ${this._accessToken}`,
96+
};
97+
}
98+
99+
/**
100+
* Creates a Token from a JWT string, extracting the expiration time from the payload.
101+
* @param jwt - The JWT token string
102+
* @param options - Additional token options (tokenType, refreshToken, scopes)
103+
* @returns A new Token instance with expiration extracted from the JWT
104+
* @throws Error if the JWT cannot be decoded
105+
*/
106+
static fromJWT(
107+
jwt: string,
108+
options?: {
109+
tokenType?: string;
110+
refreshToken?: string;
111+
scopes?: string[];
112+
},
113+
): Token {
114+
let expiresAt: Date | undefined;
115+
116+
try {
117+
const parts = jwt.split('.');
118+
if (parts.length >= 2) {
119+
const payload = Buffer.from(parts[1], 'base64').toString('utf8');
120+
const decoded = JSON.parse(payload);
121+
if (typeof decoded.exp === 'number') {
122+
expiresAt = new Date(decoded.exp * 1000);
123+
}
124+
}
125+
} catch {
126+
// If we can't decode the JWT, we'll proceed without expiration info
127+
// The server will validate the token anyway
128+
}
129+
130+
return new Token(jwt, {
131+
tokenType: options?.tokenType,
132+
expiresAt,
133+
refreshToken: options?.refreshToken,
134+
scopes: options?.scopes,
135+
});
136+
}
137+
138+
/**
139+
* Converts the token to a plain object for serialization.
140+
*/
141+
toJSON(): Record<string, unknown> {
142+
return {
143+
accessToken: this._accessToken,
144+
tokenType: this._tokenType,
145+
expiresAt: this._expiresAt?.toISOString(),
146+
refreshToken: this._refreshToken,
147+
scopes: this._scopes,
148+
};
149+
}
150+
}

‎lib/connection/auth/tokenProvider/TokenProviderAuthenticator.ts‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
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(
25+
tokenProvider: ITokenProvider,
26+
context: IClientContext,
27+
headers?: HeadersInit,
28+
) {
29+
this.tokenProvider = tokenProvider;
30+
this.context = context;
31+
this.headers = headers ?? {};
32+
}
33+
34+
async authenticate(): Promise<HeadersInit> {
35+
const logger = this.context.getLogger();
36+
const providerName = this.tokenProvider.getName();
37+
38+
logger.log(LogLevel.debug, `TokenProviderAuthenticator: getting token from ${providerName}`);
39+
40+
const token = await this.tokenProvider.getToken();
41+
42+
if (token.isExpired()) {
43+
logger.log(LogLevel.warn, `TokenProviderAuthenticator: token from ${providerName} is expired`);
44+
}
45+
46+
return token.setAuthHeader(this.headers);
47+
}
48+
}

‎lib/connection/auth/tokenProvider/index.ts‎

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: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,17 @@ 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';
67

78
export interface ClientOptions {
89
logger?: IDBSQLLogger;
910
}
1011

12+
/**
13+
* Type for the callback function that retrieves tokens from external sources.
14+
*/
15+
export type TokenCallback = () => Promise<string>;
16+
1117
type AuthOptions =
1218
| {
1319
authType?: 'access-token';
@@ -24,6 +30,18 @@ type AuthOptions =
2430
| {
2531
authType: 'custom';
2632
provider: IAuthentication;
33+
}
34+
| {
35+
authType: 'token-provider';
36+
tokenProvider: ITokenProvider;
37+
}
38+
| {
39+
authType: 'external-token';
40+
getToken: TokenCallback;
41+
}
42+
| {
43+
authType: 'static-token';
44+
staticToken: string;
2745
};
2846

2947
export type ConnectionOptions = {

0 commit comments

Comments
 (0)