Add metric view metadata support

lib/DBSQLClient.ts

@@ -172,6 +172,11 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
       }
     }
 
+    // Store enableMetricViewMetadata configuration
+    if (options.enableMetricViewMetadata !== undefined) {
+      this.config.enableMetricViewMetadata = options.enableMetricViewMetadata;
+    }
+
     this.authProvider = this.createAuthProvider(options, authProvider);
 
     this.connectionProvider = this.createConnectionProvider(options);
@@ -218,10 +223,18 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
    * const session = await client.openSession();
    */
   public async openSession(request: OpenSessionRequest = {}): Promise<IDBSQLSession> {
+    // Prepare session configuration
+    const configuration = request.configuration ? { ...request.configuration } : {};
+
+    // Add metric view metadata config if enabled
+    if (this.config.enableMetricViewMetadata) {
+      configuration['spark.sql.thriftserver.metadata.metricview.enabled'] = 'true';
+    }
+
     const response = await this.driver.openSession({
       client_protocol_i64: new Int64(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8),
       ...getInitialNamespaceOptions(request.initialCatalog, request.initialSchema),
-      configuration: request.configuration,
+      configuration,
       canUseMultipleCatalogs: true,
     });
 

lib/contracts/IClientContext.ts

@@ -21,6 +21,7 @@ export interface ClientConfig {
   cloudFetchSpeedThresholdMBps: number;
 
   useLZ4Compression: boolean;
+  enableMetricViewMetadata?: boolean;
 }
 
 export default interface IClientContext {

lib/contracts/IDBSQLClient.ts

@@ -33,6 +33,7 @@ export type ConnectionOptions = {
   userAgentEntry?: string;
   socketTimeout?: number;
   proxy?: ProxyOptions;
+  enableMetricViewMetadata?: boolean;
 } & AuthOptions;
 
 export interface OpenSessionRequest {

tests/unit/DBSQLClient.test.ts

@@ -560,3 +560,80 @@ describe('DBSQLClient.createAuthProvider', () => {
     expect(provider).to.be.equal(customProvider);
   });
 });
+
+describe('DBSQLClient.enableMetricViewMetadata', () => {
+  it('should store enableMetricViewMetadata config when enabled', async () => {
+    const client = new DBSQLClient();
+
+    expect(client.getConfig().enableMetricViewMetadata).to.be.undefined;
+
+    await client.connect({ ...connectOptions, enableMetricViewMetadata: true });
+
+    expect(client.getConfig().enableMetricViewMetadata).to.be.true;
+  });
+
+  it('should not store enableMetricViewMetadata config when disabled', async () => {
+    const client = new DBSQLClient();
+
+    expect(client.getConfig().enableMetricViewMetadata).to.be.undefined;
+
+    await client.connect({ ...connectOptions, enableMetricViewMetadata: false });
+
+    expect(client.getConfig().enableMetricViewMetadata).to.be.false;
+  });
+
+  it('should inject session parameter when enableMetricViewMetadata is true', async () => {
+    const client = new DBSQLClient();
+    const thriftClient = new ThriftClientStub();
+    sinon.stub(client, 'getClient').returns(Promise.resolve(thriftClient));
+
+    await client.connect({ ...connectOptions, enableMetricViewMetadata: true });
+    await client.openSession();
+
+    expect(thriftClient.openSessionReq?.configuration).to.have.property(
+      'spark.sql.thriftserver.metadata.metricview.enabled',
+      'true',
+    );
+  });
+
+  it('should not inject session parameter when enableMetricViewMetadata is false', async () => {
+    const client = new DBSQLClient();
+    const thriftClient = new ThriftClientStub();
+    sinon.stub(client, 'getClient').returns(Promise.resolve(thriftClient));
+
+    await client.connect({ ...connectOptions, enableMetricViewMetadata: false });
+    await client.openSession();
+
+    expect(thriftClient.openSessionReq?.configuration).to.not.have.property(
+      'spark.sql.thriftserver.metadata.metricview.enabled',
+    );
+  });
+
+  it('should not inject session parameter when enableMetricViewMetadata is not set', async () => {
+    const client = new DBSQLClient();
+    const thriftClient = new ThriftClientStub();
+    sinon.stub(client, 'getClient').returns(Promise.resolve(thriftClient));
+
+    await client.connect(connectOptions);
+    await client.openSession();
+
+    expect(thriftClient.openSessionReq?.configuration).to.not.have.property(
+      'spark.sql.thriftserver.metadata.metricview.enabled',
+    );
+  });
+
+  it('should preserve user-provided session configuration', async () => {
+    const client = new DBSQLClient();
+    const thriftClient = new ThriftClientStub();
+    sinon.stub(client, 'getClient').returns(Promise.resolve(thriftClient));
+
+    await client.connect({ ...connectOptions, enableMetricViewMetadata: true });
+    const userConfig = { QUERY_TAGS: 'team:engineering', ansi_mode: 'true' };
+    await client.openSession({ configuration: userConfig });
+
+    expect(thriftClient.openSessionReq?.configuration).to.deep.equal({
+      ...userConfig,
+      'spark.sql.thriftserver.metadata.metricview.enabled': 'true',
+    });
+  });
+});