Skip to content

Commit 8d8b704

Browse files
committed
Forward kernel telemetry options
1 parent 2406f31 commit 8d8b704

7 files changed

Lines changed: 227 additions & 5 deletions

File tree

lib/DBSQLClient.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -729,7 +729,8 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
729729
// doesn't ship in the public `.d.ts`. Mirrors Python's `kwargs.get("use_kernel")`
730730
// pattern (see databricks-sql-python/src/databricks/sql/session.py).
731731
const internalOptions = options as ConnectionOptions & InternalConnectionOptions;
732-
const backend = internalOptions.useKernel
732+
const useKernel = internalOptions.useKernel === true;
733+
const backend = useKernel
733734
? new KernelBackend({ context: this })
734735
: new ThriftBackend({
735736
context: this,
@@ -777,7 +778,7 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
777778
`Telemetry remains controlled by the runtime config and feature flag.`,
778779
);
779780
}
780-
if (this.config.telemetryEnabled && !envDisabled) {
781+
if (!useKernel && this.config.telemetryEnabled && !envDisabled) {
781782
await this.initializeTelemetry();
782783
}
783784

lib/kernel/KernelAuth.ts

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,16 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15+
import os from 'os';
1516
import { ConnectionOptions } from '../contracts/IDBSQLClient';
17+
import { ClientConfig } from '../contracts/IClientContext';
1618
import { InternalConnectionOptions } from '../contracts/InternalConnectionOptions';
1719
import AuthenticationError from '../errors/AuthenticationError';
1820
import HiveDriverError from '../errors/HiveDriverError';
1921
import { buildUserAgentString, normalizePemBytes } from '../utils';
22+
import driverVersion from '../version';
23+
import { DRIVER_NAME } from '../telemetry/types';
24+
import { sanitizeProcessName } from '../telemetry/telemetryUtils';
2025

2126
/**
2227
* Default local listener port for the U2M authorization-code callback.
@@ -127,6 +132,25 @@ export interface KernelSessionDefaults {
127132
retryOverallTimeoutSecs?: number;
128133
}
129134

135+
export interface KernelTelemetryOptions {
136+
/** Driver/runtime identity forwarded to kernel-owned telemetry. */
137+
driverName?: string;
138+
driverVersion?: string;
139+
runtimeName?: string;
140+
runtimeVersion?: string;
141+
runtimeVendor?: string;
142+
osName?: string;
143+
osVersion?: string;
144+
osArch?: string;
145+
clientAppName?: string;
146+
localeName?: string;
147+
charSetEncoding?: string;
148+
processName?: string;
149+
/** Kernel-owned telemetry switch and batching. */
150+
telemetryEnabled?: boolean;
151+
telemetryBatchSize?: number;
152+
}
153+
130154
/**
131155
* TLS options shared across all auth-mode variants. Mirror the napi
132156
* binding's `ConnectionOptions.checkServerCertificate` / `.customCaCert`
@@ -215,6 +239,7 @@ export interface KernelProxyOptions {
215239
export type KernelNativeConnectionOptions = KernelSessionDefaults &
216240
KernelTlsOptions &
217241
KernelHttpOptions &
242+
KernelTelemetryOptions &
218243
KernelProxyOptions &
219244
(
220245
| {
@@ -520,6 +545,61 @@ export function buildKernelRetryOptions(config: {
520545
return out;
521546
}
522547

548+
function getLocaleName(env: NodeJS.ProcessEnv = process.env): string {
549+
try {
550+
const lang = env.LANG || env.LC_ALL || env.LC_MESSAGES || '';
551+
const match = lang.match(/^([a-z]{2}_[A-Z]{2})/);
552+
return match?.[1] ?? 'en_US';
553+
} catch {
554+
return 'en_US';
555+
}
556+
}
557+
558+
function getProcessName(): string {
559+
try {
560+
if (process.title && process.title !== 'node') {
561+
return sanitizeProcessName(process.title) || 'node';
562+
}
563+
const scriptPath = process.argv?.[1];
564+
if (scriptPath) {
565+
return sanitizeProcessName(scriptPath).replace(/\.[^.]*$/, '') || 'node';
566+
}
567+
return 'node';
568+
} catch {
569+
return 'node';
570+
}
571+
}
572+
573+
export function isTelemetryDisabledByEnv(env: NodeJS.ProcessEnv = process.env): boolean {
574+
const raw = env.DATABRICKS_TELEMETRY_DISABLED;
575+
const trimmed = typeof raw === 'string' ? raw.trim() : '';
576+
return trimmed.length > 0 && /^(1|true|yes|on)$/i.test(trimmed);
577+
}
578+
579+
export function buildKernelTelemetryOptions(config: Pick<ClientConfig, 'telemetryEnabled' | 'telemetryBatchSize'>) {
580+
const telemetry: KernelTelemetryOptions = {
581+
driverName: DRIVER_NAME,
582+
driverVersion,
583+
runtimeName: 'Node.js',
584+
runtimeVersion: process.version,
585+
runtimeVendor: 'Node.js Foundation',
586+
osName: process.platform,
587+
osVersion: os.release(),
588+
osArch: os.arch(),
589+
clientAppName: undefined,
590+
localeName: getLocaleName(),
591+
charSetEncoding: 'UTF-8',
592+
processName: getProcessName(),
593+
telemetryEnabled: (config.telemetryEnabled ?? true) && !isTelemetryDisabledByEnv(),
594+
};
595+
596+
if (Number.isFinite(config.telemetryBatchSize)) {
597+
telemetry.telemetryBatchSize = config.telemetryBatchSize;
598+
}
599+
600+
return telemetry;
601+
}
602+
523603
/**
524604
* Map the public `ConnectionOptions.proxy` (`{protocol, host, port, auth}` —
525605
* the same shape the Thrift backend accepts) onto the kernel's structured napi

lib/kernel/KernelBackend.ts

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,12 @@ import HiveDriverError from '../errors/HiveDriverError';
2222
import { serializeQueryTags } from '../utils';
2323
import { getKernelNative, KernelNativeBinding, KernelConnection } from './KernelNativeLoader';
2424
import { decodeNapiKernelError } from './KernelErrorMapping';
25-
import { buildKernelConnectionOptions, buildKernelRetryOptions, KernelNativeConnectionOptions } from './KernelAuth';
25+
import {
26+
buildKernelConnectionOptions,
27+
buildKernelRetryOptions,
28+
buildKernelTelemetryOptions,
29+
KernelNativeConnectionOptions,
30+
} from './KernelAuth';
2631
import { installKernelLogBridge } from './KernelLogging';
2732
import KernelSessionBackend from './KernelSessionBackend';
2833

@@ -94,6 +99,7 @@ export default class KernelBackend implements IBackend {
9499
this.nativeOptions = {
95100
...buildKernelConnectionOptions(options),
96101
...buildKernelRetryOptions(this.context.getConfig()),
102+
...buildKernelTelemetryOptions(this.context.getConfig()),
97103
};
98104

99105
// Bridge the Rust kernel's `tracing` logs into the SAME `DBSQLLogger` the

native/kernel/index.d.ts

Lines changed: 27 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

tests/unit/DBSQLClient.test.ts

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import fs from 'fs';
44
import DBSQLClient, { ThriftLibrary } from '../../lib/DBSQLClient';
55
import DBSQLSession from '../../lib/DBSQLSession';
66
import ThriftBackend from '../../lib/thrift-backend/ThriftBackend';
7+
import KernelBackend from '../../lib/kernel/KernelBackend';
78

89
import PlainHttpAuthentication from '../../lib/connection/auth/PlainHttpAuthentication';
910
import DatabricksOAuth from '../../lib/connection/auth/DatabricksOAuth';
@@ -957,6 +958,22 @@ describe('DBSQLClient telemetry paths', () => {
957958
.filter((c) => c.args[0] === LogLevel.warn && /DATABRICKS_TELEMETRY_DISABLED/.test(c.args[1] as string));
958959
expect(warnCalls.length).to.equal(0);
959960
});
961+
962+
it('does not initialize Node telemetry on the kernel path', async () => {
963+
delete process.env.DATABRICKS_TELEMETRY_DISABLED;
964+
const client = new DBSQLClient();
965+
const initStub = sinon.stub(client as any, 'initializeTelemetry').resolves();
966+
sinon.stub(KernelBackend.prototype, 'connect').resolves();
967+
sinon.stub(KernelBackend.prototype, 'close').resolves();
968+
969+
try {
970+
await client.connect({ ...connectOptions, telemetryEnabled: true, useKernel: true } as any);
971+
972+
expect(initStub.callCount).to.equal(0);
973+
} finally {
974+
await client.close();
975+
}
976+
});
960977
});
961978

962979
describe('extractWorkspaceId', () => {

tests/unit/kernel/_helpers/nativeOptions.ts

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,24 @@ export default function expectNativeConnectionOptions(actual: unknown, expectedR
3333
const { customHeaders, ...rest } = actual as Record<string, unknown> & {
3434
customHeaders?: Array<{ name: string; value: string }>;
3535
};
36+
for (const key of [
37+
'driverName',
38+
'driverVersion',
39+
'runtimeName',
40+
'runtimeVersion',
41+
'runtimeVendor',
42+
'osName',
43+
'osVersion',
44+
'osArch',
45+
'clientAppName',
46+
'localeName',
47+
'charSetEncoding',
48+
'processName',
49+
'telemetryEnabled',
50+
'telemetryBatchSize',
51+
]) {
52+
delete rest[key];
53+
}
3654
expect(rest).to.deep.equal(expectedRest);
3755
expect(customHeaders, 'customHeaders').to.be.an('array').with.lengthOf(1);
3856
expect(customHeaders?.[0].name).to.equal('User-Agent');

tests/unit/kernel/execution.test.ts

Lines changed: 75 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -418,13 +418,13 @@ function makeBinding(connection: KernelConnection): KernelNativeBinding & {
418418
return Object.assign(binding, { openSessionStub });
419419
}
420420

421-
function makeContext(logger?: IDBSQLLogger): IClientContext {
421+
function makeContext(logger?: IDBSQLLogger, configOverrides: Partial<ClientConfig> = {}): IClientContext {
422422
const log: IDBSQLLogger = logger ?? {
423423
log(_level: LogLevel, _message: string): void {
424424
// no-op
425425
},
426426
};
427-
const config = {} as ClientConfig;
427+
const config = configOverrides as ClientConfig;
428428
return {
429429
getConfig: () => config,
430430
getLogger: () => log,
@@ -551,6 +551,79 @@ describe('KernelBackend', () => {
551551
});
552552
});
553553

554+
it('openSession() forwards kernel-owned telemetry config and runtime identity to napi binding', async () => {
555+
const savedEnv = process.env.DATABRICKS_TELEMETRY_DISABLED;
556+
delete process.env.DATABRICKS_TELEMETRY_DISABLED;
557+
558+
const connection = new FakeNativeConnection();
559+
const binding = makeBinding(connection);
560+
const backend = new KernelBackend({
561+
context: makeContext(undefined, { telemetryEnabled: false, telemetryBatchSize: 17 }),
562+
nativeBinding: binding,
563+
});
564+
565+
try {
566+
await backend.connect({
567+
host: 'workspace.example',
568+
path: '/sql/1.0/warehouses/xyz',
569+
token: 'dapi-token',
570+
} as ConnectionOptions);
571+
572+
await backend.openSession({});
573+
574+
const args = binding.openSessionStub.firstCall.args[0] as Record<string, unknown>;
575+
expect(args.driverName).to.equal('nodejs-sql-driver');
576+
expect(args.driverVersion).to.be.a('string').and.not.equal('');
577+
expect(args.runtimeName).to.equal('Node.js');
578+
expect(args.runtimeVersion).to.equal(process.version);
579+
expect(args.runtimeVendor).to.equal('Node.js Foundation');
580+
expect(args.osName).to.equal(process.platform);
581+
expect(args.osVersion).to.be.a('string').and.not.equal('');
582+
expect(args.osArch).to.be.a('string').and.not.equal('');
583+
expect(args.localeName).to.be.a('string').and.not.equal('');
584+
expect(args.charSetEncoding).to.equal('UTF-8');
585+
expect(args.processName).to.be.a('string').and.not.equal('');
586+
expect(args.telemetryEnabled).to.equal(false);
587+
expect(args.telemetryBatchSize).to.equal(17);
588+
} finally {
589+
if (savedEnv === undefined) {
590+
delete process.env.DATABRICKS_TELEMETRY_DISABLED;
591+
} else {
592+
process.env.DATABRICKS_TELEMETRY_DISABLED = savedEnv;
593+
}
594+
}
595+
});
596+
597+
it('openSession() forwards env-disabled kernel telemetry even when config enables telemetry', async () => {
598+
const savedEnv = process.env.DATABRICKS_TELEMETRY_DISABLED;
599+
process.env.DATABRICKS_TELEMETRY_DISABLED = 'true';
600+
601+
const connection = new FakeNativeConnection();
602+
const binding = makeBinding(connection);
603+
const backend = new KernelBackend({
604+
context: makeContext(undefined, { telemetryEnabled: true }),
605+
nativeBinding: binding,
606+
});
607+
608+
try {
609+
await backend.connect({
610+
host: 'workspace.example',
611+
path: '/sql/1.0/warehouses/xyz',
612+
token: 'dapi-token',
613+
} as ConnectionOptions);
614+
await backend.openSession({});
615+
616+
const args = binding.openSessionStub.firstCall.args[0] as { telemetryEnabled?: boolean };
617+
expect(args.telemetryEnabled).to.equal(false);
618+
} finally {
619+
if (savedEnv === undefined) {
620+
delete process.env.DATABRICKS_TELEMETRY_DISABLED;
621+
} else {
622+
process.env.DATABRICKS_TELEMETRY_DISABLED = savedEnv;
623+
}
624+
}
625+
});
626+
554627
it('openSession() serializes session-level queryTags into sessionConf.QUERY_TAGS', async () => {
555628
const connection = new FakeNativeConnection();
556629
const binding = makeBinding(connection);

0 commit comments

Comments
 (0)