Support setting session level tags as dict

Author: jiabin-hu

examples/session_params.js

@@ -10,8 +10,12 @@ client
   .connect({ host, path, token })
   .then(async (client) => {
     const session = await client.openSession({
+      queryTags: {
+        team: 'engineering',
+        test: 'session-params',
+        driver: 'node',
+      },
       configuration: {
-        QUERY_TAGS: 'team:engineering,test:session-params,driver:node',
         ansi_mode: 'false',
       },
     });

lib/DBSQLClient.ts

@@ -16,7 +16,7 @@ import HttpConnection from './connection/connections/HttpConnection';
 import IConnectionOptions from './connection/contracts/IConnectionOptions';
 import Status from './dto/Status';
 import HiveDriverError from './errors/HiveDriverError';
-import { buildUserAgentString, definedOrError } from './utils';
+import { buildUserAgentString, definedOrError, serializeQueryTags } from './utils';
 import PlainHttpAuthentication from './connection/auth/PlainHttpAuthentication';
 import DatabricksOAuth, { OAuthFlow } from './connection/auth/DatabricksOAuth';
 import {
@@ -292,6 +292,16 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
       configuration['spark.sql.thriftserver.metadata.metricview.enabled'] = 'true';
     }
 
+    // Serialize queryTags dict and set in configuration; takes precedence over configuration.QUERY_TAGS
+    if (request.queryTags !== undefined) {
+      const serialized = serializeQueryTags(request.queryTags);
+      if (serialized) {
+        configuration['QUERY_TAGS'] = serialized;
+      } else {
+        delete configuration['QUERY_TAGS'];
+      }
+    }
+
     const response = await this.driver.openSession({
       client_protocol_i64: new Int64(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8),
       ...getInitialNamespaceOptions(request.initialCatalog, request.initialSchema),

lib/contracts/IDBSQLClient.ts

@@ -60,6 +60,12 @@ export interface OpenSessionRequest {
   initialCatalog?: string;
   initialSchema?: string;
   configuration?: { [key: string]: string };
+  /**
+   * Session-level query tags as key-value pairs. Serialized and passed via session configuration
+   * as "QUERY_TAGS". Values may be null/undefined to include a key without a value.
+   * If both queryTags and configuration.QUERY_TAGS are specified, queryTags takes precedence.
+   */
+  queryTags?: Record<string, string | null | undefined>;
 }
 
 export default interface IDBSQLClient {

tests/unit/DBSQLClient.test.ts

@@ -636,4 +636,49 @@ describe('DBSQLClient.enableMetricViewMetadata', () => {
       'spark.sql.thriftserver.metadata.metricview.enabled': 'true',
     });
   });
+
+  it('should serialize queryTags dict and set in session configuration', async () => {
+    const client = new DBSQLClient();
+    const thriftClient = new ThriftClientStub();
+    sinon.stub(client, 'getClient').returns(Promise.resolve(thriftClient));
+
+    await client.openSession({
+      queryTags: { team: 'data-eng', project: 'etl' },
+    });
+
+    expect(thriftClient.openSessionReq?.configuration).to.deep.equal({
+      QUERY_TAGS: 'team:data-eng,project:etl',
+    });
+  });
+
+  it('should let queryTags take precedence over configuration.QUERY_TAGS', async () => {
+    const client = new DBSQLClient();
+    const thriftClient = new ThriftClientStub();
+    sinon.stub(client, 'getClient').returns(Promise.resolve(thriftClient));
+
+    await client.openSession({
+      queryTags: { team: 'new-team' },
+      configuration: { QUERY_TAGS: 'team:old-team,other:value', ansi_mode: 'true' },
+    });
+
+    expect(thriftClient.openSessionReq?.configuration).to.deep.equal({
+      QUERY_TAGS: 'team:new-team',
+      ansi_mode: 'true',
+    });
+  });
+
+  it('should remove QUERY_TAGS from configuration when queryTags is empty', async () => {
+    const client = new DBSQLClient();
+    const thriftClient = new ThriftClientStub();
+    sinon.stub(client, 'getClient').returns(Promise.resolve(thriftClient));
+
+    await client.openSession({
+      queryTags: {},
+      configuration: { QUERY_TAGS: 'team:old-team', ansi_mode: 'true' },
+    });
+
+    expect(thriftClient.openSessionReq?.configuration).to.deep.equal({
+      ansi_mode: 'true',
+    });
+  });
 });