From 8c072b7b1b455c208dd1cf9596188120526cea14 Mon Sep 17 00:00:00 2001 From: BumSeok Hwang Date: Wed, 2 Sep 2026 23:48:12 +0000 Subject: [PATCH] feat: implement OAuth M2M authentication and access token management --- README.md | 31 +++- package.json | 2 +- src/auth.ts | 135 +++++++++++++++ src/http.ts | 42 ++++- src/types.ts | 24 ++- test/auth.spec.ts | 325 ++++++++++++++++++++++++++++++++++++ test/s3.integration.spec.ts | 25 ++- 7 files changed, 569 insertions(+), 15 deletions(-) create mode 100644 src/auth.ts create mode 100644 test/auth.spec.ts diff --git a/README.md b/README.md index 90dc905..2f4fbd0 100644 --- a/README.md +++ b/README.md @@ -20,6 +20,7 @@ The goal is simple: stream big results with stable memory usage and without forc ## Highlights - Direct REST calls to Statement Execution API. +- Personal Access Token and OAuth M2M authentication. - Optimized polling with server-side wait (up to 50s) before falling back to client polling. - Query metrics support via Query History API (`enableMetrics` option). - Efficient external link handling: merge chunks into a single stream. @@ -48,6 +49,30 @@ const rows = await fetchAll(result, auth, { format: 'JSON_OBJECT' }) console.log(rows) // [{ value: 1 }] ``` +## Sample (OAuth M2M) +```ts +import { executeStatement, fetchAll } from '@bitofsky/databricks-sql' + +const auth = { + clientId: process.env.DATABRICKS_CLIENT_ID!, + clientSecret: process.env.DATABRICKS_CLIENT_SECRET!, + host: process.env.DATABRICKS_HOST!, + httpPath: process.env.DATABRICKS_HTTP_PATH!, +} + +const result = await executeStatement('SELECT 1 AS value', auth) +const rows = await fetchAll(result, auth, { format: 'JSON_OBJECT' }) +console.log(rows) +``` + +OAuth access tokens are requested from the workspace `/oidc/v1/token` endpoint, +cached in memory, and renewed five minutes before expiration. A rejected OAuth +access token (`401`, or `403` with `Invalid access token`) invalidates the cache +and retries the request once. +Provide either `token` or both `clientId` and `clientSecret`, but not both modes. +OAuth requests always include the required `sql` and `query-history` scopes. +Use `scopes` to request any additional scopes. + ## Sample (Streaming + Presigned URL) Stream external links into S3 with gzip compression, then return a single presigned URL. @@ -178,10 +203,12 @@ try { ### AuthInfo ```ts type AuthInfo = { - token: string host: string httpPath: string -} +} & ( + | { token: string } + | { clientId: string; clientSecret: string; scopes?: string[] } +) ``` ### executeStatement(query, auth, options?) diff --git a/package.json b/package.json index dbeff40..c62a431 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@bitofsky/databricks-sql", - "version": "1.0.5", + "version": "1.0.6", "description": "Databricks SQL client for Node.js - Direct REST API without SDK", "main": "dist/index.cjs", "module": "dist/index.js", diff --git a/src/auth.ts b/src/auth.ts new file mode 100644 index 0000000..b1dbf6b --- /dev/null +++ b/src/auth.ts @@ -0,0 +1,135 @@ +import type { AuthInfo, OAuthM2MAuthInfo } from './types.js' + +import { DatabricksSqlError, HttpError } from './errors.js' +import { buildUrl } from './util.js' + +const OAUTH_TOKEN_PATH = '/oidc/v1/token' +const REQUIRED_OAUTH_SCOPES = ['sql', 'query-history'] +const DEFAULT_TOKEN_EXPIRES_IN_SECONDS = 3600 +const TOKEN_REFRESH_BUFFER_MS = 5 * 60 * 1000 + +type OAuthTokenResponse = { + access_token?: string + expires_in?: number +} + +type CachedOAuthToken = { + accessToken: string + refreshAt: number +} + +const oauthTokenCache = new Map() +const oauthTokenRequests = new Map>() + +export function isOAuthM2MAuthInfo(auth: AuthInfo): auth is OAuthM2MAuthInfo { + return !auth.token && Boolean(auth.clientId && auth.clientSecret) +} + +export async function getAccessToken(auth: AuthInfo): Promise { + if (auth.token) { + if (auth.clientId || auth.clientSecret || auth.scopes) + throw new DatabricksSqlError( + 'AuthInfo cannot contain token with OAuth settings', + 'INVALID_AUTH_CONFIG' + ) + + return auth.token + } + + if (!isOAuthM2MAuthInfo(auth)) + throw new DatabricksSqlError( + 'AuthInfo must contain either token or clientId and clientSecret', + 'INVALID_AUTH_CONFIG' + ) + + return getOAuthAccessToken(auth) +} + +export function invalidateAccessToken(auth: OAuthM2MAuthInfo): void { + oauthTokenCache.delete(getOAuthTokenCacheKey(auth)) +} + +async function getOAuthAccessToken(auth: OAuthM2MAuthInfo): Promise { + const cacheKey = getOAuthTokenCacheKey(auth) + const cachedToken = oauthTokenCache.get(cacheKey) + + if (cachedToken && Date.now() < cachedToken.refreshAt) + return cachedToken.accessToken + + const pendingRequest = oauthTokenRequests.get(cacheKey) + if (pendingRequest) + return (await pendingRequest).accessToken + + const request = requestOAuthToken(auth) + oauthTokenRequests.set(cacheKey, request) + + try { + const token = await request + oauthTokenCache.set(cacheKey, token) + return token.accessToken + } finally { + oauthTokenRequests.delete(cacheKey) + } +} + +async function requestOAuthToken( + auth: OAuthM2MAuthInfo +): Promise { + const credentials = Buffer.from( + `${encodeURIComponent(auth.clientId)}:${encodeURIComponent(auth.clientSecret)}` + ).toString('base64') + const body = new URLSearchParams({ + grant_type: 'client_credentials', + scope: getOAuthScope(auth), + }) + const response = await fetch(buildUrl(auth.host, OAUTH_TOKEN_PATH), { + method: 'POST', + headers: { + Authorization: `Basic ${credentials}`, + 'Content-Type': 'application/x-www-form-urlencoded', + Accept: 'application/json', + }, + body, + }) + + if (!response.ok) { + const errorBody = await response.text().catch(() => '') + throw new HttpError( + response.status, + response.statusText, + `OAuth token request failed: ${errorBody || response.statusText}` + ) + } + + const result = await response.json() as OAuthTokenResponse + if (!result.access_token) + throw new DatabricksSqlError( + 'OAuth token response does not contain access_token', + 'INVALID_OAUTH_TOKEN_RESPONSE' + ) + + const expiresInSeconds = + typeof result.expires_in === 'number' && Number.isFinite(result.expires_in) + ? result.expires_in + : DEFAULT_TOKEN_EXPIRES_IN_SECONDS + const refreshInMs = Math.max( + 0, + expiresInSeconds * 1000 - TOKEN_REFRESH_BUFFER_MS + ) + + return { + accessToken: result.access_token, + refreshAt: Date.now() + refreshInMs, + } +} + +function getOAuthTokenCacheKey(auth: OAuthM2MAuthInfo): string { + return `${new URL(buildUrl(auth.host, '/')).origin}:${auth.clientId}:${getOAuthScope(auth)}` +} + +function getOAuthScope(auth: OAuthM2MAuthInfo): string { + return [...new Set([...REQUIRED_OAUTH_SCOPES, ...(auth.scopes ?? [])])] + .filter(Boolean) + .sort() + .join(' ') +} diff --git a/src/http.ts b/src/http.ts index 03f531a..8336d48 100644 --- a/src/http.ts +++ b/src/http.ts @@ -1,4 +1,9 @@ import type { AuthInfo } from './types.js' +import { + getAccessToken, + invalidateAccessToken, + isOAuthM2MAuthInfo, +} from './auth.js' import { HttpError, AuthenticationError, @@ -31,18 +36,20 @@ export async function httpRequest( let lastError: Error | undefined let retryDelay = INITIAL_RETRY_DELAY_MS + let oauthRetryAttempted = false for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) { if (signal?.aborted) throw new AbortError() try { + const accessToken = await getAccessToken(auth) // Build a minimal fetch init, skipping undefined values. const fetchInit = Object.fromEntries( Object.entries({ method, headers: { - Authorization: `Bearer ${auth.token}`, + Authorization: `Bearer ${accessToken}`, 'Content-Type': 'application/json', Accept: 'application/json', }, @@ -57,7 +64,27 @@ export async function httpRequest( if (response.ok) return (await response.json()) as T - // Authentication error (no retry) + const isTokenRejected = + response.status === 401 || + ( + response.status === 403 && + (await response.clone().text().catch(() => '')).includes( + 'Invalid access token' + ) + ) + + if ( + isTokenRejected && + isOAuthM2MAuthInfo(auth) && + !oauthRetryAttempted && + attempt < MAX_RETRIES + ) { + invalidateAccessToken(auth) + oauthRetryAttempted = true + continue + } + + // Authentication error if (response.status === 401) throw new AuthenticationError() @@ -98,6 +125,17 @@ export async function httpRequest( throw new HttpError(response.status, response.statusText, errorBody) } catch (err) { + if ( + err instanceof HttpError && + err.status >= 500 && + attempt < MAX_RETRIES + ) { + lastError = err + await delay(retryDelay, signal) + retryDelay *= 2 + continue + } + // Re-throw known errors if ( err instanceof AbortError || diff --git a/src/types.ts b/src/types.ts index 30ba52d..00605be 100644 --- a/src/types.ts +++ b/src/types.ts @@ -1,15 +1,31 @@ import type { Readable } from 'node:stream' -/** Authentication information for Databricks API */ -export type AuthInfo = { - /** Databricks Personal Access Token */ - token: string +type AuthInfoBase = { /** Databricks workspace host (e.g., ...cloud.databricks.com) */ host: string /** SQL warehouse HTTP path (e.g., /sql/1.0/warehouses/abc123) */ httpPath: string } +/** Authentication information using a Databricks Personal Access Token */ +export type PersonalAccessTokenAuthInfo = AuthInfoBase & { + token: string + clientId?: never + clientSecret?: never + scopes?: never +} + +/** Authentication information using Databricks OAuth M2M */ +export type OAuthM2MAuthInfo = AuthInfoBase & { + token?: never + clientId: string + clientSecret: string + scopes?: string[] +} + +/** Authentication information for Databricks API */ +export type AuthInfo = PersonalAccessTokenAuthInfo | OAuthM2MAuthInfo + /** Statement execution states */ export type StatementState = | 'PENDING' diff --git a/test/auth.spec.ts b/test/auth.spec.ts new file mode 100644 index 0000000..974397e --- /dev/null +++ b/test/auth.spec.ts @@ -0,0 +1,325 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { executeStatement } from '../src/api/index.js' +import type { AuthInfo, OAuthM2MAuthInfo } from '../src/types.js' +import { mockInlineResult } from './mocks.js' + +function oauthAuth(id: string): OAuthM2MAuthInfo { + return { + host: 'test.cloud.databricks.com', + httpPath: '/sql/1.0/warehouses/abc123def456', + clientId: `client-${id}`, + clientSecret: `secret-${id}`, + } +} + +function jsonResponse(value: unknown) { + return { + ok: true, + status: 200, + statusText: 'OK', + json: () => Promise.resolve(value), + } +} + +function errorResponse( + status: number, + statusText: string, + value: unknown +): Response { + return new Response(JSON.stringify(value), { + status, + statusText, + headers: { 'Content-Type': 'application/json' }, + }) +} + +afterEach(() => { + vi.restoreAllMocks() + vi.useRealTimers() +}) + +describe('authentication', () => { + it('keeps Personal Access Token authentication unchanged', async () => { + const mockFetch = vi.fn().mockResolvedValueOnce(jsonResponse(mockInlineResult)) + vi.stubGlobal('fetch', mockFetch) + + await executeStatement('SELECT 1', { + host: 'test.cloud.databricks.com', + httpPath: '/sql/1.0/warehouses/abc123def456', + token: 'pat-token', + }) + + expect(mockFetch).toHaveBeenCalledTimes(1) + expect(mockFetch).toHaveBeenCalledWith( + expect.stringContaining('/api/2.0/sql/statements'), + expect.objectContaining({ + headers: expect.objectContaining({ + Authorization: 'Bearer pat-token', + }), + }) + ) + }) + + it('exchanges client credentials for an OAuth access token', async () => { + const auth = oauthAuth('exchange') + const mockFetch = vi + .fn() + .mockResolvedValueOnce(jsonResponse({ + access_token: 'oauth-token', + token_type: 'Bearer', + expires_in: 3600, + })) + .mockResolvedValueOnce(jsonResponse(mockInlineResult)) + vi.stubGlobal('fetch', mockFetch) + + await executeStatement('SELECT 1', auth) + + const expectedCredentials = Buffer.from( + `${encodeURIComponent(auth.clientId)}:${encodeURIComponent(auth.clientSecret)}` + ).toString('base64') + expect(mockFetch).toHaveBeenNthCalledWith( + 1, + 'https://test.cloud.databricks.com/oidc/v1/token', + expect.objectContaining({ + method: 'POST', + headers: expect.objectContaining({ + Authorization: `Basic ${expectedCredentials}`, + 'Content-Type': 'application/x-www-form-urlencoded', + }), + body: expect.any(URLSearchParams), + }) + ) + expect(mockFetch).toHaveBeenNthCalledWith( + 2, + expect.stringContaining('/api/2.0/sql/statements'), + expect.objectContaining({ + headers: expect.objectContaining({ + Authorization: 'Bearer oauth-token', + }), + }) + ) + const tokenRequest = mockFetch.mock.calls[0]?.[1] as RequestInit + expect(tokenRequest.body?.toString()).toBe( + 'grant_type=client_credentials&scope=query-history+sql' + ) + }) + + it('adds configured OAuth scopes to the required scopes', async () => { + const auth: OAuthM2MAuthInfo = { + ...oauthAuth('scopes'), + scopes: ['jobs', 'sql'], + } + const mockFetch = vi + .fn() + .mockResolvedValueOnce( + jsonResponse({ access_token: 'scoped-token', expires_in: 3600 }) + ) + .mockResolvedValueOnce(jsonResponse(mockInlineResult)) + vi.stubGlobal('fetch', mockFetch) + + await executeStatement('SELECT 1', auth) + + const tokenRequest = mockFetch.mock.calls[0]?.[1] as RequestInit + expect(tokenRequest.body?.toString()).toBe( + 'grant_type=client_credentials&scope=jobs+query-history+sql' + ) + }) + + it('reuses an OAuth token until its refresh time', async () => { + const auth = oauthAuth('cache') + const mockFetch = vi + .fn() + .mockResolvedValueOnce(jsonResponse({ access_token: 'cached-token', expires_in: 3600 })) + .mockResolvedValueOnce(jsonResponse(mockInlineResult)) + .mockResolvedValueOnce(jsonResponse(mockInlineResult)) + vi.stubGlobal('fetch', mockFetch) + + await executeStatement('SELECT 1', auth) + await executeStatement('SELECT 2', auth) + + expect(mockFetch).toHaveBeenCalledTimes(3) + expect( + mockFetch.mock.calls.filter(([url]) => + String(url).includes('/oidc/v1/token') + ) + ).toHaveLength(1) + }) + + it('shares one OAuth token request across concurrent calls', async () => { + const auth = oauthAuth('concurrent') + let resolveToken: + | ((response: ReturnType) => void) + | undefined + const tokenResponse = new Promise>((resolve) => { + resolveToken = resolve + }) + const mockFetch = vi.fn((url: string | URL | Request) => { + if (String(url).includes('/oidc/v1/token')) + return tokenResponse + return Promise.resolve(jsonResponse(mockInlineResult)) + }) + vi.stubGlobal('fetch', mockFetch) + + const first = executeStatement('SELECT 1', auth) + const second = executeStatement('SELECT 2', auth) + + await vi.waitFor(() => expect(mockFetch).toHaveBeenCalledTimes(1)) + resolveToken?.(jsonResponse({ access_token: 'shared-token', expires_in: 3600 })) + await Promise.all([first, second]) + + expect(mockFetch).toHaveBeenCalledTimes(3) + expect( + mockFetch.mock.calls.filter(([url]) => + String(url).includes('/oidc/v1/token') + ) + ).toHaveLength(1) + }) + + it('refreshes an OAuth token five minutes before expiration', async () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')) + const auth = oauthAuth('expiry') + const mockFetch = vi + .fn() + .mockResolvedValueOnce(jsonResponse({ access_token: 'first-token', expires_in: 3600 })) + .mockResolvedValueOnce(jsonResponse(mockInlineResult)) + .mockResolvedValueOnce(jsonResponse({ access_token: 'second-token', expires_in: 3600 })) + .mockResolvedValueOnce(jsonResponse(mockInlineResult)) + vi.stubGlobal('fetch', mockFetch) + + await executeStatement('SELECT 1', auth) + await vi.advanceTimersByTimeAsync(55 * 60 * 1000) + await executeStatement('SELECT 2', auth) + + expect(mockFetch).toHaveBeenCalledTimes(4) + expect(mockFetch.mock.calls[3]?.[1]).toEqual(expect.objectContaining({ + headers: expect.objectContaining({ Authorization: 'Bearer second-token' }), + })) + }) + + it('refreshes once when an OAuth access token receives 401', async () => { + const auth = oauthAuth('unauthorized') + const mockFetch = vi + .fn() + .mockResolvedValueOnce(jsonResponse({ access_token: 'rejected-token', expires_in: 3600 })) + .mockResolvedValueOnce(errorResponse(401, 'Unauthorized', { + error_code: 401, + message: 'Unauthorized', + })) + .mockResolvedValueOnce(jsonResponse({ access_token: 'renewed-token', expires_in: 3600 })) + .mockResolvedValueOnce(jsonResponse(mockInlineResult)) + vi.stubGlobal('fetch', mockFetch) + + await executeStatement('SELECT 1', auth) + + expect(mockFetch).toHaveBeenCalledTimes(4) + expect(mockFetch.mock.calls[3]?.[1]).toEqual(expect.objectContaining({ + headers: expect.objectContaining({ Authorization: 'Bearer renewed-token' }), + })) + }) + + it('refreshes once when Databricks rejects an OAuth token with 403', async () => { + const auth = oauthAuth('forbidden-token') + const mockFetch = vi + .fn() + .mockResolvedValueOnce( + jsonResponse({ access_token: 'rejected-token', expires_in: 3600 }) + ) + .mockResolvedValueOnce(errorResponse(403, 'Forbidden', { + error_code: 403, + message: 'Invalid access token.', + })) + .mockResolvedValueOnce( + jsonResponse({ access_token: 'renewed-token', expires_in: 3600 }) + ) + .mockResolvedValueOnce(jsonResponse(mockInlineResult)) + vi.stubGlobal('fetch', mockFetch) + + await executeStatement('SELECT 1', auth) + + expect(mockFetch).toHaveBeenCalledTimes(4) + expect(mockFetch.mock.calls[3]?.[1]).toEqual(expect.objectContaining({ + headers: expect.objectContaining({ Authorization: 'Bearer renewed-token' }), + })) + }) + + it('does not refresh an OAuth token for a permission 403', async () => { + const auth = oauthAuth('forbidden-permission') + const mockFetch = vi + .fn() + .mockResolvedValueOnce( + jsonResponse({ access_token: 'valid-token', expires_in: 3600 }) + ) + .mockResolvedValueOnce(errorResponse(403, 'Forbidden', { + error_code: 403, + message: 'Permission denied.', + })) + vi.stubGlobal('fetch', mockFetch) + + await expect(executeStatement('SELECT 1', auth)).rejects.toMatchObject({ + status: 403, + }) + expect(mockFetch).toHaveBeenCalledTimes(2) + }) + + it('includes the OAuth error response in token request failures', async () => { + const auth = oauthAuth('invalid-client') + const mockFetch = vi.fn().mockResolvedValueOnce( + errorResponse(401, 'Unauthorized', { + error: 'invalid_client', + error_description: 'Client authentication failed', + }) + ) + vi.stubGlobal('fetch', mockFetch) + + await expect(executeStatement('SELECT 1', auth)).rejects.toThrow( + /invalid_client.*Client authentication failed/ + ) + }) + + it('retries OAuth token endpoint server errors', async () => { + vi.useFakeTimers() + const auth = oauthAuth('token-server-error') + const mockFetch = vi + .fn() + .mockResolvedValueOnce( + errorResponse(503, 'Service Unavailable', { error: 'unavailable' }) + ) + .mockResolvedValueOnce( + jsonResponse({ access_token: 'oauth-token', expires_in: 3600 }) + ) + .mockResolvedValueOnce(jsonResponse(mockInlineResult)) + vi.stubGlobal('fetch', mockFetch) + + const resultPromise = executeStatement('SELECT 1', auth) + await vi.advanceTimersByTimeAsync(1000) + + await expect(resultPromise).resolves.toMatchObject({ + status: { state: 'SUCCEEDED' }, + }) + expect(mockFetch).toHaveBeenCalledTimes(3) + }) + + it('rejects incomplete or ambiguous authentication settings', async () => { + vi.stubGlobal('fetch', vi.fn()) + const incompleteAuth = { + host: 'test.cloud.databricks.com', + httpPath: '/sql/1.0/warehouses/abc123def456', + clientId: 'client-only', + } as unknown as AuthInfo + const ambiguousAuth = { + host: 'test.cloud.databricks.com', + httpPath: '/sql/1.0/warehouses/abc123def456', + token: 'pat-token', + clientId: 'client', + clientSecret: 'secret', + } as unknown as AuthInfo + + await expect(executeStatement('SELECT 1', incompleteAuth)).rejects.toMatchObject({ + code: 'INVALID_AUTH_CONFIG', + }) + await expect(executeStatement('SELECT 1', ambiguousAuth)).rejects.toMatchObject({ + code: 'INVALID_AUTH_CONFIG', + }) + }) +}) diff --git a/test/s3.integration.spec.ts b/test/s3.integration.spec.ts index affbb4e..b3db690 100644 --- a/test/s3.integration.spec.ts +++ b/test/s3.integration.spec.ts @@ -16,20 +16,33 @@ import { executeStatement, fetchAll, mergeExternalLinks } from '../src/index.js' import type { AuthInfo, MergeExternalLinksResult } from '../src/index.js' const DATABRICKS_TOKEN = process.env.DATABRICKS_TOKEN +const DATABRICKS_CLIENT_ID = process.env.DATABRICKS_CLIENT_ID +const DATABRICKS_CLIENT_SECRET = process.env.DATABRICKS_CLIENT_SECRET const DATABRICKS_HOST = process.env.DATABRICKS_HOST const DATABRICKS_HTTP_PATH = process.env.DATABRICKS_HTTP_PATH const AWS_REGION = process.env.AWS_REGION const S3_BUCKET = process.env.DATABRICKS_SQL_S3_BUCKET const S3_PREFIX = process.env.DATABRICKS_SQL_S3_PREFIX || '' -const auth: AuthInfo = { - token: DATABRICKS_TOKEN || '', - host: DATABRICKS_HOST || '', - httpPath: DATABRICKS_HTTP_PATH || '', -} +const auth: AuthInfo = DATABRICKS_TOKEN + ? { + token: DATABRICKS_TOKEN, + host: DATABRICKS_HOST || '', + httpPath: DATABRICKS_HTTP_PATH || '', + } + : { + clientId: DATABRICKS_CLIENT_ID || '', + clientSecret: DATABRICKS_CLIENT_SECRET || '', + host: DATABRICKS_HOST || '', + httpPath: DATABRICKS_HTTP_PATH || '', + } + +const hasDatabricksAuth = Boolean( + DATABRICKS_TOKEN || (DATABRICKS_CLIENT_ID && DATABRICKS_CLIENT_SECRET) +) const shouldSkip = - !DATABRICKS_TOKEN || + !hasDatabricksAuth || !DATABRICKS_HOST || !DATABRICKS_HTTP_PATH || !AWS_REGION ||