From 651159ec72f3009e7213d8efdd5b22bd25745cf3 Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Thu, 20 Aug 2026 00:12:51 +0000 Subject: [PATCH] feat(kernel): support JWT private-key M2M auth on useKernel Add JWT private-key client-assertion auth (RFC 7523) to the kernel backend. On `authType: 'databricks-oauth'`, supplying `oauthJwtKeyFile` selects the JWT flow: the kernel signs a short-lived assertion with the private key instead of sending a client secret and owns the token lifecycle (`authMode: 'OAuthM2mJwt'`). - KernelAuth: new JWT branch in buildKernelConnectionOptions (checked before the U2M/M2M-secret split; a private-key file is unambiguous JWT M2M intent), plus the OAuthM2mJwt native option shape. Requires oauthClientId + oauthJwtKid; optional oauthJwtPassphrase / oauthJwtAlgorithm / oauthScopes / tokenUrl. Mutually exclusive with oauthClientSecret. Also threads tokenUrl through the existing M2m branch. - IDBSQLClient: new oauthJwt* + tokenUrl fields on the databricks-oauth ConnectionOptions member. - DBSQLClient: on the useKernel path, do not build the connector's own OAuth provider (it eagerly starts the U2M browser flow / M2M exchange before the kernel is consulted); hand over a minimal PAT provider only when a token is present. Mirrors the Python connector. - tests: 9 unit tests for JWT routing / precedence / validation. Verified end-to-end: SELECT 1 via useKernel against an Azure Databricks warehouse, authenticated by Entra ID with a JWT private-key assertion (tokenUrl pointed at the Entra token endpoint). Requires a @databricks/databricks-sql-kernel build with JWT + tokenUrl support (kernel PRs #249 merged, #275 for tokenUrl). Signed-off-by: Rahul Singhal --- CHANGELOG.md | 4 + lib/DBSQLClient.ts | 25 ++++- lib/contracts/IDBSQLClient.ts | 16 ++++ lib/kernel/KernelAuth.ts | 72 +++++++++++++- tests/unit/kernel/auth-m2m-jwt.test.ts | 124 +++++++++++++++++++++++++ 5 files changed, 235 insertions(+), 6 deletions(-) create mode 100644 tests/unit/kernel/auth-m2m-jwt.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 036c71a4..7a03779d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Release History +## Unreleased + +- Kernel backend (`useKernel: true`): OAuth **M2M with a JWT private-key client assertion** (RFC 7523) is now supported. On the `databricks-oauth` auth type, supplying `oauthJwtKeyFile` (with `oauthClientId` + `oauthJwtKid`, optional `oauthJwtPassphrase` / `oauthJwtAlgorithm` / `oauthScopes`, and `tokenUrl` for the IdP token endpoint) selects the JWT client-assertion flow: the kernel signs a short-lived assertion with the private key instead of sending a client secret, and owns the token lifecycle. A private-key file is treated as unambiguous JWT M2M intent and is mutually exclusive with `oauthClientSecret`. `tokenUrl` points the grant at the workspace's OAuth IdP (e.g. Entra ID for Azure Databricks), which is required because Databricks-native OIDC does not advertise the `private_key_jwt` method. Also fixes the kernel path to not eagerly build the connector's own OAuth provider (which could start the U2M browser flow before the kernel is consulted). Verified end-to-end against an Azure Databricks warehouse via Entra ID. Requires a `@databricks/databricks-sql-kernel` build with JWT + `tokenUrl` support. + ## 2.0.0 **Breaking changes — completes the security cleanup that 1.17.0 could not do without breaking changes.** diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index f021edf9..e72a548d 100644 --- a/lib/DBSQLClient.ts +++ b/lib/DBSQLClient.ts @@ -721,14 +721,31 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I // hit endpoints that don't carry the workspace in their URL path. this.config.customHeaders = this.buildCustomHeaders(options.path, options.customHeaders); - this.authProvider = this.createAuthProvider(options, authProvider); - - this.connectionProvider = this.createConnectionProvider(options); - // M0: `useKernel` is consumed via a non-exported internal-options cast so it // doesn't ship in the public `.d.ts`. Mirrors Python's `kwargs.get("use_kernel")` // pattern (see databricks-sql-python/src/databricks/sql/session.py). const internalOptions = options as ConnectionOptions & InternalConnectionOptions; + + // On the kernel path the kernel owns the full auth lifecycle (it resolves + // M2M / U2M / JWT purely from the raw options via `buildKernelConnectionOptions`). + // We must NOT build the connector's own OAuth provider here: for OAuth it + // eagerly runs the U2M browser flow / M2M token exchange at connect() time + // (a telemetry / feature-flag client calls `authProvider.authenticate()`), + // racing — and conflicting with — the kernel's auth. So for `useKernel` we + // hand over only a minimal PAT provider when a `token` is present, and + // `undefined` otherwise. Mirrors Python's use_kernel auth-provider handling. + if (internalOptions.useKernel) { + const { token } = options as { token?: string }; + this.authProvider = + typeof token === 'string' && token.length > 0 + ? new PlainHttpAuthentication({ username: 'token', password: token, context: this }) + : undefined; + } else { + this.authProvider = this.createAuthProvider(options, authProvider); + } + + this.connectionProvider = this.createConnectionProvider(options); + const backend = internalOptions.useKernel ? new KernelBackend({ context: this }) : new ThriftBackend({ diff --git a/lib/contracts/IDBSQLClient.ts b/lib/contracts/IDBSQLClient.ts index bbaa4c69..f2e1621a 100644 --- a/lib/contracts/IDBSQLClient.ts +++ b/lib/contracts/IDBSQLClient.ts @@ -26,6 +26,22 @@ type AuthOptions = // U2M flow to `['sql', 'offline_access']` (parity with the Thrift driver's // `defaultOAuthScopes`), overriding the kernel's bare `all-apis offline_access`. oauthScopes?: Array; + // JWT private-key M2M (RFC 7523 client assertion) — KERNEL BACKEND ONLY + // (`useKernel: true`). Supplying `oauthJwtKeyFile` selects the JWT + // client-assertion flow: the kernel signs a short-lived assertion with the + // private key instead of sending a client secret. Requires `oauthClientId` + // and `oauthJwtKid`. Optional `oauthJwtPassphrase` (encrypted PKCS#8 key), + // `oauthJwtAlgorithm` (default `RS256`), `oauthScopes`, and `tokenUrl` (the + // IdP token endpoint — required when auth is against an external IdP such as + // Entra ID, which is where `private_key_jwt` is supported). Mutually + // exclusive with `oauthClientSecret`. + oauthJwtKeyFile?: string; + oauthJwtKid?: string; + oauthJwtPassphrase?: string; + oauthJwtAlgorithm?: string; + // OAuth token endpoint override (kernel backend). Points the M2M / + // JWT client-assertion grant at the workspace's IdP token endpoint. + tokenUrl?: string; } | { authType: 'custom'; diff --git a/lib/kernel/KernelAuth.ts b/lib/kernel/KernelAuth.ts index 7cf99afa..4e3abd03 100644 --- a/lib/kernel/KernelAuth.ts +++ b/lib/kernel/KernelAuth.ts @@ -230,6 +230,19 @@ export type KernelNativeConnectionOptions = KernelSessionDefaults & oauthClientId: string; oauthClientSecret: string; oauthScopes?: Array; + tokenUrl?: string; + } + | { + hostName: string; + httpPath: string; + authMode: 'OAuthM2mJwt'; + oauthClientId: string; + jwtKeyFile: string; + jwtKid: string; + jwtPassphrase?: string; + jwtAlgorithm?: string; + oauthScopes?: Array; + tokenUrl?: string; } | { hostName: string; @@ -602,6 +615,11 @@ export function buildKernelConnectionOptions(options: ConnectionOptions): Kernel azureTenantId?: string; useDatabricksOAuthInAzure?: boolean; persistence?: unknown; + oauthJwtKeyFile?: string; + oauthJwtKid?: string; + oauthJwtPassphrase?: string; + oauthJwtAlgorithm?: string; + tokenUrl?: string; }; if (authType === undefined || authType === 'access-token') { @@ -637,6 +655,55 @@ export function buildKernelConnectionOptions(options: ConnectionOptions): Kernel ); } + // JWT private-key M2M (RFC 7523 client assertion). A private-key file is + // unambiguous JWT M2M intent, so this is checked before the U2M/M2M + // secret split. The kernel signs a short-lived assertion with the key + // (`authMode: 'OAuthM2mJwt'`) instead of sending a client secret. Requires + // `oauthClientId` (assertion issuer/subject) and `oauthJwtKid` (key id). + // Mutually exclusive with `oauthClientSecret`. + if (oauth.oauthJwtKeyFile !== undefined) { + if (oauth.oauthClientSecret !== undefined) { + throw new HiveDriverError( + 'kernel backend: cannot supply both `oauthJwtKeyFile` (JWT private-key M2M) ' + + 'and `oauthClientSecret` (shared-secret M2M). Pick one.', + ); + } + if (oauth.persistence !== undefined) { + throw new HiveDriverError( + 'kernel backend: `persistence` is not supported on JWT private-key M2M ' + + '(M2M tokens have no refresh token; the kernel re-issues on expiry).', + ); + } + if (oauth.oauthClientId === undefined) { + throw new AuthenticationError( + 'kernel backend: JWT private-key M2M (`oauthJwtKeyFile`) requires `oauthClientId` ' + + '(the service principal / OAuth client id used as the assertion issuer and subject).', + ); + } + if (oauth.oauthJwtKid === undefined) { + throw new AuthenticationError( + 'kernel backend: JWT private-key M2M (`oauthJwtKeyFile`) requires `oauthJwtKid` ' + + '(the key id written into the JWT header so the IdP can select the registered public key).', + ); + } + const jwt = { + ...base, + authMode: 'OAuthM2mJwt' as const, + oauthClientId: oauth.oauthClientId, + jwtKeyFile: oauth.oauthJwtKeyFile, + jwtKid: oauth.oauthJwtKid, + // Configurable (parity with pyo3); defaults to `['all-apis']` in the kernel. + oauthScopes: + Array.isArray(oauth.oauthScopes) && oauth.oauthScopes.length > 0 ? oauth.oauthScopes : M2M_DEFAULT_SCOPES, + }; + return { + ...jwt, + ...(oauth.oauthJwtPassphrase !== undefined ? { jwtPassphrase: oauth.oauthJwtPassphrase } : {}), + ...(oauth.oauthJwtAlgorithm !== undefined ? { jwtAlgorithm: oauth.oauthJwtAlgorithm } : {}), + ...(oauth.tokenUrl !== undefined ? { tokenUrl: oauth.tokenUrl } : {}), + }; + } + // Flow selector + client-id resolution mirror the Thrift driver EXACTLY // (`DBSQLClient.createAuthProvider`, DBSQLClient.ts:220): // flow = oauthClientSecret === undefined ? U2M : M2M (strict undefined) @@ -680,9 +747,9 @@ export function buildKernelConnectionOptions(options: ConnectionOptions): Kernel '(M2M tokens have no refresh token; the kernel re-issues on expiry).', ); } - return { + const m2m = { ...base, - authMode: 'OAuthM2m', + authMode: 'OAuthM2m' as const, // Thrift: `getClientId()` = `oauthClientId ?? defaultClientId`. oauthClientId: oauth.oauthClientId ?? DEFAULT_OAUTH_CLIENT_ID, oauthClientSecret: oauth.oauthClientSecret, @@ -690,6 +757,7 @@ export function buildKernelConnectionOptions(options: ConnectionOptions): Kernel oauthScopes: Array.isArray(oauth.oauthScopes) && oauth.oauthScopes.length > 0 ? oauth.oauthScopes : M2M_DEFAULT_SCOPES, }; + return oauth.tokenUrl !== undefined ? { ...m2m, tokenUrl: oauth.tokenUrl } : m2m; } throw new HiveDriverError( diff --git a/tests/unit/kernel/auth-m2m-jwt.test.ts b/tests/unit/kernel/auth-m2m-jwt.test.ts new file mode 100644 index 00000000..7c8ae216 --- /dev/null +++ b/tests/unit/kernel/auth-m2m-jwt.test.ts @@ -0,0 +1,124 @@ +// Copyright (c) 2026 Databricks, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import { expect } from 'chai'; +import { buildKernelConnectionOptions } from '../../../lib/kernel/KernelAuth'; +import { ConnectionOptions } from '../../../lib/contracts/IDBSQLClient'; +import HiveDriverError from '../../../lib/errors/HiveDriverError'; +import AuthenticationError from '../../../lib/errors/AuthenticationError'; + +// A private-key file selects JWT client-assertion M2M (RFC 7523); the kernel +// signs a short-lived assertion with the key instead of sending a secret. +const baseJwt = { + host: 'example.azuredatabricks.net', + path: '/sql/1.0/warehouses/abc', + authType: 'databricks-oauth' as const, + oauthClientId: 'sp-uuid', + oauthJwtKeyFile: '/keys/jwt.pem', + oauthJwtKid: 'kid-1', +}; + +describe('KernelAuth — OAuth M2M JWT private-key auth flow', () => { + it('routes oauthJwtKeyFile to authMode OAuthM2mJwt with the required fields', () => { + const native = buildKernelConnectionOptions(baseJwt as ConnectionOptions); + expect(native.authMode).to.equal('OAuthM2mJwt'); + const jwt = native as { + oauthClientId?: string; + jwtKeyFile?: string; + jwtKid?: string; + oauthScopes?: string[]; + }; + expect(jwt.oauthClientId).to.equal('sp-uuid'); + expect(jwt.jwtKeyFile).to.equal('/keys/jwt.pem'); + expect(jwt.jwtKid).to.equal('kid-1'); + // Defaults to the M2M scope (parity with pyo3 / the secret M2M path). + expect(jwt.oauthScopes).to.deep.equal(['all-apis']); + }); + + it('forwards optional passphrase / algorithm / tokenUrl / scopes when present', () => { + const native = buildKernelConnectionOptions({ + ...baseJwt, + oauthJwtPassphrase: 'pw', + oauthJwtAlgorithm: 'ES256', + tokenUrl: 'https://login.microsoftonline.com/tenant/oauth2/v2.0/token', + oauthScopes: ['2ff814a6-.../.default'], + } as ConnectionOptions); + const jwt = native as { + jwtPassphrase?: string; + jwtAlgorithm?: string; + tokenUrl?: string; + oauthScopes?: string[]; + }; + expect(jwt.jwtPassphrase).to.equal('pw'); + expect(jwt.jwtAlgorithm).to.equal('ES256'); + expect(jwt.tokenUrl).to.equal('https://login.microsoftonline.com/tenant/oauth2/v2.0/token'); + expect(jwt.oauthScopes).to.deep.equal(['2ff814a6-.../.default']); + }); + + it('omits optional fields when not supplied', () => { + const native = buildKernelConnectionOptions(baseJwt as ConnectionOptions); + expect(native).to.not.have.property('jwtPassphrase'); + expect(native).to.not.have.property('jwtAlgorithm'); + expect(native).to.not.have.property('tokenUrl'); + }); + + it('takes precedence over the shared-secret M2M / U2M split', () => { + // A private key present makes this JWT M2M regardless of anything else + // (no secret ⇒ would otherwise be U2M). + const native = buildKernelConnectionOptions(baseJwt as ConnectionOptions); + expect(native.authMode).to.equal('OAuthM2mJwt'); + }); + + it('rejects oauthJwtKeyFile together with oauthClientSecret (ambiguous)', () => { + expect(() => + buildKernelConnectionOptions({ + ...baseJwt, + oauthClientSecret: 'shh', + } as ConnectionOptions), + ).to.throw(HiveDriverError, /both `oauthJwtKeyFile`.*`oauthClientSecret`/); + }); + + it('requires oauthClientId', () => { + const { oauthClientId, ...noClientId } = baseJwt; + expect(() => buildKernelConnectionOptions(noClientId as ConnectionOptions)).to.throw( + AuthenticationError, + /requires `oauthClientId`/, + ); + }); + + it('requires oauthJwtKid', () => { + const { oauthJwtKid, ...noKid } = baseJwt; + expect(() => buildKernelConnectionOptions(noKid as ConnectionOptions)).to.throw( + AuthenticationError, + /requires `oauthJwtKid`/, + ); + }); + + it('rejects persistence on the JWT M2M path', () => { + expect(() => + buildKernelConnectionOptions({ + ...baseJwt, + persistence: {} as never, + } as ConnectionOptions), + ).to.throw(HiveDriverError, /persistence/); + }); + + it('prepends `/` to the path on the JWT branch too', () => { + const native = buildKernelConnectionOptions({ + ...baseJwt, + path: 'sql/1.0/warehouses/abc', + } as ConnectionOptions); + expect((native as { httpPath: string }).httpPath).to.equal('/sql/1.0/warehouses/abc'); + }); +});