From 47b625aeaa7a1930811c24093d8803a478c4de7d Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Thu, 16 Oct 2025 00:55:09 +0530 Subject: [PATCH] Add metric view metadata support Signed-off-by: Shivam Raj --- lib/DBSQLClient.ts | 15 ++++++- lib/contracts/IClientContext.ts | 1 + lib/contracts/IDBSQLClient.ts | 1 + tests/unit/DBSQLClient.test.ts | 77 +++++++++++++++++++++++++++++++++ 4 files changed, 93 insertions(+), 1 deletion(-) diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index 79cd936d..00496463 100644 --- a/lib/DBSQLClient.ts +++ b/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 { + // 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, }); diff --git a/lib/contracts/IClientContext.ts b/lib/contracts/IClientContext.ts index cc13d326..9f0998fc 100644 --- a/lib/contracts/IClientContext.ts +++ b/lib/contracts/IClientContext.ts @@ -21,6 +21,7 @@ export interface ClientConfig { cloudFetchSpeedThresholdMBps: number; useLZ4Compression: boolean; + enableMetricViewMetadata?: boolean; } export default interface IClientContext { diff --git a/lib/contracts/IDBSQLClient.ts b/lib/contracts/IDBSQLClient.ts index e18b45fb..26588031 100644 --- a/lib/contracts/IDBSQLClient.ts +++ b/lib/contracts/IDBSQLClient.ts @@ -33,6 +33,7 @@ export type ConnectionOptions = { userAgentEntry?: string; socketTimeout?: number; proxy?: ProxyOptions; + enableMetricViewMetadata?: boolean; } & AuthOptions; export interface OpenSessionRequest { diff --git a/tests/unit/DBSQLClient.test.ts b/tests/unit/DBSQLClient.test.ts index b81dbae1..931615f7 100644 --- a/tests/unit/DBSQLClient.test.ts +++ b/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', + }); + }); +});