|
| 1 | +/** |
| 2 | + * Example: Custom Token Provider Implementation |
| 3 | + * |
| 4 | + * This example demonstrates how to create a custom token provider by |
| 5 | + * implementing the ITokenProvider interface. This gives you full control |
| 6 | + * over token management, including custom caching, refresh logic, and |
| 7 | + * error handling. |
| 8 | + */ |
| 9 | + |
| 10 | +import { DBSQLClient } from '@databricks/sql'; |
| 11 | +import { ITokenProvider, Token } from '../../lib/connection/auth/tokenProvider'; |
| 12 | + |
| 13 | +/** |
| 14 | + * Custom token provider that refreshes tokens from a custom OAuth server. |
| 15 | + */ |
| 16 | +class CustomOAuthTokenProvider implements ITokenProvider { |
| 17 | + private readonly oauthServerUrl: string; |
| 18 | + |
| 19 | + private readonly clientId: string; |
| 20 | + |
| 21 | + private readonly clientSecret: string; |
| 22 | + |
| 23 | + constructor(oauthServerUrl: string, clientId: string, clientSecret: string) { |
| 24 | + this.oauthServerUrl = oauthServerUrl; |
| 25 | + this.clientId = clientId; |
| 26 | + this.clientSecret = clientSecret; |
| 27 | + } |
| 28 | + |
| 29 | + async getToken(): Promise<Token> { |
| 30 | + // eslint-disable-next-line no-console |
| 31 | + console.log('Fetching token from custom OAuth server...'); |
| 32 | + return this.fetchTokenWithRetry(0); |
| 33 | + } |
| 34 | + |
| 35 | + /** |
| 36 | + * Recursively attempts to fetch a token with exponential backoff. |
| 37 | + */ |
| 38 | + private async fetchTokenWithRetry(attempt: number): Promise<Token> { |
| 39 | + const maxRetries = 3; |
| 40 | + |
| 41 | + try { |
| 42 | + return await this.fetchToken(); |
| 43 | + } catch (error) { |
| 44 | + // Don't retry client errors (4xx) |
| 45 | + if (error instanceof Error && error.message.includes('OAuth token request failed: 4')) { |
| 46 | + throw error; |
| 47 | + } |
| 48 | + |
| 49 | + if (attempt >= maxRetries) { |
| 50 | + throw error; |
| 51 | + } |
| 52 | + |
| 53 | + // Exponential backoff: 1s, 2s, 4s |
| 54 | + const delay = 1000 * 2 ** attempt; |
| 55 | + await new Promise<void>((resolve) => { |
| 56 | + setTimeout(resolve, delay); |
| 57 | + }); |
| 58 | + |
| 59 | + return this.fetchTokenWithRetry(attempt + 1); |
| 60 | + } |
| 61 | + } |
| 62 | + |
| 63 | + private async fetchToken(): Promise<Token> { |
| 64 | + const response = await fetch(`${this.oauthServerUrl}/oauth/token`, { |
| 65 | + method: 'POST', |
| 66 | + headers: { |
| 67 | + 'Content-Type': 'application/x-www-form-urlencoded', |
| 68 | + }, |
| 69 | + body: new URLSearchParams({ |
| 70 | + grant_type: 'client_credentials', |
| 71 | + client_id: this.clientId, |
| 72 | + client_secret: this.clientSecret, |
| 73 | + scope: 'sql', |
| 74 | + }).toString(), |
| 75 | + }); |
| 76 | + |
| 77 | + if (!response.ok) { |
| 78 | + throw new Error(`OAuth token request failed: ${response.status}`); |
| 79 | + } |
| 80 | + |
| 81 | + const data = (await response.json()) as { |
| 82 | + access_token: string; |
| 83 | + token_type?: string; |
| 84 | + expires_in?: number; |
| 85 | + }; |
| 86 | + |
| 87 | + // Calculate expiration |
| 88 | + let expiresAt: Date | undefined; |
| 89 | + if (typeof data.expires_in === 'number') { |
| 90 | + expiresAt = new Date(Date.now() + data.expires_in * 1000); |
| 91 | + } |
| 92 | + |
| 93 | + return new Token(data.access_token, { |
| 94 | + tokenType: data.token_type ?? 'Bearer', |
| 95 | + expiresAt, |
| 96 | + }); |
| 97 | + } |
| 98 | + |
| 99 | + getName(): string { |
| 100 | + return 'CustomOAuthTokenProvider'; |
| 101 | + } |
| 102 | +} |
| 103 | + |
| 104 | +/** |
| 105 | + * Simple token provider that reads from a file (for development/testing). |
| 106 | + */ |
| 107 | +// exported for use as an alternative example provider |
| 108 | +// eslint-disable-next-line @typescript-eslint/no-unused-vars |
| 109 | +class FileTokenProvider implements ITokenProvider { |
| 110 | + private readonly filePath: string; |
| 111 | + |
| 112 | + constructor(filePath: string) { |
| 113 | + this.filePath = filePath; |
| 114 | + } |
| 115 | + |
| 116 | + async getToken(): Promise<Token> { |
| 117 | + const fs = await import('fs/promises'); |
| 118 | + const tokenData = await fs.readFile(this.filePath, 'utf-8'); |
| 119 | + const parsed = JSON.parse(tokenData); |
| 120 | + |
| 121 | + return Token.fromJWT(parsed.access_token, { |
| 122 | + refreshToken: parsed.refresh_token, |
| 123 | + }); |
| 124 | + } |
| 125 | + |
| 126 | + getName(): string { |
| 127 | + return 'FileTokenProvider'; |
| 128 | + } |
| 129 | +} |
| 130 | + |
| 131 | +async function main() { |
| 132 | + const host = process.env.DATABRICKS_HOST!; |
| 133 | + const path = process.env.DATABRICKS_HTTP_PATH!; |
| 134 | + |
| 135 | + const client = new DBSQLClient(); |
| 136 | + |
| 137 | + // Option 1: Use a custom OAuth token provider (shown below) |
| 138 | + // Option 2: Use a file-based token provider for development: |
| 139 | + // const fileProvider = new FileTokenProvider('/path/to/token.json'); |
| 140 | + const oauthProvider = new CustomOAuthTokenProvider( |
| 141 | + process.env.OAUTH_SERVER_URL!, |
| 142 | + process.env.OAUTH_CLIENT_ID!, |
| 143 | + process.env.OAUTH_CLIENT_SECRET!, |
| 144 | + ); |
| 145 | + |
| 146 | + await client.connect({ |
| 147 | + host, |
| 148 | + path, |
| 149 | + authType: 'token-provider', |
| 150 | + tokenProvider: oauthProvider, |
| 151 | + // Optionally enable federation if your OAuth server issues non-Databricks tokens |
| 152 | + enableTokenFederation: true, |
| 153 | + }); |
| 154 | + |
| 155 | + console.log('Connected successfully with custom token provider'); |
| 156 | + |
| 157 | + // Open a session and run a query |
| 158 | + const session = await client.openSession(); |
| 159 | + const operation = await session.executeStatement('SELECT 1 AS result'); |
| 160 | + const result = await operation.fetchAll(); |
| 161 | + |
| 162 | + console.log('Query result:', result); |
| 163 | + |
| 164 | + await operation.close(); |
| 165 | + await session.close(); |
| 166 | + await client.close(); |
| 167 | +} |
| 168 | + |
| 169 | +main().catch(console.error); |
0 commit comments