From 21d31dd8d7bb7db09a3138b7bb089ecb6690d5fe Mon Sep 17 00:00:00 2001
From: Madhavendra Rathore <madhavendra.rathore@databricks.com>
Date: Mon, 8 Jun 2026 20:47:52 +0000
Subject: [PATCH] feat(sea): wire session-level query tags (statement-level
 already forwarded)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Brings SEA query-tag handling to parity with the Thrift backend across
both scopes:

- Session-level (`openSession({ queryTags })`): NEW — `SeaBackend.openSession`
  serializes `request.queryTags` into the reserved `QUERY_TAGS` session conf
  (the kernel allowlists `QUERY_TAGS` and forwards it onto the SEA
  `CreateSession` `session_confs`), mirroring `ThriftBackend.openSession`.
  Runs after the `configuration` merge so `queryTags` takes precedence over
  an explicit `configuration.QUERY_TAGS`, per the documented contract.
  Verified on a live warehouse that `QUERY_TAGS` reaches the `/sessions` wire.

- Statement-level (`executeStatement({ queryTags })`): already forwarded by
  `SeaSessionBackend` into `statementConf.query_tags`. Previously a no-op
  because the kernel dropped `statement_conf` before the SEA wire; the
  companion kernel PR (databricks-sql-kernel#150) adds the native
  `query_tags` array so it now reaches the server end-to-end. Functional
  once that lands + KERNEL_REV is bumped.

Tests: openSession serializes session-level queryTags into
sessionConf.QUERY_TAGS, and queryTags wins over an explicit
configuration.QUERY_TAGS. (Statement-level forwarding is already covered.)

Co-authored-by: Isaac
Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>
---
 lib/sea/SeaBackend.ts            | 15 +++++++++++++++
 tests/unit/sea/execution.test.ts | 30 ++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/lib/sea/SeaBackend.ts b/lib/sea/SeaBackend.ts
index 3ee00288..bcca8159 100644
--- a/lib/sea/SeaBackend.ts
+++ b/lib/sea/SeaBackend.ts
@@ -23,6 +23,7 @@ import { getSeaNative, SeaNativeBinding, SeaConnection } from './SeaNativeLoader
 import { decodeNapiKernelError } from './SeaErrorMapping';
 import { buildSeaConnectionOptions, buildSeaRetryOptions, SeaNativeConnectionOptions } from './SeaAuth';
 import { installKernelLogBridge } from './SeaLogging';
+import { serializeQueryTags } from '../utils';
 import SeaSessionBackend from './SeaSessionBackend';
 
 export interface SeaBackendOptions {
@@ -145,6 +146,20 @@ export default class SeaBackend 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` 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: SeaConnection;
     try {
diff --git a/tests/unit/sea/execution.test.ts b/tests/unit/sea/execution.test.ts
index 083af3a9..7aec2297 100644
--- a/tests/unit/sea/execution.test.ts
+++ b/tests/unit/sea/execution.test.ts
@@ -602,6 +602,36 @@ describe('SeaBackend', () => {
     expect(connection.lastSql).to.equal('SELECT 1');
   });
 
+  it('openSession() serializes session-level queryTags into sessionConf.QUERY_TAGS', async () => {
+    const connection = new FakeNativeConnection();
+    const binding = makeBinding(connection);
+    const backend = new SeaBackend({ 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 SeaBackend({ 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('close() clears connection state without throwing', async () => {
     const connection = new FakeNativeConnection();
     const binding = makeBinding(connection);