Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions lib/DBSQLSession.ts
Original file line number Diff line number Diff line change
Expand Up @@ -236,6 +236,13 @@ export default class DBSQLSession implements IDBSQLSession {
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;
}
Expand Down
16 changes: 16 additions & 0 deletions lib/contracts/IDBSQLSession.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,22 @@ export type ExecuteStatementOptions = {
* These tags apply only to this statement and do not persist across queries.
*/
queryTags?: Record<string, string | null | undefined>;
/**
* 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 = {
Expand Down
58 changes: 58 additions & 0 deletions tests/unit/DBSQLSession.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -297,6 +297,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 = new DBSQLSession({ 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 = new DBSQLSession({ 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 = new DBSQLSession({ 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 = new DBSQLSession({ 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 = new DBSQLSession({ handle: sessionHandleStub, context: new ClientContextStub() });
Expand Down
Loading