diff --git a/lib/contracts/IDBSQLSession.ts b/lib/contracts/IDBSQLSession.ts index 392f3108..08ebad26 100644 --- a/lib/contracts/IDBSQLSession.ts +++ b/lib/contracts/IDBSQLSession.ts @@ -27,6 +27,22 @@ export type ExecuteStatementOptions = { * These tags apply only to this statement and do not persist across queries. */ queryTags?: Record; + /** + * Enable metric-view metadata expansion for this statement. When `true`, the + * Spark conf `spark.databricks.optimizer.enableMetricViewMetadata=true` is + * forwarded via the Thrift `confOverlay`. Applies to a single statement only; + * does not persist across queries. + * + * Setting `false` is equivalent to omitting the option — both leave + * `confOverlay` untouched and the server's session default applies. The + * option is opt-in only; there is no way to explicitly clear a server + * default from the client. + * + * Equivalent SEA wiring will route the same key through napi `statementConf` + * once the kernel statement-options surface lands — until then this knob is + * honored only on the Thrift backend. + */ + metricViewMetadata?: boolean; }; export type TypeInfoRequest = { diff --git a/lib/thrift-backend/ThriftSessionBackend.ts b/lib/thrift-backend/ThriftSessionBackend.ts index c103ab4f..4a53a057 100644 --- a/lib/thrift-backend/ThriftSessionBackend.ts +++ b/lib/thrift-backend/ThriftSessionBackend.ts @@ -188,6 +188,13 @@ export default class ThriftSessionBackend implements ISessionBackend { request.confOverlay = { ...request.confOverlay, query_tags: serializedQueryTags }; } + if (options.metricViewMetadata === true) { + request.confOverlay = { + ...request.confOverlay, + 'spark.databricks.optimizer.enableMetricViewMetadata': 'true', + }; + } + if (ProtocolVersion.supportsCloudFetch(this.serverProtocolVersion)) { request.canDownloadResult = options.useCloudFetch ?? clientConfig.useCloudFetch; } diff --git a/tests/unit/DBSQLSession.test.ts b/tests/unit/DBSQLSession.test.ts index 51b27133..ffc3f551 100644 --- a/tests/unit/DBSQLSession.test.ts +++ b/tests/unit/DBSQLSession.test.ts @@ -298,6 +298,64 @@ describe('DBSQLSession', () => { }); }); + describe('executeStatement with metricViewMetadata', () => { + const metricViewConfKey = 'spark.databricks.optimizer.enableMetricViewMetadata'; + + it('should forward the metric-view conf via confOverlay when metricViewMetadata is true', async () => { + const context = new ClientContextStub(); + const driver = sinon.spy(context.driver); + const session = createSessionForTest({ handle: sessionHandleStub, context }); + + await session.executeStatement('SELECT * FROM my_metric_view', { metricViewMetadata: true }); + + expect(driver.executeStatement.callCount).to.eq(1); + const req = driver.executeStatement.firstCall.args[0]; + expect(req.confOverlay).to.deep.include({ [metricViewConfKey]: 'true' }); + }); + + it('should not set the metric-view conf when metricViewMetadata is omitted', async () => { + const context = new ClientContextStub(); + const driver = sinon.spy(context.driver); + const session = createSessionForTest({ 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?.[metricViewConfKey]).to.be.undefined; + }); + + it('should not set the metric-view conf when metricViewMetadata is false', async () => { + const context = new ClientContextStub(); + const driver = sinon.spy(context.driver); + const session = createSessionForTest({ handle: sessionHandleStub, context }); + + await session.executeStatement('SELECT 1', { metricViewMetadata: false }); + + expect(driver.executeStatement.callCount).to.eq(1); + const req = driver.executeStatement.firstCall.args[0]; + expect(req.confOverlay?.[metricViewConfKey]).to.be.undefined; + }); + + it('should coexist with queryTags in the same confOverlay', async () => { + const context = new ClientContextStub(); + const driver = sinon.spy(context.driver); + const session = createSessionForTest({ handle: sessionHandleStub, context }); + + await session.executeStatement('SELECT 1', { + metricViewMetadata: true, + queryTags: { team: 'eng' }, + }); + + expect(driver.executeStatement.callCount).to.eq(1); + const req = driver.executeStatement.firstCall.args[0]; + expect(req.confOverlay).to.deep.include({ + [metricViewConfKey]: 'true', + query_tags: 'team:eng', + }); + }); + }); + describe('getTypeInfo', () => { it('should run operation', async () => { const session = createSessionForTest({ handle: sessionHandleStub, context: new ClientContextStub() });