From baa55ac2fa7ad8107c16974c40f8080f2350494a Mon Sep 17 00:00:00 2001 From: Madhavendra Rathore Date: Tue, 9 Jun 2026 15:03:48 +0000 Subject: [PATCH 1/2] feat(kernel): wire session-level query tags into KernelBackend.openSession MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports the session-level query-tags wiring onto the post-#428 lib/kernel path (originally lib/sea/SeaBackend, before the SEA→kernel rename). openSession serializes request.queryTags into the reserved QUERY_TAGS session conf, which the kernel allowlists (SESSION_CONF_ALLOWLIST) and forwards onto the SEA CreateSession session_confs — mirroring ThriftBackend.openSession. queryTags takes precedence over an explicit configuration.QUERY_TAGS. Verified end-to-end against a live warehouse: the tag lands in system.query.history.query_tags. Co-authored-by: Isaac Signed-off-by: Madhavendra Rathore --- lib/kernel/KernelBackend.ts | 15 +++++++++++++++ tests/unit/kernel/execution.test.ts | 30 +++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/lib/kernel/KernelBackend.ts b/lib/kernel/KernelBackend.ts index 927771df..e0ee420f 100644 --- a/lib/kernel/KernelBackend.ts +++ b/lib/kernel/KernelBackend.ts @@ -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'; @@ -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 { diff --git a/tests/unit/kernel/execution.test.ts b/tests/unit/kernel/execution.test.ts index d94c6e2e..dd3521b9 100644 --- a/tests/unit/kernel/execution.test.ts +++ b/tests/unit/kernel/execution.test.ts @@ -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 }).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 }).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); From d6c0d815c4a4748def69250e2fae097c0d9e8ef5 Mon Sep 17 00:00:00 2001 From: Madhavendra Rathore Date: Tue, 9 Jun 2026 15:03:48 +0000 Subject: [PATCH 2/2] fix(kernel): request Thrift-parity OAuth scopes, configurable via oauthScopes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel U2M flow passed no scopes, so it fell through to the kernel's bare default (all-apis offline_access). The databricks-sql-connector OAuth app is registered for `sql`, so U2M auth used the wrong scope set. Pass scopes explicitly from the driver: - U2M defaults to ['sql', 'offline_access'] (matches the Thrift driver's defaultOAuthScopes), overriding the kernel's all-apis default. - M2M defaults to ['all-apis'] (matches Thrift + the kernel's M2M default). - Both overridable via a new `oauthScopes` connect option — closing the configurability gap with pyo3, which already forwards `scopes` on M2M. Driver-only change: the napi binding already forwards oauth_scopes and the kernel's u2m.rs/m2m.rs feed them into the authorize/token request. Co-authored-by: Isaac Signed-off-by: Madhavendra Rathore --- lib/contracts/IDBSQLClient.ts | 4 ++++ lib/kernel/KernelAuth.ts | 30 ++++++++++++++++++++++++++ tests/unit/kernel/auth-m2m.test.ts | 26 +++++++++++++++++++++++ tests/unit/kernel/auth-u2m.test.ts | 34 ++++++++++++++++++++++++++++++ 4 files changed, 94 insertions(+) diff --git a/lib/contracts/IDBSQLClient.ts b/lib/contracts/IDBSQLClient.ts index 80aa2763..de67070e 100644 --- a/lib/contracts/IDBSQLClient.ts +++ b/lib/contracts/IDBSQLClient.ts @@ -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; } | { authType: 'custom'; diff --git a/lib/kernel/KernelAuth.ts b/lib/kernel/KernelAuth.ts index 11b50a87..c4450742 100644 --- a/lib/kernel/KernelAuth.ts +++ b/lib/kernel/KernelAuth.ts @@ -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 @@ -189,12 +203,14 @@ export type KernelNativeConnectionOptions = KernelSessionDefaults & authMode: 'OAuthM2m'; oauthClientId: string; oauthClientSecret: string; + oauthScopes?: Array; } | { hostName: string; httpPath: string; authMode: 'OAuthU2m'; oauthRedirectPort: number; + oauthScopes?: Array; } ); @@ -541,6 +557,7 @@ export function buildKernelConnectionOptions(options: ConnectionOptions): Kernel const oauth = options as { oauthClientId?: string; oauthClientSecret?: string; + oauthScopes?: Array; azureTenantId?: string; useDatabricksOAuthInAzure?: boolean; persistence?: unknown; @@ -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, }; } @@ -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, }; } diff --git a/tests/unit/kernel/auth-m2m.test.ts b/tests/unit/kernel/auth-m2m.test.ts index 365e1e68..078008ab 100644 --- a/tests/unit/kernel/auth-m2m.test.ts +++ b/tests/unit/kernel/auth-m2m.test.ts @@ -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', @@ -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(); diff --git a/tests/unit/kernel/auth-u2m.test.ts b/tests/unit/kernel/auth-u2m.test.ts index 4cd7ef40..18e79ce3 100644 --- a/tests/unit/kernel/auth-u2m.test.ts +++ b/tests/unit/kernel/auth-u2m.test.ts @@ -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), @@ -143,6 +176,7 @@ describe('KernelAuth + KernelBackend — OAuth U2M auth flow', () => { intervalsAsString: true, authMode: 'OAuthU2m', oauthRedirectPort: 8030, + oauthScopes: ['sql', 'offline_access'], }); await session.close();