diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java index 482faa2b8d..f61b5140f8 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java @@ -34,7 +34,9 @@ public static TelemetryClientFactory getInstance() { public ITelemetryClient getTelemetryClient( IDatabricksConnectionContext connectionContext, DatabricksConfig databricksConfig) { - if (connectionContext != null && connectionContext.isTelemetryEnabled()) { + if (connectionContext != null + && connectionContext.isTelemetryEnabled() + && databricksConfig != null) { return telemetryClients.computeIfAbsent( connectionContext.getConnectionUuid(), k -> diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java index 4673fdd2d9..0640ac04ac 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java @@ -113,44 +113,54 @@ static void exportLatencyLog( long latencyMilliseconds, SqlExecutionEvent executionEvent, StatementId statementId) { - TelemetryEvent telemetryEvent = - new TelemetryEvent() - .setLatency(latencyMilliseconds) - .setSqlOperation(executionEvent) - .setDriverConnectionParameters(getDriverConnectionParameter(connectionContext)); - if (statementId != null) { - telemetryEvent.setSqlStatementId(statementId.toString()); + // Though we already handle null connectionContext in the downstream implementation, + // we are adding this check for extra sanity + if (connectionContext != null) { + TelemetryEvent telemetryEvent = + new TelemetryEvent() + .setLatency(latencyMilliseconds) + .setSqlOperation(executionEvent) + .setDriverConnectionParameters(getDriverConnectionParameter(connectionContext)); + if (statementId != null) { + telemetryEvent.setSqlStatementId(statementId.toString()); + } + TelemetryFrontendLog telemetryFrontendLog = + new TelemetryFrontendLog() + .setFrontendLogEventId(getEventUUID()) + .setContext(getLogContext()) + .setEntry(new FrontendLogEntry().setSqlDriverLog(telemetryEvent)); + TelemetryClientFactory.getInstance() + .getTelemetryClient( + connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) + .exportEvent(telemetryFrontendLog); } - TelemetryFrontendLog telemetryFrontendLog = - new TelemetryFrontendLog() - .setFrontendLogEventId(getEventUUID()) - .setContext(getLogContext()) - .setEntry(new FrontendLogEntry().setSqlDriverLog(telemetryEvent)); - TelemetryClientFactory.getInstance() - .getTelemetryClient(connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) - .exportEvent(telemetryFrontendLog); } public static void exportLatencyLog( IDatabricksConnectionContext connectionContext, long latencyMilliseconds, DriverVolumeOperation volumeOperationEvent) { - TelemetryFrontendLog telemetryFrontendLog = - new TelemetryFrontendLog() - .setFrontendLogEventId(getEventUUID()) - .setContext(getLogContext()) - .setEntry( - new FrontendLogEntry() - .setSqlDriverLog( - new TelemetryEvent() - .setLatency(latencyMilliseconds) - .setVolumeOperation(volumeOperationEvent) - .setDriverConnectionParameters( - getDriverConnectionParameter(connectionContext)))); + // Though we already handle null connectionContext in the downstream implementation, + // we are adding this check for extra sanity + if (connectionContext != null) { + TelemetryFrontendLog telemetryFrontendLog = + new TelemetryFrontendLog() + .setFrontendLogEventId(getEventUUID()) + .setContext(getLogContext()) + .setEntry( + new FrontendLogEntry() + .setSqlDriverLog( + new TelemetryEvent() + .setLatency(latencyMilliseconds) + .setVolumeOperation(volumeOperationEvent) + .setDriverConnectionParameters( + getDriverConnectionParameter(connectionContext)))); - TelemetryClientFactory.getInstance() - .getTelemetryClient(connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) - .exportEvent(telemetryFrontendLog); + TelemetryClientFactory.getInstance() + .getTelemetryClient( + connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) + .exportEvent(telemetryFrontendLog); + } } private static DriverConnectionParameters getDriverConnectionParameter(