From 8d8b70488c6ed29ac49396ee1a883aeb3e50a596 Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Fri, 21 Aug 2026 21:36:36 +0000 Subject: [PATCH] Forward kernel telemetry options --- lib/DBSQLClient.ts | 5 +- lib/kernel/KernelAuth.ts | 80 +++++++++++++++++++++ lib/kernel/KernelBackend.ts | 8 ++- native/kernel/index.d.ts | 27 +++++++ tests/unit/DBSQLClient.test.ts | 17 +++++ tests/unit/kernel/_helpers/nativeOptions.ts | 18 +++++ tests/unit/kernel/execution.test.ts | 77 +++++++++++++++++++- 7 files changed, 227 insertions(+), 5 deletions(-) diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index f021edf9..af813e6a 100644 --- a/lib/DBSQLClient.ts +++ b/lib/DBSQLClient.ts @@ -729,7 +729,8 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I // doesn't ship in the public `.d.ts`. Mirrors Python's `kwargs.get("use_kernel")` // pattern (see databricks-sql-python/src/databricks/sql/session.py). const internalOptions = options as ConnectionOptions & InternalConnectionOptions; - const backend = internalOptions.useKernel + const useKernel = internalOptions.useKernel === true; + const backend = useKernel ? new KernelBackend({ context: this }) : new ThriftBackend({ context: this, @@ -777,7 +778,7 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I `Telemetry remains controlled by the runtime config and feature flag.`, ); } - if (this.config.telemetryEnabled && !envDisabled) { + if (!useKernel && this.config.telemetryEnabled && !envDisabled) { await this.initializeTelemetry(); } diff --git a/lib/kernel/KernelAuth.ts b/lib/kernel/KernelAuth.ts index 7cf99afa..8dbbcd4c 100644 --- a/lib/kernel/KernelAuth.ts +++ b/lib/kernel/KernelAuth.ts @@ -12,11 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. +import os from 'os'; import { ConnectionOptions } from '../contracts/IDBSQLClient'; +import { ClientConfig } from '../contracts/IClientContext'; import { InternalConnectionOptions } from '../contracts/InternalConnectionOptions'; import AuthenticationError from '../errors/AuthenticationError'; import HiveDriverError from '../errors/HiveDriverError'; import { buildUserAgentString, normalizePemBytes } from '../utils'; +import driverVersion from '../version'; +import { DRIVER_NAME } from '../telemetry/types'; +import { sanitizeProcessName } from '../telemetry/telemetryUtils'; /** * Default local listener port for the U2M authorization-code callback. @@ -127,6 +132,25 @@ export interface KernelSessionDefaults { retryOverallTimeoutSecs?: number; } +export interface KernelTelemetryOptions { + /** Driver/runtime identity forwarded to kernel-owned telemetry. */ + driverName?: string; + driverVersion?: string; + runtimeName?: string; + runtimeVersion?: string; + runtimeVendor?: string; + osName?: string; + osVersion?: string; + osArch?: string; + clientAppName?: string; + localeName?: string; + charSetEncoding?: string; + processName?: string; + /** Kernel-owned telemetry switch and batching. */ + telemetryEnabled?: boolean; + telemetryBatchSize?: number; +} + /** * TLS options shared across all auth-mode variants. Mirror the napi * binding's `ConnectionOptions.checkServerCertificate` / `.customCaCert` @@ -215,6 +239,7 @@ export interface KernelProxyOptions { export type KernelNativeConnectionOptions = KernelSessionDefaults & KernelTlsOptions & KernelHttpOptions & + KernelTelemetryOptions & KernelProxyOptions & ( | { @@ -520,6 +545,61 @@ export function buildKernelRetryOptions(config: { return out; } +function getLocaleName(env: NodeJS.ProcessEnv = process.env): string { + try { + const lang = env.LANG || env.LC_ALL || env.LC_MESSAGES || ''; + const match = lang.match(/^([a-z]{2}_[A-Z]{2})/); + return match?.[1] ?? 'en_US'; + } catch { + return 'en_US'; + } +} + +function getProcessName(): string { + try { + if (process.title && process.title !== 'node') { + return sanitizeProcessName(process.title) || 'node'; + } + const scriptPath = process.argv?.[1]; + if (scriptPath) { + return sanitizeProcessName(scriptPath).replace(/\.[^.]*$/, '') || 'node'; + } + return 'node'; + } catch { + return 'node'; + } +} + +export function isTelemetryDisabledByEnv(env: NodeJS.ProcessEnv = process.env): boolean { + const raw = env.DATABRICKS_TELEMETRY_DISABLED; + const trimmed = typeof raw === 'string' ? raw.trim() : ''; + return trimmed.length > 0 && /^(1|true|yes|on)$/i.test(trimmed); +} + +export function buildKernelTelemetryOptions(config: Pick) { + const telemetry: KernelTelemetryOptions = { + driverName: DRIVER_NAME, + driverVersion, + runtimeName: 'Node.js', + runtimeVersion: process.version, + runtimeVendor: 'Node.js Foundation', + osName: process.platform, + osVersion: os.release(), + osArch: os.arch(), + clientAppName: undefined, + localeName: getLocaleName(), + charSetEncoding: 'UTF-8', + processName: getProcessName(), + telemetryEnabled: (config.telemetryEnabled ?? true) && !isTelemetryDisabledByEnv(), + }; + + if (Number.isFinite(config.telemetryBatchSize)) { + telemetry.telemetryBatchSize = config.telemetryBatchSize; + } + + return telemetry; +} + /** * Map the public `ConnectionOptions.proxy` (`{protocol, host, port, auth}` — * the same shape the Thrift backend accepts) onto the kernel's structured napi diff --git a/lib/kernel/KernelBackend.ts b/lib/kernel/KernelBackend.ts index 221e7beb..49b67029 100644 --- a/lib/kernel/KernelBackend.ts +++ b/lib/kernel/KernelBackend.ts @@ -22,7 +22,12 @@ import HiveDriverError from '../errors/HiveDriverError'; import { serializeQueryTags } from '../utils'; import { getKernelNative, KernelNativeBinding, KernelConnection } from './KernelNativeLoader'; import { decodeNapiKernelError } from './KernelErrorMapping'; -import { buildKernelConnectionOptions, buildKernelRetryOptions, KernelNativeConnectionOptions } from './KernelAuth'; +import { + buildKernelConnectionOptions, + buildKernelRetryOptions, + buildKernelTelemetryOptions, + KernelNativeConnectionOptions, +} from './KernelAuth'; import { installKernelLogBridge } from './KernelLogging'; import KernelSessionBackend from './KernelSessionBackend'; @@ -94,6 +99,7 @@ export default class KernelBackend implements IBackend { this.nativeOptions = { ...buildKernelConnectionOptions(options), ...buildKernelRetryOptions(this.context.getConfig()), + ...buildKernelTelemetryOptions(this.context.getConfig()), }; // Bridge the Rust kernel's `tracing` logs into the SAME `DBSQLLogger` the diff --git a/native/kernel/index.d.ts b/native/kernel/index.d.ts index 0b042121..eed3508f 100644 --- a/native/kernel/index.d.ts +++ b/native/kernel/index.d.ts @@ -406,6 +406,33 @@ export interface ConnectionOptions { * (~49.7 days) far exceeds any sensible socket timeout. */ socketTimeoutMs?: number + /** + * Driver/system metadata for kernel-owned telemetry. These fields mirror + * the connector telemetry system configuration and are applied when the + * kernel constructs its telemetry exporter. + */ + driverName?: string + driverVersion?: string + runtimeName?: string + runtimeVersion?: string + runtimeVendor?: string + osName?: string + osVersion?: string + osArch?: string + clientAppName?: string + localeName?: string + charSetEncoding?: string + processName?: string + /** + * Whether kernel-owned telemetry is enabled for this session. The wrapper + * should pass the user-facing telemetry opt-out here and avoid emitting + * duplicate wrapper telemetry on the kernel path. + */ + telemetryEnabled?: boolean + /** + * Maximum number of kernel telemetry metrics to batch before flushing. + */ + telemetryBatchSize?: number } /** * Open a Databricks SQL session and return an opaque `Connection` diff --git a/tests/unit/DBSQLClient.test.ts b/tests/unit/DBSQLClient.test.ts index 312cf603..86d4d21a 100644 --- a/tests/unit/DBSQLClient.test.ts +++ b/tests/unit/DBSQLClient.test.ts @@ -4,6 +4,7 @@ import fs from 'fs'; import DBSQLClient, { ThriftLibrary } from '../../lib/DBSQLClient'; import DBSQLSession from '../../lib/DBSQLSession'; import ThriftBackend from '../../lib/thrift-backend/ThriftBackend'; +import KernelBackend from '../../lib/kernel/KernelBackend'; import PlainHttpAuthentication from '../../lib/connection/auth/PlainHttpAuthentication'; import DatabricksOAuth from '../../lib/connection/auth/DatabricksOAuth'; @@ -957,6 +958,22 @@ describe('DBSQLClient telemetry paths', () => { .filter((c) => c.args[0] === LogLevel.warn && /DATABRICKS_TELEMETRY_DISABLED/.test(c.args[1] as string)); expect(warnCalls.length).to.equal(0); }); + + it('does not initialize Node telemetry on the kernel path', async () => { + delete process.env.DATABRICKS_TELEMETRY_DISABLED; + const client = new DBSQLClient(); + const initStub = sinon.stub(client as any, 'initializeTelemetry').resolves(); + sinon.stub(KernelBackend.prototype, 'connect').resolves(); + sinon.stub(KernelBackend.prototype, 'close').resolves(); + + try { + await client.connect({ ...connectOptions, telemetryEnabled: true, useKernel: true } as any); + + expect(initStub.callCount).to.equal(0); + } finally { + await client.close(); + } + }); }); describe('extractWorkspaceId', () => { diff --git a/tests/unit/kernel/_helpers/nativeOptions.ts b/tests/unit/kernel/_helpers/nativeOptions.ts index 995265e2..4fd091ea 100644 --- a/tests/unit/kernel/_helpers/nativeOptions.ts +++ b/tests/unit/kernel/_helpers/nativeOptions.ts @@ -33,6 +33,24 @@ export default function expectNativeConnectionOptions(actual: unknown, expectedR const { customHeaders, ...rest } = actual as Record & { customHeaders?: Array<{ name: string; value: string }>; }; + for (const key of [ + 'driverName', + 'driverVersion', + 'runtimeName', + 'runtimeVersion', + 'runtimeVendor', + 'osName', + 'osVersion', + 'osArch', + 'clientAppName', + 'localeName', + 'charSetEncoding', + 'processName', + 'telemetryEnabled', + 'telemetryBatchSize', + ]) { + delete rest[key]; + } expect(rest).to.deep.equal(expectedRest); expect(customHeaders, 'customHeaders').to.be.an('array').with.lengthOf(1); expect(customHeaders?.[0].name).to.equal('User-Agent'); diff --git a/tests/unit/kernel/execution.test.ts b/tests/unit/kernel/execution.test.ts index 101d698a..d9d780dc 100644 --- a/tests/unit/kernel/execution.test.ts +++ b/tests/unit/kernel/execution.test.ts @@ -418,13 +418,13 @@ function makeBinding(connection: KernelConnection): KernelNativeBinding & { return Object.assign(binding, { openSessionStub }); } -function makeContext(logger?: IDBSQLLogger): IClientContext { +function makeContext(logger?: IDBSQLLogger, configOverrides: Partial = {}): IClientContext { const log: IDBSQLLogger = logger ?? { log(_level: LogLevel, _message: string): void { // no-op }, }; - const config = {} as ClientConfig; + const config = configOverrides as ClientConfig; return { getConfig: () => config, getLogger: () => log, @@ -551,6 +551,79 @@ describe('KernelBackend', () => { }); }); + it('openSession() forwards kernel-owned telemetry config and runtime identity to napi binding', async () => { + const savedEnv = process.env.DATABRICKS_TELEMETRY_DISABLED; + delete process.env.DATABRICKS_TELEMETRY_DISABLED; + + const connection = new FakeNativeConnection(); + const binding = makeBinding(connection); + const backend = new KernelBackend({ + context: makeContext(undefined, { telemetryEnabled: false, telemetryBatchSize: 17 }), + nativeBinding: binding, + }); + + try { + await backend.connect({ + host: 'workspace.example', + path: '/sql/1.0/warehouses/xyz', + token: 'dapi-token', + } as ConnectionOptions); + + await backend.openSession({}); + + const args = binding.openSessionStub.firstCall.args[0] as Record; + expect(args.driverName).to.equal('nodejs-sql-driver'); + expect(args.driverVersion).to.be.a('string').and.not.equal(''); + expect(args.runtimeName).to.equal('Node.js'); + expect(args.runtimeVersion).to.equal(process.version); + expect(args.runtimeVendor).to.equal('Node.js Foundation'); + expect(args.osName).to.equal(process.platform); + expect(args.osVersion).to.be.a('string').and.not.equal(''); + expect(args.osArch).to.be.a('string').and.not.equal(''); + expect(args.localeName).to.be.a('string').and.not.equal(''); + expect(args.charSetEncoding).to.equal('UTF-8'); + expect(args.processName).to.be.a('string').and.not.equal(''); + expect(args.telemetryEnabled).to.equal(false); + expect(args.telemetryBatchSize).to.equal(17); + } finally { + if (savedEnv === undefined) { + delete process.env.DATABRICKS_TELEMETRY_DISABLED; + } else { + process.env.DATABRICKS_TELEMETRY_DISABLED = savedEnv; + } + } + }); + + it('openSession() forwards env-disabled kernel telemetry even when config enables telemetry', async () => { + const savedEnv = process.env.DATABRICKS_TELEMETRY_DISABLED; + process.env.DATABRICKS_TELEMETRY_DISABLED = 'true'; + + const connection = new FakeNativeConnection(); + const binding = makeBinding(connection); + const backend = new KernelBackend({ + context: makeContext(undefined, { telemetryEnabled: true }), + nativeBinding: binding, + }); + + try { + await backend.connect({ + host: 'workspace.example', + path: '/sql/1.0/warehouses/xyz', + token: 'dapi-token', + } as ConnectionOptions); + await backend.openSession({}); + + const args = binding.openSessionStub.firstCall.args[0] as { telemetryEnabled?: boolean }; + expect(args.telemetryEnabled).to.equal(false); + } finally { + if (savedEnv === undefined) { + delete process.env.DATABRICKS_TELEMETRY_DISABLED; + } else { + process.env.DATABRICKS_TELEMETRY_DISABLED = savedEnv; + } + } + }); + it('openSession() serializes session-level queryTags into sessionConf.QUERY_TAGS', async () => { const connection = new FakeNativeConnection(); const binding = makeBinding(connection);