feat(kernel): session-level query tags + Thrift-parity OAuth scopes (configurable)

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<string>;
     }
   | {
       authType: 'custom';

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<string>;
       }
     | {
         hostName: string;
         httpPath: string;
         authMode: 'OAuthU2m';
         oauthRedirectPort: number;
+        oauthScopes?: Array<string>;
       }
   );
 
@@ -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;
@@ -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,
     };
   }
 

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 {

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();

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();

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