add statement level query tags support

Author: jiabin-hu

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

lib/contracts/IDBSQLSession.ts

@@ -21,6 +21,12 @@ export type ExecuteStatementOptions = {
   stagingAllowedLocalPath?: string | string[];
   namedParameters?: Record<string, DBSQLParameter | DBSQLParameterValue>;
   ordinalParameters?: Array<DBSQLParameter | DBSQLParameterValue>;
+  /**
+   * 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<string, string | null | undefined>;
 };
 
 export type TypeInfoRequest = {

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,
+};

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<string, string | null | undefined> | 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(',');
+}

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

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');
+  });
+});