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
2 changes: 1 addition & 1 deletion java-bigquery/google-cloud-bigquery-jdbc/.gitignore
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
drivers/**
target-it/**
*logs/**
*logs**/**
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
* Allocates a dedicated, independent InheritableThreadLocal object per concrete BigQueryConnection
* instance.
*/
public class BigQueryJdbcMdc {
class BigQueryJdbcMdc {
private static final AtomicLong nextId = new AtomicLong(1);
private static final ConcurrentHashMap<BigQueryConnection, InheritableThreadLocal<String>>
instanceLocals = new ConcurrentHashMap<>();
Expand All @@ -35,7 +35,7 @@ public class BigQueryJdbcMdc {
private static final InheritableThreadLocal<String> currentConnectionId =
new InheritableThreadLocal<>();

public static void registerInstance(BigQueryConnection connection, String id) {
static void registerInstance(BigQueryConnection connection, String id) {
if (connection != null) {
String cleanId =
instanceIds.computeIfAbsent(
Expand All @@ -56,12 +56,12 @@ public static void registerInstance(BigQueryConnection connection, String id) {
/**
* Returns the connection ID carried by any registered active connection on the current thread.
*/
public static String getConnectionId() {
static String getConnectionId() {
return currentConnectionId.get();
}

/** Clears the connection ID context from all active connection contexts on the current thread. */
public static void removeInstance(BigQueryConnection connection) {
static void removeInstance(BigQueryConnection connection) {
if (connection != null) {
InheritableThreadLocal<String> local = instanceLocals.remove(connection);
if (local != null) {
Expand All @@ -71,7 +71,7 @@ public static void removeInstance(BigQueryConnection connection) {
}
}

public static void clear() {
static void clear() {
currentConnectionId.remove();
for (InheritableThreadLocal<String> local : instanceLocals.values()) {
local.remove();
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,137 @@
/*
* Copyright 2026 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package com.google.cloud.bigquery.jdbc;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.concurrent.ConcurrentHashMap;
import java.util.logging.FileHandler;
import java.util.logging.Handler;
import java.util.logging.Level;
import java.util.logging.LogRecord;

/**
* Custom logging handler that dynamically creates and routes log records to per-connection specific
* log files using the connection ID retrieved from BigQueryJdbcMdc.
*/
class PerConnectionFileHandler extends Handler {
private final String baseLogPath;
private final Level level;
private final ConcurrentHashMap<String, FileHandler> handlers = new ConcurrentHashMap<>();
Comment thread
Neenu1995 marked this conversation as resolved.
private FileHandler defaultHandler;

PerConnectionFileHandler(String baseLogPath, Level level) {
this.baseLogPath = baseLogPath != null ? baseLogPath : "";
this.level = level;

try {
if (!baseLogPath.isEmpty()) {
Path dir = Paths.get(baseLogPath);
if (!Files.exists(dir)) {
Files.createDirectories(dir);
}
}

String defaultPath = getLogFilePath("Jdbc-default");
this.defaultHandler = new FileHandler(defaultPath, 0, 1, true);
this.defaultHandler.setLevel(level);
this.defaultHandler.setFormatter(BigQueryJdbcRootLogger.getFormatter());
} catch (IOException e) {
reportError(
"Failed to initialize default log file", e, java.util.logging.ErrorManager.OPEN_FAILURE);
}
}

private String getLogFilePath(String id) {
String path = baseLogPath;
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we store baseLogPath as a Path (maybe using toAbsolutePath when initialized and use `baseLogPath.resolve("BigQuery-" + id + ".log")?

if (!path.isEmpty() && !path.endsWith("/")) {
path = path + "/";
}
return path + "BigQuery-" + id + ".log";
}

@Override
public void publish(LogRecord record) {
if (!isLoggable(record)) {
return;
}

String connectionId = BigQueryJdbcMdc.getConnectionId();
FileHandler handler = defaultHandler;

if (connectionId != null && !connectionId.isEmpty()) {
handler =
handlers.computeIfAbsent(
connectionId,
id -> {
try {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you move it to a helper function instead of lambda? (And can be merged with setup for defaultHandler)

String filePath = getLogFilePath(id);
FileHandler fh = new FileHandler(filePath, 0, 1, true);
fh.setLevel(level);
fh.setFormatter(BigQueryJdbcRootLogger.getFormatter());
return fh;
} catch (IOException e) {
reportError(
"Failed to create log file for connection " + id,
e,
java.util.logging.ErrorManager.OPEN_FAILURE);
return defaultHandler;
}
});
}

if (handler != null) {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can handler be null?

handler.publish(record);
}
}

@Override
public void flush() {
if (defaultHandler != null) {
defaultHandler.flush();
}
for (FileHandler h : handlers.values()) {
h.flush();
}
}

@Override
public void close() throws SecurityException {
for (FileHandler h : handlers.values()) {
try {
h.close();
} catch (Exception e) {
}
}
try {
if (defaultHandler != null) defaultHandler.close();
} finally {
handlers.clear();
}
}
Comment thread
Neenu1995 marked this conversation as resolved.

public void closeHandler(String connectionId) {
if (connectionId != null) {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: I think this check is not required. If it is null, .remove() is going to return null which is handled.

FileHandler handler = handlers.remove(connectionId);
if (handler != null) {
handler.close();
}
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
/*
* Copyright 2026 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package com.google.cloud.bigquery.jdbc;

import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.logging.Level;
import java.util.logging.LogRecord;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import org.mockito.Mockito;

public class PerConnectionFileHandlerTest {

@TempDir Path tempDir;

private PerConnectionFileHandler handler;
private BigQueryConnection mockConnection;

@BeforeEach
public void setUp() {
handler = new PerConnectionFileHandler(tempDir.toString(), Level.INFO);
mockConnection = Mockito.mock(BigQueryConnection.class);
BigQueryJdbcMdc.clear();
}

@AfterEach
public void tearDown() {
if (handler != null) {
handler.close();
}
BigQueryJdbcMdc.clear();
}

@Test
public void testInitialization() {
Path defaultLog = tempDir.resolve("BigQuery-Jdbc-default.log");
assertTrue(Files.exists(defaultLog));
}

@Test
public void testPublishDefault() throws IOException {
LogRecord record = new LogRecord(Level.INFO, "Test message default");
handler.publish(record);
handler.flush();

Path defaultLog = tempDir.resolve("BigQuery-Jdbc-default.log");
String content = new String(Files.readAllBytes(defaultLog));
assertTrue(content.contains("Test message default"));
}

@Test
public void testPublishConnectionSpecific() throws IOException {
BigQueryJdbcMdc.registerInstance(mockConnection, "123");

LogRecord record = new LogRecord(Level.INFO, "Test message connection 123");
handler.publish(record);
handler.flush();

Path connLog = tempDir.resolve("BigQuery-JdbcConnection-123.log");
assertTrue(Files.exists(connLog));
String content = new String(Files.readAllBytes(connLog));
assertTrue(content.contains("Test message connection 123"));
}

@Test
public void testCloseHandler() {
BigQueryJdbcMdc.registerInstance(mockConnection, "456");

LogRecord record = new LogRecord(Level.INFO, "Test message connection 456");
handler.publish(record);
handler.flush();

Path connLog = tempDir.resolve("BigQuery-JdbcConnection-456.log");
assertTrue(Files.exists(connLog));

handler.closeHandler("JdbcConnection-456");
assertTrue(Files.exists(connLog));
}
}
Loading