Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 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 lib/contracts/IDBSQLClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,10 @@ type AuthOptions =
oauthClientId?: string;
oauthClientSecret?: string;
useDatabricksOAuthInAzure?: boolean;
// OAuth scopes to request. When omitted, the kernel backend defaults the
// U2M flow to `['sql', 'offline_access']` (parity with the Thrift driver's
// `defaultOAuthScopes`), overriding the kernel's bare `all-apis offline_access`.
oauthScopes?: Array<string>;
}
| {
authType: 'custom';
Expand Down
30 changes: 30 additions & 0 deletions lib/kernel/KernelAuth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,20 @@ import { buildUserAgentString } from '../utils';
*/
const U2M_DEFAULT_REDIRECT_PORT = 8030;

// U2M OAuth scopes default. Matches the standalone Thrift driver's
// `defaultOAuthScopes` (lib/connection/auth/DatabricksOAuth/OAuthScope.ts):
// `['sql', 'offline_access']`. The kernel's bare default is
// `['all-apis', 'offline_access']`; the `databricks-sql-connector` OAuth app is
// registered for the `sql` scope, so we pass the Thrift-parity scopes explicitly
// unless the caller overrides via `oauthScopes`.
const U2M_DEFAULT_SCOPES = ['sql', 'offline_access'];

// M2M OAuth scopes default. Matches the standalone Thrift driver (`getScopes`
// forces `['all-apis']` for the client-credentials flow) and the kernel's own
// M2M default (`m2m.rs` → `['all-apis']`). Overridable via `oauthScopes`
// (parity with pyo3, which forwards `scopes` on M2M).
const M2M_DEFAULT_SCOPES = ['all-apis'];

/**
* Shape consumed by the napi-binding's `openSession()` (see
* `native/kernel/index.d.ts`). Mirrors `ConnectionOptions` in the binding's
Expand Down Expand Up @@ -189,12 +203,14 @@ export type KernelNativeConnectionOptions = KernelSessionDefaults &
authMode: 'OAuthM2m';
oauthClientId: string;
oauthClientSecret: string;
oauthScopes?: Array<string>;
}
| {
hostName: string;
httpPath: string;
authMode: 'OAuthU2m';
oauthRedirectPort: number;
oauthScopes?: Array<string>;
}
);

Expand Down Expand Up @@ -541,6 +557,7 @@ export function buildKernelConnectionOptions(options: ConnectionOptions): Kernel
const oauth = options as {
oauthClientId?: string;
oauthClientSecret?: string;
oauthScopes?: Array<string>;
azureTenantId?: string;
useDatabricksOAuthInAzure?: boolean;
persistence?: unknown;
Expand Down Expand Up @@ -627,6 +644,13 @@ export function buildKernelConnectionOptions(options: ConnectionOptions): Kernel
...base,
authMode: 'OAuthU2m',
oauthRedirectPort: U2M_DEFAULT_REDIRECT_PORT,
// Pass scopes explicitly so the kernel requests the same set as the
// Thrift driver (`sql offline_access`) rather than its bare-Rust
// `all-apis offline_access` default. Caller can override via `oauthScopes`.
oauthScopes:
Array.isArray(oauth.oauthScopes) && oauth.oauthScopes.length > 0
? oauth.oauthScopes
: U2M_DEFAULT_SCOPES,
};
}

Expand All @@ -652,6 +676,12 @@ export function buildKernelConnectionOptions(options: ConnectionOptions): Kernel
authMode: 'OAuthM2m',
oauthClientId: oauth.oauthClientId,
oauthClientSecret: oauth.oauthClientSecret,
// Configurable (parity with pyo3); defaults to `['all-apis']` — the only
// scope the client-credentials flow allows, matching Thrift + the kernel.
oauthScopes:
Array.isArray(oauth.oauthScopes) && oauth.oauthScopes.length > 0
? oauth.oauthScopes
: M2M_DEFAULT_SCOPES,
};
}

Expand Down
15 changes: 15 additions & 0 deletions lib/kernel/KernelBackend.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import { ConnectionOptions, OpenSessionRequest } from '../contracts/IDBSQLClient
import { InternalConnectionOptions } from '../contracts/InternalConnectionOptions';
import { LogLevel } from '../contracts/IDBSQLLogger';
import HiveDriverError from '../errors/HiveDriverError';
import { serializeQueryTags } from '../utils';
import { getKernelNative, KernelNativeBinding, KernelConnection } from './KernelNativeLoader';
import { decodeNapiKernelError } from './KernelErrorMapping';
import { buildKernelConnectionOptions, buildKernelRetryOptions, KernelNativeConnectionOptions } from './KernelAuth';
Expand Down Expand Up @@ -145,6 +146,20 @@ export default class KernelBackend implements IBackend {
if (request.configuration !== undefined) {
sessionOptions.sessionConf = { ...request.configuration };
}
// Session-level query tags: serialize into the reserved `QUERY_TAGS`
// session conf. The kernel allowlists `QUERY_TAGS` (SESSION_CONF_ALLOWLIST)
// and forwards it onto the SEA `CreateSession` `session_confs`, mirroring
// the Thrift backend's `ThriftBackend.openSession`. Runs after the
// `configuration` merge so `queryTags` takes precedence over an explicit
// `configuration.QUERY_TAGS`, matching the documented contract.
if (request.queryTags !== undefined) {
const serialized = serializeQueryTags(request.queryTags);
if (serialized) {
sessionOptions.sessionConf = { ...(sessionOptions.sessionConf ?? {}), QUERY_TAGS: serialized };
} else if (sessionOptions.sessionConf) {
delete sessionOptions.sessionConf.QUERY_TAGS;
}
}

let nativeConnection: KernelConnection;
try {
Expand Down
26 changes: 26 additions & 0 deletions tests/unit/kernel/auth-m2m.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,34 @@ describe('KernelAuth + KernelBackend — OAuth M2M auth flow', () => {
authMode: 'OAuthM2m',
oauthClientId: 'client-uuid',
oauthClientSecret: 'dose-fake-secret',
oauthScopes: ['all-apis'],
});
});

it('defaults M2M oauthScopes to all-apis (Thrift + kernel parity)', () => {
const native = buildKernelConnectionOptions({
host: 'example.cloud.databricks.com',
path: '/sql/1.0/warehouses/abc',
authType: 'databricks-oauth',
oauthClientId: 'client-uuid',
oauthClientSecret: 'dose-fake-secret',
});
expect(native.authMode).to.equal('OAuthM2m');
expect((native as { oauthScopes?: string[] }).oauthScopes).to.deep.equal(['all-apis']);
});

it('honors a caller-supplied M2M oauthScopes override (parity with pyo3)', () => {
const native = buildKernelConnectionOptions({
host: 'example.cloud.databricks.com',
path: '/sql/1.0/warehouses/abc',
authType: 'databricks-oauth',
oauthClientId: 'client-uuid',
oauthClientSecret: 'dose-fake-secret',
oauthScopes: ['sql', 'offline_access'],
} as ConnectionOptions);
expect((native as { oauthScopes?: string[] }).oauthScopes).to.deep.equal(['sql', 'offline_access']);
});

it('prepends `/` to the path on the M2M branch too', () => {
const opts: ConnectionOptions = {
host: 'example.cloud.databricks.com',
Expand Down Expand Up @@ -177,6 +202,7 @@ describe('KernelAuth + KernelBackend — OAuth M2M auth flow', () => {
authMode: 'OAuthM2m',
oauthClientId: 'client-uuid',
oauthClientSecret: 'dose-fake-secret',
oauthScopes: ['all-apis'],
});

await session.close();
Expand Down
34 changes: 34 additions & 0 deletions tests/unit/kernel/auth-u2m.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,42 @@ describe('KernelAuth + KernelBackend — OAuth U2M auth flow', () => {
intervalsAsString: true,
authMode: 'OAuthU2m',
oauthRedirectPort: 8030,
oauthScopes: ['sql', 'offline_access'],
});
});

it('defaults U2M oauthScopes to Thrift parity (sql offline_access)', () => {
const native = buildKernelConnectionOptions({
host: 'example.cloud.databricks.com',
path: '/sql/1.0/warehouses/abc',
authType: 'databricks-oauth',
});
expect(native.authMode).to.equal('OAuthU2m');
// Matches the standalone Thrift driver's defaultOAuthScopes, NOT the
// kernel's bare `all-apis offline_access` default.
expect((native as { oauthScopes?: string[] }).oauthScopes).to.deep.equal(['sql', 'offline_access']);
});

it('honors a caller-supplied U2M oauthScopes override', () => {
const native = buildKernelConnectionOptions({
host: 'example.cloud.databricks.com',
path: '/sql/1.0/warehouses/abc',
authType: 'databricks-oauth',
oauthScopes: ['all-apis'],
} as ConnectionOptions);
expect((native as { oauthScopes?: string[] }).oauthScopes).to.deep.equal(['all-apis']);
});

it('falls back to the default U2M scopes when oauthScopes is an empty array', () => {
const native = buildKernelConnectionOptions({
host: 'example.cloud.databricks.com',
path: '/sql/1.0/warehouses/abc',
authType: 'databricks-oauth',
oauthScopes: [],
} as ConnectionOptions);
expect((native as { oauthScopes?: string[] }).oauthScopes).to.deep.equal(['sql', 'offline_access']);
});

it('rejects oauthClientId without oauthClientSecret as M2M-with-missing-secret', () => {
// Round-4 NF3-2: presence of `oauthClientId` signals M2M intent.
// Routing now keys off the id (the "do I have an id?" signal),
Expand Down Expand Up @@ -143,6 +176,7 @@ describe('KernelAuth + KernelBackend — OAuth U2M auth flow', () => {
intervalsAsString: true,
authMode: 'OAuthU2m',
oauthRedirectPort: 8030,
oauthScopes: ['sql', 'offline_access'],
});

await session.close();
Expand Down
30 changes: 30 additions & 0 deletions tests/unit/kernel/execution.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -551,6 +551,36 @@ describe('KernelBackend', () => {
});
});

it('openSession() serializes session-level queryTags into sessionConf.QUERY_TAGS', async () => {
const connection = new FakeNativeConnection();
const binding = makeBinding(connection);
const backend = new KernelBackend({ context: makeContext(), nativeBinding: binding });
await backend.connect({ host: 'h', path: '/p', token: 't' } as ConnectionOptions);

await backend.openSession({ queryTags: { team: 'eng', env: 'prod' } });

// Session-level tags land in the reserved QUERY_TAGS session conf (the
// kernel allowlists it → SEA CreateSession session_confs), mirroring Thrift.
const conf = (binding.openSessionStub.firstCall.args[0] as { sessionConf?: Record<string, string> }).sessionConf;
expect(conf?.QUERY_TAGS).to.be.a('string');
expect(conf?.QUERY_TAGS).to.contain('team:eng').and.to.contain('env:prod');
});

it('openSession() queryTags takes precedence over an explicit configuration.QUERY_TAGS', async () => {
const connection = new FakeNativeConnection();
const binding = makeBinding(connection);
const backend = new KernelBackend({ context: makeContext(), nativeBinding: binding });
await backend.connect({ host: 'h', path: '/p', token: 't' } as ConnectionOptions);

await backend.openSession({
configuration: { QUERY_TAGS: 'manual-raw-value' },
queryTags: { team: 'eng' },
});

const conf = (binding.openSessionStub.firstCall.args[0] as { sessionConf?: Record<string, string> }).sessionConf;
expect(conf?.QUERY_TAGS).to.contain('team:eng').and.to.not.equal('manual-raw-value');
});

it('openSession() returns a KernelSessionBackend wrapping the napi Connection', async () => {
const connection = new FakeNativeConnection();
const binding = makeBinding(connection);
Expand Down
Loading