diff --git a/java-bigquery/google-cloud-bigquery-jdbc/.gitignore b/java-bigquery/google-cloud-bigquery-jdbc/.gitignore index 75b2db720ed5..a8633c476dc5 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/.gitignore +++ b/java-bigquery/google-cloud-bigquery-jdbc/.gitignore @@ -1,3 +1,3 @@ drivers/** target-it/** -*logs/** \ No newline at end of file +*logs**/** \ No newline at end of file diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java index 016b3ca470aa..0b6ffc985d81 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java @@ -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> instanceLocals = new ConcurrentHashMap<>(); @@ -35,7 +35,7 @@ public class BigQueryJdbcMdc { private static final InheritableThreadLocal 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( @@ -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 local = instanceLocals.remove(connection); if (local != null) { @@ -71,7 +71,7 @@ public static void removeInstance(BigQueryConnection connection) { } } - public static void clear() { + static void clear() { currentConnectionId.remove(); for (InheritableThreadLocal local : instanceLocals.values()) { local.remove(); diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandler.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandler.java new file mode 100644 index 000000000000..6b8ca7458181 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandler.java @@ -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 handlers = new ConcurrentHashMap<>(); + 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; + 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 { + 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) { + 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(); + } + } + + public void closeHandler(String connectionId) { + if (connectionId != null) { + FileHandler handler = handlers.remove(connectionId); + if (handler != null) { + handler.close(); + } + } + } +} diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java new file mode 100644 index 000000000000..94f8b4f5e6b2 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java @@ -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)); + } +}