feat: implement OAuth M2M authentication and access token management

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?)

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",

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<string, CachedOAuthToken>()
+const oauthTokenRequests = new Map<string, Promise<CachedOAuthToken>>()
+
+export function isOAuthM2MAuthInfo(auth: AuthInfo): auth is OAuthM2MAuthInfo {
+  return !auth.token && Boolean(auth.clientId && auth.clientSecret)
+}
+
+export async function getAccessToken(auth: AuthInfo): Promise<string> {
+  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<string> {
+  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<CachedOAuthToken> {
+  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(' ')
+}

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<T>(
 
   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<T>(
       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<T>(
       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 ||

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'