Skip to content
Merged
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,129 @@
/*
* 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 Path 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 = Paths.get(baseLogPath != null ? baseLogPath : "").toAbsolutePath();
this.level = level;

try {
if (!Files.exists(this.baseLogPath)) {
Files.createDirectories(this.baseLogPath);
}

this.defaultHandler = createFileHandler("Jdbc-default");
} catch (IOException e) {
reportError(
"Failed to initialize default log file", e, java.util.logging.ErrorManager.OPEN_FAILURE);
}
}

private String getLogFilePath(String id) {
return baseLogPath.resolve("BigQuery-" + id + ".log").toString();
}

private FileHandler createFileHandler(String id) {
try {
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 null;
}
}

@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, this::createFileHandler);
if (handler == null) {
handler = defaultHandler;
}
}

if (handler != null) {
Comment thread
Neenu1995 marked this conversation as resolved.
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) {
Comment thread
Neenu1995 marked this conversation as resolved.
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