Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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.**
Expand Down
51 changes: 46 additions & 5 deletions lib/DBSQLClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -497,8 +497,20 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
*/
private mapAuthType(options: ConnectionOptions): string {
switch (options.authType) {
case 'databricks-oauth':
case 'databricks-oauth': {
// JWT private-key M2M (kernel-only) presents no `oauthClientSecret`,
// so without this check it would misreport as `external-browser`
// (U2M) — the opposite of its machine-to-machine nature. The field
// lives on the internal options surface (see InternalConnectionOptions)
// and is only honored on the kernel path; gate the label on `useKernel`
// so a Thrift-path connection (which has no JWT branch and would run
// the U2M browser flow) isn't mislabeled `oauth-m2m-jwt`.
const { oauthJwtKeyFile, useKernel } = options as ConnectionOptions & InternalConnectionOptions;
if (useKernel && oauthJwtKeyFile !== undefined) {
return 'oauth-m2m-jwt';
Comment thread
rahuls-db marked this conversation as resolved.
}
return options.oauthClientSecret === undefined ? 'external-browser' : 'oauth-m2m';
}
case 'custom':
return 'custom';
case 'token-provider':
Expand Down Expand Up @@ -721,14 +733,43 @@ 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) {
Comment thread
rahuls-db marked this conversation as resolved.
Comment thread
rahuls-db marked this conversation as resolved.
// The kernel owns auth via the native binding, so a JS-side custom
// `authProvider` (deprecated arg) genuinely can't be plumbed through.
// Warn rather than drop it silently, so a caller who passes one alongside
// `useKernel` can diagnose why their provider isn't used.
if (authProvider) {
this.logger.log(
LogLevel.warn,
'DBSQLClient: a custom authProvider was supplied with useKernel; it is ignored because the ' +
'kernel backend owns authentication via the native binding. Configure auth through the ' +
'connection options (token / OAuth fields) instead.',
);
}
const { token } = options as { token?: string };
Comment thread
rahuls-db marked this conversation as resolved.
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({
Expand Down
51 changes: 51 additions & 0 deletions lib/contracts/InternalConnectionOptions.ts
Original file line number Diff line number Diff line change
Expand Up @@ -74,4 +74,55 @@ export interface InternalConnectionOptions {
* @internal kernel path only.
*/
clientKeyPem?: Buffer | string;

/**
* kernel-only: JWT private-key M2M (RFC 7523 client assertion). Supplying
* `oauthJwtKeyFile` (alongside `authType: 'databricks-oauth'`) 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` (the assertion issuer/subject) and `oauthJwtKid` (the key
* id written into the JWT header). Mutually exclusive with
* `oauthClientSecret`.
*
* These live on the internal options surface — NOT the public
* `databricks-oauth` `AuthOptions` — because the Thrift backend has no
* JWT client-assertion path; exposing them publicly would let a Thrift
* caller set them and have them silently ignored. The kernel path reads
* them via the `InternalConnectionOptions` cast, exactly like `useKernel`
* and the TLS knobs above.
* @internal kernel path only.
*/
oauthJwtKeyFile?: string;

/**
* kernel-only: key id written into the JWT assertion header so the IdP can
* select the registered public key. Required when `oauthJwtKeyFile` is set.
* @internal kernel path only.
*/
oauthJwtKid?: string;

/**
* kernel-only: passphrase for an encrypted PKCS#8 private key
* (`oauthJwtKeyFile`). Omit for an unencrypted key.
* @internal kernel path only.
*/
oauthJwtPassphrase?: string;

/**
* kernel-only: JWT signing algorithm for the client assertion. Defaults to
* `RS256` in the kernel when omitted.
* @internal kernel path only.
*/
oauthJwtAlgorithm?: string;

/**
* kernel-only: OAuth token-endpoint override. Points the M2M /
* JWT client-assertion grant at the workspace's IdP token endpoint —
* required when auth is against an external IdP such as Entra ID, which is
* where `private_key_jwt` is supported. Applies to both shared-secret M2M
* and JWT M2M (auth-method-agnostic, matching JDBC's
* `OAuth2ConnAuthTokenEndpoint`).
* @internal kernel path only.
*/
tokenUrl?: string;
}
80 changes: 76 additions & 4 deletions lib/kernel/KernelAuth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,19 @@ export type KernelNativeConnectionOptions = KernelSessionDefaults &
oauthClientId: string;
oauthClientSecret: string;
oauthScopes?: Array<string>;
tokenUrl?: string;
}
| {
hostName: string;
httpPath: string;
authMode: 'OAuthM2mJwt';
oauthClientId: string;
jwtKeyFile: string;
jwtKid: string;
jwtPassphrase?: string;
jwtAlgorithm?: string;
oauthScopes?: Array<string>;
tokenUrl?: string;
}
| {
hostName: string;
Expand Down Expand Up @@ -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') {
Expand All @@ -611,9 +629,13 @@ export function buildKernelConnectionOptions(options: ConnectionOptions): Kernel
"kernel backend: a non-empty PAT must be supplied via `token` when using `authType: 'access-token'`.",
);
}
if (oauth.oauthClientId !== undefined || oauth.oauthClientSecret !== undefined) {
if (
oauth.oauthClientId !== undefined ||
oauth.oauthClientSecret !== undefined ||
oauth.oauthJwtKeyFile !== undefined
) {
throw new HiveDriverError(
'kernel backend: cannot supply both `token` and `oauthClientId`/`oauthClientSecret` ' +
'kernel backend: cannot supply both `token` and `oauthClientId`/`oauthClientSecret`/`oauthJwtKeyFile` ' +
"on the same connection. Pick one: 'access-token' (PAT) uses `token`; " +
"'databricks-oauth' uses the OAuth fields.",
);
Expand All @@ -637,6 +659,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,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Low — The JWT M2M branch validates oauthClientId and oauthJwtKid but treats tokenUrl as optional. Per this PR's own description, private_key_jwt cannot succeed against Databricks-native OIDC (it doesn't advertise the method), so tokenUrl (pointing at an external IdP like Entra) is effectively required for this flow to work at all. As written, omitting tokenUrl on the JWT branch passes all TS-side validation and then fails opaquely inside the kernel at connect time. Consider validating tokenUrl presence here with a clear, actionable error — matching the early-validation pattern already used for oauthClientId / oauthJwtKid a few lines above — so callers get the same crisp signal instead of a downstream kernel invalid_client.

oauthClientId: oauth.oauthClientId,
jwtKeyFile: oauth.oauthJwtKeyFile,
jwtKid: oauth.oauthJwtKid,
// Configurable (parity with pyo3); defaults to `['all-apis']` in the kernel.
oauthScopes:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Low — The JWT branch defaults oauthScopes to M2M_DEFAULT_SCOPES (['all-apis']), inherited verbatim from the shared-secret M2M path. But all-apis is a Databricks-native scope; the JWT private_key_jwt flow targets an external IdP (Entra), which rejects all-apis and expects <resource-id>/.default (as the usage example itself shows). Because the value is forwarded explicitly, the kernel's own default is never consulted — so a caller who omits oauthScopes on the JWT branch gets a scope that is essentially guaranteed to be wrong for the only IdP this flow supports. Consider either omitting oauthScopes when unset (letting the kernel default apply) or documenting that oauthScopes is de-facto required for the JWT flow.

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)
Expand Down Expand Up @@ -680,16 +751,17 @@ 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,
// Configurable (parity with pyo3); defaults to `['all-apis']`.
oauthScopes:
Array.isArray(oauth.oauthScopes) && oauth.oauthScopes.length > 0 ? oauth.oauthScopes : M2M_DEFAULT_SCOPES,
};
return oauth.tokenUrl !== undefined ? { ...m2m, tokenUrl: oauth.tokenUrl } : m2m;
Comment thread
rahuls-db marked this conversation as resolved.
}

