From 6f1118f2bdeb82841b1a19b1ad622708c773f325 Mon Sep 17 00:00:00 2001 From: Jiabin Hu Date: Wed, 25 Mar 2026 13:56:46 -0700 Subject: [PATCH 1/2] add statement level query tags support --- lib/DBSQLSession.ts | 7 +++- lib/contracts/IDBSQLSession.ts | 6 ++++ lib/utils/index.ts | 11 +++++- lib/utils/queryTags.ts | 34 ++++++++++++++++++ tests/unit/DBSQLSession.test.ts | 38 ++++++++++++++++++++ tests/unit/utils/queryTags.test.ts | 58 ++++++++++++++++++++++++++++++ 6 files changed, 152 insertions(+), 2 deletions(-) create mode 100644 lib/utils/queryTags.ts create mode 100644 tests/unit/utils/queryTags.test.ts diff --git a/lib/DBSQLSession.ts b/lib/DBSQLSession.ts index 9b4245c3..95715e1b 100644 --- a/lib/DBSQLSession.ts +++ b/lib/DBSQLSession.ts @@ -31,7 +31,7 @@ import IOperation from './contracts/IOperation'; import DBSQLOperation from './DBSQLOperation'; import Status from './dto/Status'; import InfoValue from './dto/InfoValue'; -import { definedOrError, LZ4, ProtocolVersion } from './utils'; +import { definedOrError, LZ4, ProtocolVersion, serializeQueryTags } from './utils'; import CloseableCollection from './utils/CloseableCollection'; import { LogLevel } from './contracts/IDBSQLLogger'; import HiveDriverError from './errors/HiveDriverError'; @@ -227,6 +227,11 @@ export default class DBSQLSession implements IDBSQLSession { request.parameters = getQueryParameters(options.namedParameters, options.ordinalParameters); } + const serializedQueryTags = serializeQueryTags(options.queryTags); + if (serializedQueryTags !== undefined) { + request.confOverlay = { ...request.confOverlay, query_tags: serializedQueryTags }; + } + if (ProtocolVersion.supportsCloudFetch(this.serverProtocolVersion)) { request.canDownloadResult = options.useCloudFetch ?? clientConfig.useCloudFetch; } diff --git a/lib/contracts/IDBSQLSession.ts b/lib/contracts/IDBSQLSession.ts index 0f751714..392f3108 100644 --- a/lib/contracts/IDBSQLSession.ts +++ b/lib/contracts/IDBSQLSession.ts @@ -21,6 +21,12 @@ export type ExecuteStatementOptions = { stagingAllowedLocalPath?: string | string[]; namedParameters?: Record; ordinalParameters?: Array; + /** + * Per-statement query tags as key-value pairs. Serialized and passed via confOverlay + * as "query_tags". Values may be null/undefined to include a key without a value. + * These tags apply only to this statement and do not persist across queries. + */ + queryTags?: Record; }; export type TypeInfoRequest = { diff --git a/lib/utils/index.ts b/lib/utils/index.ts index b8203c4d..00fa2363 100644 --- a/lib/utils/index.ts +++ b/lib/utils/index.ts @@ -3,5 +3,14 @@ import buildUserAgentString from './buildUserAgentString'; import formatProgress, { ProgressUpdateTransformer } from './formatProgress'; import LZ4 from './lz4'; import * as ProtocolVersion from './protocolVersion'; +import { serializeQueryTags } from './queryTags'; -export { definedOrError, buildUserAgentString, formatProgress, ProgressUpdateTransformer, LZ4, ProtocolVersion }; +export { + definedOrError, + buildUserAgentString, + formatProgress, + ProgressUpdateTransformer, + LZ4, + ProtocolVersion, + serializeQueryTags, +}; diff --git a/lib/utils/queryTags.ts b/lib/utils/queryTags.ts new file mode 100644 index 00000000..d788e75b --- /dev/null +++ b/lib/utils/queryTags.ts @@ -0,0 +1,34 @@ +/** + * Serializes a query tags dictionary into a string for use in confOverlay. + * + * Format: comma-separated key:value pairs, e.g. "key1:value1,key2:value2" + * - If a value is null or undefined, the key is included without a colon or value + * - Special characters (backslash, colon, comma) in values are backslash-escaped + * - Keys are not escaped + * + * @param queryTags - dictionary of query tag key-value pairs + * @returns serialized string, or undefined if input is empty/null/undefined + */ +export function serializeQueryTags( + queryTags: Record | null | undefined, +): string | undefined { + if (queryTags == null) { + return undefined; + } + + const keys = Object.keys(queryTags); + if (keys.length === 0) { + return undefined; + } + + return keys + .map((key) => { + const value = queryTags[key]; + if (value == null) { + return key; + } + const escapedValue = value.replace(/[\\:,]/g, (c) => `\\${c}`); + return `${key}:${escapedValue}`; + }) + .join(','); +} diff --git a/tests/unit/DBSQLSession.test.ts b/tests/unit/DBSQLSession.test.ts index ddf843dc..0dc79037 100644 --- a/tests/unit/DBSQLSession.test.ts +++ b/tests/unit/DBSQLSession.test.ts @@ -259,6 +259,44 @@ describe('DBSQLSession', () => { }); }); + describe('executeStatement with queryTags', () => { + it('should set confOverlay with query_tags when queryTags are provided', async () => { + const context = new ClientContextStub(); + const driver = sinon.spy(context.driver); + const session = new DBSQLSession({ handle: sessionHandleStub, context }); + + await session.executeStatement('SELECT 1', { queryTags: { team: 'eng', app: 'etl' } }); + + expect(driver.executeStatement.callCount).to.eq(1); + const req = driver.executeStatement.firstCall.args[0]; + expect(req.confOverlay).to.deep.include({ query_tags: 'team:eng,app:etl' }); + }); + + it('should not set confOverlay query_tags when queryTags is not provided', async () => { + const context = new ClientContextStub(); + const driver = sinon.spy(context.driver); + const session = new DBSQLSession({ handle: sessionHandleStub, context }); + + await session.executeStatement('SELECT 1'); + + expect(driver.executeStatement.callCount).to.eq(1); + const req = driver.executeStatement.firstCall.args[0]; + expect(req.confOverlay?.query_tags).to.be.undefined; + }); + + it('should not set confOverlay query_tags when queryTags is empty', async () => { + const context = new ClientContextStub(); + const driver = sinon.spy(context.driver); + const session = new DBSQLSession({ handle: sessionHandleStub, context }); + + await session.executeStatement('SELECT 1', { queryTags: {} }); + + expect(driver.executeStatement.callCount).to.eq(1); + const req = driver.executeStatement.firstCall.args[0]; + expect(req.confOverlay?.query_tags).to.be.undefined; + }); + }); + describe('getTypeInfo', () => { it('should run operation', async () => { const session = new DBSQLSession({ handle: sessionHandleStub, context: new ClientContextStub() }); diff --git a/tests/unit/utils/queryTags.test.ts b/tests/unit/utils/queryTags.test.ts new file mode 100644 index 00000000..3904b098 --- /dev/null +++ b/tests/unit/utils/queryTags.test.ts @@ -0,0 +1,58 @@ +import { expect } from 'chai'; +import { serializeQueryTags } from '../../../lib/utils/queryTags'; + +describe('serializeQueryTags', () => { + it('should return undefined for null input', () => { + expect(serializeQueryTags(null)).to.be.undefined; + }); + + it('should return undefined for undefined input', () => { + expect(serializeQueryTags(undefined)).to.be.undefined; + }); + + it('should return undefined for empty object', () => { + expect(serializeQueryTags({})).to.be.undefined; + }); + + it('should serialize a single tag', () => { + expect(serializeQueryTags({ team: 'engineering' })).to.equal('team:engineering'); + }); + + it('should serialize multiple tags', () => { + const result = serializeQueryTags({ team: 'engineering', app: 'etl' }); + expect(result).to.equal('team:engineering,app:etl'); + }); + + it('should omit colon for null value', () => { + expect(serializeQueryTags({ team: null })).to.equal('team'); + }); + + it('should omit colon for undefined value', () => { + expect(serializeQueryTags({ team: undefined })).to.equal('team'); + }); + + it('should mix null and non-null values', () => { + const result = serializeQueryTags({ team: 'eng', flag: null, app: 'etl' }); + expect(result).to.equal('team:eng,flag,app:etl'); + }); + + it('should escape backslash in value', () => { + expect(serializeQueryTags({ path: 'a\\b' })).to.equal('path:a\\\\b'); + }); + + it('should escape colon in value', () => { + expect(serializeQueryTags({ url: 'http://host' })).to.equal('url:http\\://host'); + }); + + it('should escape comma in value', () => { + expect(serializeQueryTags({ list: 'a,b' })).to.equal('list:a\\,b'); + }); + + it('should escape multiple special characters in value', () => { + expect(serializeQueryTags({ val: 'a\\b:c,d' })).to.equal('val:a\\\\b\\:c\\,d'); + }); + + it('should not escape special characters in keys', () => { + expect(serializeQueryTags({ 'key:name': 'value' })).to.equal('key:name:value'); + }); +}); From 740fe7c199e067d71ff489fdff80d5bd86a83d3f Mon Sep 17 00:00:00 2001 From: Jiabin Hu Date: Wed, 25 Mar 2026 14:07:25 -0700 Subject: [PATCH 2/2] Escape backslashes in query tag keys Co-authored-by: Isaac --- lib/utils/queryTags.ts | 7 ++++--- tests/unit/utils/queryTags.test.ts | 10 +++++++++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/lib/utils/queryTags.ts b/lib/utils/queryTags.ts index d788e75b..eb4a92db 100644 --- a/lib/utils/queryTags.ts +++ b/lib/utils/queryTags.ts @@ -3,8 +3,8 @@ * * Format: comma-separated key:value pairs, e.g. "key1:value1,key2:value2" * - If a value is null or undefined, the key is included without a colon or value + * - Backslashes in keys are escaped; other special characters in keys are not escaped * - Special characters (backslash, colon, comma) in values are backslash-escaped - * - Keys are not escaped * * @param queryTags - dictionary of query tag key-value pairs * @returns serialized string, or undefined if input is empty/null/undefined @@ -23,12 +23,13 @@ export function serializeQueryTags( return keys .map((key) => { + const escapedKey = key.replace(/\\/g, '\\\\'); const value = queryTags[key]; if (value == null) { - return key; + return escapedKey; } const escapedValue = value.replace(/[\\:,]/g, (c) => `\\${c}`); - return `${key}:${escapedValue}`; + return `${escapedKey}:${escapedValue}`; }) .join(','); } diff --git a/tests/unit/utils/queryTags.test.ts b/tests/unit/utils/queryTags.test.ts index 3904b098..d804894a 100644 --- a/tests/unit/utils/queryTags.test.ts +++ b/tests/unit/utils/queryTags.test.ts @@ -52,7 +52,15 @@ describe('serializeQueryTags', () => { expect(serializeQueryTags({ val: 'a\\b:c,d' })).to.equal('val:a\\\\b\\:c\\,d'); }); - it('should not escape special characters in keys', () => { + it('should escape backslash in key', () => { + expect(serializeQueryTags({ 'a\\b': 'value' })).to.equal('a\\\\b:value'); + }); + + it('should escape backslash in key with null value', () => { + expect(serializeQueryTags({ 'a\\b': null })).to.equal('a\\\\b'); + }); + + it('should not escape other special characters in keys', () => { expect(serializeQueryTags({ 'key:name': 'value' })).to.equal('key:name:value'); }); });