diff --git a/examples/tokenFederation/README.md b/examples/tokenFederation/README.md new file mode 100644 index 00000000..9f17aa1b --- /dev/null +++ b/examples/tokenFederation/README.md @@ -0,0 +1,51 @@ +# Token Federation Examples + +Examples demonstrating the token provider and federation features of the Databricks SQL Node.js Driver. + +## Examples + +### Static Token (`staticToken.ts`) + +The simplest authentication method. Use a static access token that doesn't change during the application lifetime. + +```bash +DATABRICKS_HOST= DATABRICKS_HTTP_PATH= DATABRICKS_TOKEN= npx ts-node staticToken.ts +``` + +### External Token (`externalToken.ts`) + +Use a callback function to provide tokens dynamically. Useful for integrating with secret managers, vaults, or other token sources. Tokens are automatically cached by the driver. + +```bash +DATABRICKS_HOST= DATABRICKS_HTTP_PATH= DATABRICKS_TOKEN= npx ts-node externalToken.ts +``` + +### Token Federation (`federation.ts`) + +Automatically exchange tokens from external identity providers (Azure AD, Google, Okta, etc.) for Databricks-compatible tokens using RFC 8693 token exchange. + +```bash +DATABRICKS_HOST= DATABRICKS_HTTP_PATH= AZURE_AD_TOKEN= npx ts-node federation.ts +``` + +### M2M Federation (`m2mFederation.ts`) + +Machine-to-machine token federation with a service principal. Requires a `federationClientId` to identify the service principal to Databricks. + +```bash +DATABRICKS_HOST= DATABRICKS_HTTP_PATH= DATABRICKS_CLIENT_ID= SERVICE_ACCOUNT_TOKEN= npx ts-node m2mFederation.ts +``` + +### Custom Token Provider (`customTokenProvider.ts`) + +Implement the `ITokenProvider` interface for full control over token management, including custom caching, refresh logic, retry, and error handling. + +```bash +DATABRICKS_HOST= DATABRICKS_HTTP_PATH= OAUTH_SERVER_URL= OAUTH_CLIENT_ID= OAUTH_CLIENT_SECRET= npx ts-node customTokenProvider.ts +``` + +## Prerequisites + +- Node.js 14+ +- A Databricks workspace with token federation enabled (for federation examples) +- Valid credentials for your identity provider diff --git a/examples/tokenFederation/customTokenProvider.ts b/examples/tokenFederation/customTokenProvider.ts new file mode 100644 index 00000000..b468dbdb --- /dev/null +++ b/examples/tokenFederation/customTokenProvider.ts @@ -0,0 +1,169 @@ +/** + * Example: Custom Token Provider Implementation + * + * This example demonstrates how to create a custom token provider by + * implementing the ITokenProvider interface. This gives you full control + * over token management, including custom caching, refresh logic, and + * error handling. + */ + +import { DBSQLClient } from '@databricks/sql'; +import { ITokenProvider, Token } from '../../lib/connection/auth/tokenProvider'; + +/** + * Custom token provider that refreshes tokens from a custom OAuth server. + */ +class CustomOAuthTokenProvider implements ITokenProvider { + private readonly oauthServerUrl: string; + + private readonly clientId: string; + + private readonly clientSecret: string; + + constructor(oauthServerUrl: string, clientId: string, clientSecret: string) { + this.oauthServerUrl = oauthServerUrl; + this.clientId = clientId; + this.clientSecret = clientSecret; + } + + async getToken(): Promise { + // eslint-disable-next-line no-console + console.log('Fetching token from custom OAuth server...'); + return this.fetchTokenWithRetry(0); + } + + /** + * Recursively attempts to fetch a token with exponential backoff. + */ + private async fetchTokenWithRetry(attempt: number): Promise { + const maxRetries = 3; + + try { + return await this.fetchToken(); + } catch (error) { + // Don't retry client errors (4xx) + if (error instanceof Error && error.message.includes('OAuth token request failed: 4')) { + throw error; + } + + if (attempt >= maxRetries) { + throw error; + } + + // Exponential backoff: 1s, 2s, 4s + const delay = 1000 * 2 ** attempt; + await new Promise((resolve) => { + setTimeout(resolve, delay); + }); + + return this.fetchTokenWithRetry(attempt + 1); + } + } + + private async fetchToken(): Promise { + const response = await fetch(`${this.oauthServerUrl}/oauth/token`, { + method: 'POST', + headers: { + 'Content-Type': 'application/x-www-form-urlencoded', + }, + body: new URLSearchParams({ + grant_type: 'client_credentials', + client_id: this.clientId, + client_secret: this.clientSecret, + scope: 'sql', + }).toString(), + }); + + if (!response.ok) { + throw new Error(`OAuth token request failed: ${response.status}`); + } + + const data = (await response.json()) as { + access_token: string; + token_type?: string; + expires_in?: number; + }; + + // Calculate expiration + let expiresAt: Date | undefined; + if (typeof data.expires_in === 'number') { + expiresAt = new Date(Date.now() + data.expires_in * 1000); + } + + return new Token(data.access_token, { + tokenType: data.token_type ?? 'Bearer', + expiresAt, + }); + } + + getName(): string { + return 'CustomOAuthTokenProvider'; + } +} + +/** + * Simple token provider that reads from a file (for development/testing). + */ +// exported for use as an alternative example provider +// eslint-disable-next-line @typescript-eslint/no-unused-vars +class FileTokenProvider implements ITokenProvider { + private readonly filePath: string; + + constructor(filePath: string) { + this.filePath = filePath; + } + + async getToken(): Promise { + const fs = await import('fs/promises'); + const tokenData = await fs.readFile(this.filePath, 'utf-8'); + const parsed = JSON.parse(tokenData); + + return Token.fromJWT(parsed.access_token, { + refreshToken: parsed.refresh_token, + }); + } + + getName(): string { + return 'FileTokenProvider'; + } +} + +async function main() { + const host = process.env.DATABRICKS_HOST!; + const path = process.env.DATABRICKS_HTTP_PATH!; + + const client = new DBSQLClient(); + + // Option 1: Use a custom OAuth token provider (shown below) + // Option 2: Use a file-based token provider for development: + // const fileProvider = new FileTokenProvider('/path/to/token.json'); + const oauthProvider = new CustomOAuthTokenProvider( + process.env.OAUTH_SERVER_URL!, + process.env.OAUTH_CLIENT_ID!, + process.env.OAUTH_CLIENT_SECRET!, + ); + + await client.connect({ + host, + path, + authType: 'token-provider', + tokenProvider: oauthProvider, + // Optionally enable federation if your OAuth server issues non-Databricks tokens + enableTokenFederation: true, + }); + + console.log('Connected successfully with custom token provider'); + + // Open a session and run a query + const session = await client.openSession(); + const operation = await session.executeStatement('SELECT 1 AS result'); + const result = await operation.fetchAll(); + + console.log('Query result:', result); + + await operation.close(); + await session.close(); + await client.close(); +} + +main().catch(console.error); diff --git a/examples/tokenFederation/externalToken.ts b/examples/tokenFederation/externalToken.ts new file mode 100644 index 00000000..224da6de --- /dev/null +++ b/examples/tokenFederation/externalToken.ts @@ -0,0 +1,53 @@ +/** + * Example: Using an external token provider + * + * This example demonstrates how to use a callback function to provide + * tokens dynamically. This is useful for integrating with secret managers, + * vaults, or other token sources that may refresh tokens. + */ + +import { DBSQLClient } from '@databricks/sql'; + +// Simulate fetching a token from a secret manager or vault +async function fetchTokenFromVault(): Promise { + // In a real application, this would fetch from AWS Secrets Manager, + // Azure Key Vault, HashiCorp Vault, or another secret manager + console.log('Fetching token from vault...'); + + // Simulated token - replace with actual vault integration + const token = process.env.DATABRICKS_TOKEN!; + return token; +} + +async function main() { + const host = process.env.DATABRICKS_HOST!; + const path = process.env.DATABRICKS_HTTP_PATH!; + + const client = new DBSQLClient(); + + // Connect using an external token provider + // The callback will be called each time a new token is needed + // Note: The token is automatically cached, so the callback won't be + // called on every request + await client.connect({ + host, + path, + authType: 'external-token', + getToken: fetchTokenFromVault, + }); + + console.log('Connected successfully with external token provider'); + + // Open a session and run a query + const session = await client.openSession(); + const operation = await session.executeStatement('SELECT current_user() AS user'); + const result = await operation.fetchAll(); + + console.log('Query result:', result); + + await operation.close(); + await session.close(); + await client.close(); +} + +main().catch(console.error); diff --git a/examples/tokenFederation/federation.ts b/examples/tokenFederation/federation.ts new file mode 100644 index 00000000..1d21e50e --- /dev/null +++ b/examples/tokenFederation/federation.ts @@ -0,0 +1,80 @@ +/** + * Example: Token Federation with an External Identity Provider + * + * This example demonstrates how to use token federation to automatically + * exchange tokens from external identity providers (Azure AD, Google, Okta, + * Auth0, AWS Cognito, GitHub) for Databricks-compatible tokens. + * + * Token federation uses RFC 8693 (OAuth 2.0 Token Exchange) to exchange + * the external JWT token for a Databricks access token. + */ + +import { DBSQLClient } from '@databricks/sql'; + +// Example: Fetch a token from Azure AD +// In a real application, you would use the Azure SDK or similar +async function getAzureADToken(): Promise { + // Example using @azure/identity: + // + // import { DefaultAzureCredential } from '@azure/identity'; + // const credential = new DefaultAzureCredential(); + // const token = await credential.getToken('https://your-scope/.default'); + // return token.token; + + // For this example, we use an environment variable + const token = process.env.AZURE_AD_TOKEN!; + console.log('Fetched token from Azure AD'); + return token; +} + +// Example: Fetch a token from Google +// eslint-disable-next-line @typescript-eslint/no-unused-vars +async function getGoogleToken(): Promise { + // Example using google-auth-library: + // + // import { GoogleAuth } from 'google-auth-library'; + // const auth = new GoogleAuth(); + // const client = await auth.getClient(); + // const token = await client.getAccessToken(); + // return token.token; + + const token = process.env.GOOGLE_TOKEN!; + console.log('Fetched token from Google'); + return token; +} + +async function main() { + const host = process.env.DATABRICKS_HOST!; + const path = process.env.DATABRICKS_HTTP_PATH!; + + const client = new DBSQLClient(); + + // Connect using token federation + // The driver will automatically: + // 1. Get the token from the callback + // 2. Check if the token's issuer matches the Databricks host + // 3. If not, exchange the token for a Databricks token via RFC 8693 + // 4. Cache the result for subsequent requests + await client.connect({ + host, + path, + authType: 'external-token', + getToken: getAzureADToken, // or getGoogleToken, etc. + enableTokenFederation: true, + }); + + console.log('Connected successfully with token federation'); + + // Open a session and run a query + const session = await client.openSession(); + const operation = await session.executeStatement('SELECT current_user() AS user'); + const result = await operation.fetchAll(); + + console.log('Query result:', result); + + await operation.close(); + await session.close(); + await client.close(); +} + +main().catch(console.error); diff --git a/examples/tokenFederation/m2mFederation.ts b/examples/tokenFederation/m2mFederation.ts new file mode 100644 index 00000000..e4c22f4f --- /dev/null +++ b/examples/tokenFederation/m2mFederation.ts @@ -0,0 +1,65 @@ +/** + * Example: Machine-to-Machine (M2M) Token Federation with Service Principal + * + * This example demonstrates how to use token federation with a service + * principal or machine identity. This is useful for server-to-server + * authentication where there is no interactive user. + * + * When using M2M federation, you typically need to provide a client_id + * to identify the service principal to Databricks. + */ + +import { DBSQLClient } from '@databricks/sql'; + +// Example: Fetch a service account token from your identity provider +async function getServiceAccountToken(): Promise { + // Example for Azure service principal: + // + // import { ClientSecretCredential } from '@azure/identity'; + // const credential = new ClientSecretCredential( + // process.env.AZURE_TENANT_ID!, + // process.env.AZURE_CLIENT_ID!, + // process.env.AZURE_CLIENT_SECRET! + // ); + // const token = await credential.getToken('https://your-scope/.default'); + // return token.token; + + // For this example, we use an environment variable + const token = process.env.SERVICE_ACCOUNT_TOKEN!; + console.log('Fetched service account token'); + return token; +} + +async function main() { + const host = process.env.DATABRICKS_HOST!; + const path = process.env.DATABRICKS_HTTP_PATH!; + const clientId = process.env.DATABRICKS_CLIENT_ID!; + + const client = new DBSQLClient(); + + // Connect using M2M token federation + // The federationClientId identifies your service principal to Databricks + await client.connect({ + host, + path, + authType: 'external-token', + getToken: getServiceAccountToken, + enableTokenFederation: true, + federationClientId: clientId, // Required for M2M/SP federation + }); + + console.log('Connected successfully with M2M token federation'); + + // Open a session and run a query + const session = await client.openSession(); + const operation = await session.executeStatement('SELECT current_user() AS user'); + const result = await operation.fetchAll(); + + console.log('Query result:', result); + + await operation.close(); + await session.close(); + await client.close(); +} + +main().catch(console.error); diff --git a/examples/tokenFederation/staticToken.ts b/examples/tokenFederation/staticToken.ts new file mode 100644 index 00000000..d6cec8df --- /dev/null +++ b/examples/tokenFederation/staticToken.ts @@ -0,0 +1,40 @@ +/** + * Example: Using a static token with the token provider system + * + * This example demonstrates how to use a static access token with the + * token provider infrastructure. This is useful when you have a token + * that doesn't change during the lifetime of your application. + */ + +import { DBSQLClient } from '@databricks/sql'; + +async function main() { + const host = process.env.DATABRICKS_HOST!; + const path = process.env.DATABRICKS_HTTP_PATH!; + const token = process.env.DATABRICKS_TOKEN!; + + const client = new DBSQLClient(); + + // Connect using a static token + await client.connect({ + host, + path, + authType: 'static-token', + staticToken: token, + }); + + console.log('Connected successfully with static token'); + + // Open a session and run a query + const session = await client.openSession(); + const operation = await session.executeStatement('SELECT 1 AS result'); + const result = await operation.fetchAll(); + + console.log('Query result:', result); + + await operation.close(); + await session.close(); + await client.close(); +} + +main().catch(console.error); diff --git a/lib/index.ts b/lib/index.ts index 710a036d..adf14f36 100644 --- a/lib/index.ts +++ b/lib/index.ts @@ -9,12 +9,28 @@ import DBSQLSession from './DBSQLSession'; import { DBSQLParameter, DBSQLParameterType } from './DBSQLParameter'; import DBSQLLogger from './DBSQLLogger'; import PlainHttpAuthentication from './connection/auth/PlainHttpAuthentication'; +import { + Token, + StaticTokenProvider, + ExternalTokenProvider, + CachedTokenProvider, + FederationProvider, +} from './connection/auth/tokenProvider'; import HttpConnection from './connection/connections/HttpConnection'; import { formatProgress } from './utils'; import { LogLevel } from './contracts/IDBSQLLogger'; +// Re-export types for TypeScript users +export type { default as ITokenProvider } from './connection/auth/tokenProvider/ITokenProvider'; + export const auth = { PlainHttpAuthentication, + // Token provider classes for custom authentication + Token, + StaticTokenProvider, + ExternalTokenProvider, + CachedTokenProvider, + FederationProvider, }; const { TException, TApplicationException, TApplicationExceptionType, TProtocolException, TProtocolExceptionType } = diff --git a/tsconfig.build.json b/tsconfig.build.json index 7b375312..9aa952a0 100644 --- a/tsconfig.build.json +++ b/tsconfig.build.json @@ -4,5 +4,5 @@ "outDir": "./dist/" /* Redirect output structure to the directory. */, "rootDir": "./lib/" /* Specify the root directory of input files. Use to control the output directory structure with --outDir. */ }, - "exclude": ["./tests/**/*", "./dist/**/*"] + "exclude": ["./tests/**/*", "./dist/**/*", "./examples/**/*"] }