throw new HiveDriverError(
Expand Down
107 changes: 107 additions & 0 deletions tests/unit/DBSQLClient.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -284,6 +284,83 @@ describe('DBSQLClient.connect', () => {
}
});

it('useKernel: true with an OAuth flow installs NO auth provider (kernel owns auth; no eager browser flow)', async () => {
const client = new DBSQLClient();

// `useKernel` + `databricks-oauth` (U2M: no secret, no token). The kernel
// owns the full auth lifecycle here, so `connect()` must NOT build the
// connector's own OAuth provider (which would eagerly open a browser /
// run the token exchange at connect() time via a telemetry client). The
// authProvider is assigned before the backend connects, so it is set even
// though the subsequent KernelBackend connect() rejects (absent native
// binding in CI / no live workspace).
const kernelOAuthOptions = {
...connectOptions,
token: undefined,
authType: 'databricks-oauth',
useKernel: true,
} as any;

try {
await client.connect(kernelOAuthOptions);
} catch (error) {
if (error instanceof AssertionError || !(error instanceof Error)) {
throw error;
}
// Expected: KernelBackend connect() rejects (native binding absent / no
// live workspace). The contract under test is the authProvider decision,
// which happened before the throw.
}

expect(client['authProvider']).to.be.undefined;
});

it('useKernel: true with a token installs a PAT-only PlainHttpAuthentication provider', async () => {
const client = new DBSQLClient();

// `useKernel` + a PAT: the connector hands the kernel a minimal PAT
// provider (for the telemetry / feature-flag clients) rather than
// undefined, and still must NOT build an OAuth provider.
const kernelPatOptions = { ...connectOptions, token: 'dapiXXXX', useKernel: true } as any;

try {
await client.connect(kernelPatOptions);
} catch (error) {
if (error instanceof AssertionError || !(error instanceof Error)) {
throw error;
}
// Expected: KernelBackend connect() rejects (native binding absent).
}

expect(client['authProvider']).to.be.instanceOf(PlainHttpAuthentication);
});

