Add metric view metadata support

Author: shivam2680

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