Skip to content

Commit 921fa0d

Browse files
committed
Merge branch 'token-federation-pr2-federation-caching' into token-federation-pr3-examples
2 parents d767c4a + bf392eb commit 921fa0d

5 files changed

Lines changed: 151 additions & 71 deletions

File tree

lib/connection/auth/tokenProvider/FederationProvider.ts

Lines changed: 105 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,21 @@ const DEFAULT_SCOPE = 'sql';
2828
*/
2929
const REQUEST_TIMEOUT_MS = 30000;
3030

31+
/**
32+
* Maximum number of retry attempts for transient errors.
33+
*/
34+
const MAX_RETRY_ATTEMPTS = 3;
35+
36+
/**
37+
* Base delay in milliseconds for exponential backoff.
38+
*/
39+
const RETRY_BASE_DELAY_MS = 1000;
40+
41+
/**
42+
* HTTP status codes that are considered retryable.
43+
*/
44+
const RETRYABLE_STATUS_CODES = new Set([429, 500, 502, 503, 504]);
45+
3146
/**
3247
* A token provider that wraps another provider with automatic token federation.
3348
* When the base provider returns a token from a different issuer, this provider
@@ -111,6 +126,7 @@ export default class FederationProvider implements ITokenProvider {
111126

112127
/**
113128
* Exchanges the token for a Databricks-compatible token using RFC 8693.
129+
* Includes retry logic for transient errors with exponential backoff.
114130
* @param token - The token to exchange
115131
* @returns The exchanged token
116132
*/
@@ -128,47 +144,102 @@ export default class FederationProvider implements ITokenProvider {
128144
params.append('client_id', this.clientId);
129145
}
130146

131-
const controller = new AbortController();
132-
const timeoutId = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
147+
let lastError: Error | undefined;
133148