it('useKernel: true warns when a custom authProvider is supplied (it cannot be plumbed through)', async () => {
const client = new DBSQLClient();
const logSpy = sinon.spy((client as any).logger, 'log');

// The kernel owns auth via the native binding, so a JS-side authProvider
// is ignored — but the drop must be warned, not silent.
const kernelOptions = { ...connectOptions, token: 'dapiXXXX', useKernel: true } as any;

try {
await client.connect(kernelOptions, new AuthProviderStub());
} catch (error) {
if (error instanceof AssertionError || !(error instanceof Error)) {
throw error;
}
// Expected: KernelBackend connect() rejects (native binding absent). The
// warning is emitted before the backend connects.
}

const warned = logSpy
.getCalls()
.some((c) => c.args[0] === LogLevel.warn && /custom authProvider was supplied with useKernel/.test(c.args[1]));
expect(warned).to.be.true;

logSpy.restore();
});

it('populates config.customHeaders with org-id parsed from ?o= (SPOG)', async () => {
const client = new DBSQLClient();
await client.connect({ ...connectOptions, path: '/sql/1.0/warehouses/abc?o=12345678901234' });
Expand All @@ -297,6 +374,36 @@ describe('DBSQLClient.connect', () => {
});
});

describe('DBSQLClient.mapAuthType (telemetry authType)', () => {
it('labels databricks-oauth + oauthJwtKeyFile as oauth-m2m-jwt ONLY on the kernel path', () => {
const client = new DBSQLClient();

const kernelJwt = {
...connectOptions,
token: undefined,
authType: 'databricks-oauth',
oauthJwtKeyFile: '/keys/jwt.pem',
useKernel: true,
} as any;
expect(client['mapAuthType'](kernelJwt)).to.equal('oauth-m2m-jwt');
});

it('does NOT label a Thrift-path connection oauth-m2m-jwt even if oauthJwtKeyFile is set (no useKernel)', () => {
const client = new DBSQLClient();

// oauthJwtKeyFile is a kernel-only internal option; on the Thrift path a
// no-secret OAuth connection actually runs U2M (external-browser), so the
// label must reflect that rather than mislabeling it oauth-m2m-jwt.
const thriftJwt = {
...connectOptions,
token: undefined,
authType: 'databricks-oauth',
oauthJwtKeyFile: '/keys/jwt.pem',
} as any;
expect(client['mapAuthType'](thriftJwt)).to.equal('external-browser');
});
});

describe('DBSQLClient.openSession', () => {
it('should successfully open session', async () => {
const { client } = makeStubbedClient();
Expand Down
Loading
Loading