Skip to content

Commit 21d31dd

Browse files
committed
feat(sea): wire session-level query tags (statement-level already forwarded)
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>
1 parent 4804f1e commit 21d31dd

2 files changed

Lines changed: 45 additions & 0 deletions

File tree

lib/sea/SeaBackend.ts

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import { getSeaNative, SeaNativeBinding, SeaConnection } from './SeaNativeLoader
2323
import { decodeNapiKernelError } from './SeaErrorMapping';
2424
import { buildSeaConnectionOptions, buildSeaRetryOptions, SeaNativeConnectionOptions } from './SeaAuth';
2525
import { installKernelLogBridge } from './SeaLogging';
26+
import { serializeQueryTags } from '../utils';
2627
import SeaSessionBackend from './SeaSessionBackend';
2728

2829
export interface SeaBackendOptions {
@@ -145,6 +146,20 @@ export default class SeaBackend implements IBackend {
145146
if (request.configuration !== undefined) {
146147
sessionOptions.sessionConf = { ...request.configuration };
147148
}
149+
// Session-level query tags: serialize into the reserved `QUERY_TAGS`
150+
// session conf (the kernel allowlists `QUERY_TAGS` and forwards it onto
151+
// the SEA `CreateSession` `session_confs`), mirroring the Thrift
152+
// backend's `ThriftBackend.openSession`. Runs after the `configuration`
153+
// merge so `queryTags` takes precedence over an explicit
154+
// `configuration.QUERY_TAGS`, matching the documented contract.
155+
if (request.queryTags !== undefined) {
156+
const serialized = serializeQueryTags(request.queryTags);
157+
if (serialized) {
158+
sessionOptions.sessionConf = { ...(sessionOptions.sessionConf ?? {}), QUERY_TAGS: serialized };
159+
} else if (sessionOptions.sessionConf) {
160+
delete sessionOptions.sessionConf.QUERY_TAGS;
161+
}
162+
}
148163

149164
let nativeConnection: SeaConnection;
150165
try {

tests/unit/sea/execution.test.ts

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -602,6 +602,36 @@ describe('SeaBackend', () => {
602602
expect(connection.lastSql).to.equal('SELECT 1');
603603
});
604604

605+
it('openSession() serializes session-level queryTags into sessionConf.QUERY_TAGS', async () => {
606+
const connection = new FakeNativeConnection();
607+
const binding = makeBinding(connection);
608+
const backend = new SeaBackend({ context: makeContext(), nativeBinding: binding });
609+
await backend.connect({ host: 'h', path: '/p', token: 't' } as ConnectionOptions);
610+
611+
await backend.openSession({ queryTags: { team: 'eng', env: 'prod' } });
612+
613+
// Session-level tags land in the reserved QUERY_TAGS session conf (the
614+
// kernel allowlists it → SEA CreateSession session_confs), mirroring Thrift.
615+
const conf = (binding.openSessionStub.firstCall.args[0] as { sessionConf?: Record<string, string> }).sessionConf;
616+
expect(conf?.QUERY_TAGS).to.be.a('string');
617+
expect(conf?.QUERY_TAGS).to.contain('team:eng').and.to.contain('env:prod');
618+
});
619+
620+
it('openSession() queryTags takes precedence over an explicit configuration.QUERY_TAGS', async () => {
621+
const connection = new FakeNativeConnection();
622+
const binding = makeBinding(connection);
623+
const backend = new SeaBackend({ context: makeContext(), nativeBinding: binding });
624+
await backend.connect({ host: 'h', path: '/p', token: 't' } as ConnectionOptions);
625+
626+
await backend.openSession({
627+
configuration: { QUERY_TAGS: 'manual-raw-value' },
628+
queryTags: { team: 'eng' },
629+
});
630+
631+
const conf = (binding.openSessionStub.firstCall.args[0] as { sessionConf?: Record<string, string> }).sessionConf;
632+
expect(conf?.QUERY_TAGS).to.contain('team:eng').and.to.not.equal('manual-raw-value');
633+
});
634+
605635
it('close() clears connection state without throwing', async () => {
606636
const connection = new FakeNativeConnection();
607637
const binding = makeBinding(connection);

0 commit comments

Comments
 (0)