134-
try {
135-
const response = await fetch(url, {
136-
method: 'POST',
137-
headers: {
138-
'Content-Type': 'application/x-www-form-urlencoded',
139-
},
140-
body: params.toString(),
141-
signal: controller.signal,
142-
});
143-
144-
if (!response.ok) {
145-
const errorText = await response.text();
146-
throw new Error(`Token exchange failed: ${response.status} ${response.statusText} - ${errorText}`);
149+
for (let attempt = 0; attempt <= MAX_RETRY_ATTEMPTS; attempt++) {
150+
if (attempt > 0) {
151+
// Exponential backoff: 1s, 2s, 4s
152+
const delay = RETRY_BASE_DELAY_MS * Math.pow(2, attempt - 1);
153+
await this.sleep(delay);
147154
}
148155

149-
const data = (await response.json()) as {
150-
access_token?: string;
151-
token_type?: string;
152-
expires_in?: number;
153-
};
154-
155-
if (!data.access_token) {
156-
throw new Error('Token exchange response missing access_token');
156+
const controller = new AbortController();
157+
const timeoutId = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
158+
159+
try {
160+
const response = await fetch(url, {
161+
method: 'POST',
162+
headers: {
163+
'Content-Type': 'application/x-www-form-urlencoded',
164+
},
165+
body: params.toString(),
166+
signal: controller.signal,
167+
});
168+
169+
if (!response.ok) {
170+
const errorText = await response.text();
171+
const error = new Error(`Token exchange failed: ${response.status} ${response.statusText} - ${errorText}`);
172+
173+
// Check if this is a retryable status code
174+
if (RETRYABLE_STATUS_CODES.has(response.status) && attempt < MAX_RETRY_ATTEMPTS) {
175+
lastError = error;
176+
continue;
177+
}
178+
179+
throw error;
180+
}
181+
182+
const data = (await response.json()) as {
183+
access_token?: string;
184+
token_type?: string;
185+
expires_in?: number;
186+
};
187+
188+
if (!data.access_token) {
189+
throw new Error('Token exchange response missing access_token');
190+
}
191+
192+
// Calculate expiration from expires_in
193+
let expiresAt: Date | undefined;
194+
if (typeof data.expires_in === 'number') {
195+
expiresAt = new Date(Date.now() + data.expires_in * 1000);
196+
}
197+
198+
return new Token(data.access_token, {
199+
tokenType: data.token_type ?? 'Bearer',
200+
expiresAt,
201+
});
202+
} catch (error) {
203+
clearTimeout(timeoutId);
204+
205+
// Retry on network errors (timeout, connection issues)
206+
if (this.isRetryableError(error) && attempt < MAX_RETRY_ATTEMPTS) {
207+
lastError = error instanceof Error ? error : new Error(String(error));
208+
continue;
209+
}
210+
211+
throw error;
212+
} finally {
213+
clearTimeout(timeoutId);
157214
}
215+
}
158216

159-
// Calculate expiration from expires_in
160-
let expiresAt: Date | undefined;
161-
if (typeof data.expires_in === 'number') {
162-
expiresAt = new Date(Date.now() + data.expires_in * 1000);
163-
}
217+
// If we exhausted all retries, throw the last error
218+
throw lastError ?? new Error('Token exchange failed after retries');
219+
}
164220

165-
return new Token(data.access_token, {
166-
tokenType: data.token_type ?? 'Bearer',
167-
expiresAt,
168-
});
169-
} finally {
170-
clearTimeout(timeoutId);
221+
/**
222+
* Determines if an error is retryable (network errors, timeouts).
223+
*/
224+
private isRetryableError(error: unknown): boolean {
225+
if (error instanceof Error) {
226+
// AbortError from timeout
227+
if (error.name === 'AbortError') {
228+
return true;
229+
}
230+
// Network errors from node-fetch
231+
if (error.name === 'FetchError') {
232+
return true;
233+
}
171234
}
235+
return false;
236+
}
237+
238+
/**
239+
* Sleeps for the specified duration.
240+
*/
241+
private sleep(ms: number): Promise<void> {
242+
return new Promise((resolve) => setTimeout(resolve, ms));
172243
}
173244

174245
/**

lib/connection/auth/tokenProvider/StaticTokenProvider.ts

Lines changed: 3 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import ITokenProvider from './ITokenProvider';
2-
import Token from './Token';
2+
import Token, { TokenOptions, TokenFromJWTOptions } from './Token';
33

44
/**
55
* A token provider that returns a static token.
@@ -13,15 +13,7 @@ export default class StaticTokenProvider implements ITokenProvider {
1313
* @param accessToken - The access token string
1414
* @param options - Optional token configuration (tokenType, expiresAt, refreshToken, scopes)
1515
*/
16-
constructor(
17-
accessToken: string,
18-
options?: {
19-
tokenType?: string;
20-
expiresAt?: Date;
21-
refreshToken?: string;
22-
scopes?: string[];
23-
},
24-
) {
16+
constructor(accessToken: string, options?: TokenOptions) {
2517
this.token = new Token(accessToken, options);
2618
}
2719

@@ -31,14 +23,7 @@ export default class StaticTokenProvider implements ITokenProvider {
3123
* @param jwt - The JWT token string
3224
* @param options - Optional token configuration
3325
*/
34-
static fromJWT(
35-
jwt: string,
36-
options?: {
37-
tokenType?: string;
38-
refreshToken?: string;
39-
scopes?: string[];
40-
},
41-
): StaticTokenProvider {
26+
static fromJWT(jwt: string, options?: TokenFromJWTOptions): StaticTokenProvider {
4227
const token = Token.fromJWT(jwt, options);
4328
return new StaticTokenProvider(token.accessToken, {
4429
tokenType: token.tokenType,

lib/connection/auth/tokenProvider/Token.ts

Lines changed: 24 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,26 @@ import { HeadersInit } from 'node-fetch';
66
*/
77
const EXPIRATION_BUFFER_SECONDS = 30;
88

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+
929
/**
1030
* Represents an access token with optional metadata and lifecycle management.
1131
*/
@@ -20,15 +40,7 @@ export default class Token {
2040

2141
private readonly _scopes?: string[];
2242

23-
constructor(
24-
accessToken: string,
25-
options?: {
26-
tokenType?: string;
27-
expiresAt?: Date;
28-
refreshToken?: string;
29-
scopes?: string[];
30-
},
31-
) {
43+
constructor(accessToken: string, options?: TokenOptions) {
3244
this._accessToken = accessToken;
3345
this._tokenType = options?.tokenType ?? 'Bearer';
3446
this._expiresAt = options?.expiresAt;
@@ -101,17 +113,11 @@ export default class Token {
101113
* If the JWT cannot be decoded, the token is created without expiration info.
102114
* The server will validate the token anyway, so decoding failures are handled gracefully.
103115
* @param jwt - The JWT token string
104-
* @param options - Additional token options (tokenType, refreshToken, scopes)
116+
* @param options - Additional token options (tokenType, refreshToken, scopes).
117+
* Note: expiresAt is not accepted here as it is extracted from the JWT payload.
105118
* @returns A new Token instance with expiration extracted from the JWT (if available)
106119
*/
107-
static fromJWT(
108-
jwt: string,
109-
options?: {
110-
tokenType?: string;
111-
refreshToken?: string;
112-
scopes?: string[];
113-
},
114-
): Token {
120+
static fromJWT(jwt: string, options?: TokenFromJWTOptions): Token {
115121
let expiresAt: Date | undefined;
116122

117123
try {

lib/connection/auth/tokenProvider/TokenProviderAuthenticator.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,9 @@ export default class TokenProviderAuthenticator implements IAuthentication {
3636
const token = await this.tokenProvider.getToken();
3737

3838
if (token.isExpired()) {
39-
logger.log(LogLevel.warn, `TokenProviderAuthenticator: token from ${providerName} is expired`);
39+
const message = `TokenProviderAuthenticator: token from ${providerName} is expired`;
40+
logger.log(LogLevel.error, message);
41+
throw new Error(message);
4042
}
4143

4244
return token.setAuthHeader(this.headers);

tests/unit/connection/auth/tokenProvider/TokenProviderAuthenticator.test.ts

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,5 +104,21 @@ describe('TokenProviderAuthenticator', () => {
104104
expect(e).to.equal(error);
105105
}
106106
});
107+
108+
it('should throw error when token is expired', async () => {
109+
const provider = new MockTokenProvider('my-access-token', 'TestProvider');
110+
const expiredDate = new Date(Date.now() - 60000); // 1 minute ago
111+
provider.setToken(new Token('expired-token', { expiresAt: expiredDate }));
112+
const authenticator = new TokenProviderAuthenticator(provider, context);
113+
114+
try {
115+
await authenticator.authenticate();
116+
expect.fail('Should have thrown an error');
117+
} catch (e) {
118+
expect(e).to.be.instanceOf(Error);
119+
expect((e as Error).message).to.include('expired');
120+
expect((e as Error).message).to.include('TestProvider');
121+
}
122+
});
107123
});
108124
});

0 commit comments

Comments
 (0)