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 new file mode 100644 index 000000000000..d0632bf79421 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java @@ -0,0 +1,97 @@ +/* + * 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.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; + +/** + * Lightweight MDC implementation for the BigQuery JDBC driver using InheritableThreadLocal. + * Allocates a dedicated, independent InheritableThreadLocal object per concrete BigQueryConnection + * instance. + */ +public class BigQueryJdbcMdc { + private static final AtomicLong nextId = new AtomicLong(1); + private static final ConcurrentHashMap> + instanceLocals = new ConcurrentHashMap<>(); + private static final ConcurrentHashMap instanceIds = + new ConcurrentHashMap<>(); + + /** Allocates an exclusive InheritableThreadLocal and registers the connection mapping. */ + private static final InheritableThreadLocal currentConnectionId = + new InheritableThreadLocal<>(); + + public static void registerInstance(BigQueryConnection connection, String id) { + if (connection != null) { + String cleanId = + instanceIds.computeIfAbsent( + connection, + k -> { + String suffix = + (id != null && !id.isEmpty()) ? id : String.valueOf(nextId.getAndIncrement()); + return "JdbcConnection-" + suffix; + }); + + currentConnectionId.set(cleanId); + InheritableThreadLocal threadLocal = + instanceLocals.computeIfAbsent(connection, k -> new InheritableThreadLocal<>()); + threadLocal.set(cleanId); + } + } + + /** Retrieves the connection ID mapped to a specific BigQueryConnection instance. */ + public static String getConnectionId(BigQueryConnection connection) { + if (connection != null) { + InheritableThreadLocal local = instanceLocals.get(connection); + if (local != null) { + String val = local.get(); + if (val != null) { + return val; + } + } + return instanceIds.get(connection); + } + return null; + } + + /** + * Returns the connection ID carried by any registered active connection on the current thread. + */ + public 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) { + if (connection != null) { + InheritableThreadLocal local = instanceLocals.remove(connection); + if (local != null) { + local.remove(); + } + instanceIds.remove(connection); + } + } + + public static void clear() { + currentConnectionId.remove(); + for (InheritableThreadLocal local : instanceLocals.values()) { + local.remove(); + } + instanceLocals.clear(); + instanceIds.clear(); + } +} diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java index a6723441550e..973857364ebb 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java @@ -22,8 +22,9 @@ import java.nio.file.Path; import java.nio.file.Paths; import java.nio.file.StandardCopyOption; -import java.text.SimpleDateFormat; -import java.util.Date; +import java.time.Instant; +import java.time.ZoneId; +import java.time.format.DateTimeFormatter; import java.util.Optional; import java.util.logging.ConsoleHandler; import java.util.logging.FileHandler; @@ -62,7 +63,8 @@ class BigQueryJdbcRootLogger { public static Formatter getFormatter() { return new Formatter() { - private static final String PATTERN = "yyyy-MM-dd HH:mm:ss.SSS"; + private final DateTimeFormatter dateFormatter = + DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS").withZone(ZoneId.systemDefault()); private static final String FORMAT = "%1$s %2$5s %3$d --- [%4$-7.15s] %5$-50s %6$-20s: %7$s%8$s"; private static final int MAX_THREAD_NAME_LENGTH = 15; @@ -81,7 +83,9 @@ Optional getThread(long threadId) { @Override public String format(LogRecord record) { - String date = new SimpleDateFormat(PATTERN).format(new Date(record.getMillis())); + + String date = dateFormatter.format(Instant.ofEpochMilli(record.getMillis())); + String threadName = getThread(record.getThreadID()) .map(Thread::getName) diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdcTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdcTest.java new file mode 100644 index 000000000000..4016f7da8105 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdcTest.java @@ -0,0 +1,69 @@ +/* + * 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.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +public class BigQueryJdbcMdcTest { + + private BigQueryConnection mockConnection1; + private BigQueryConnection mockConnection2; + + @BeforeEach + public void setUp() { + mockConnection1 = Mockito.mock(BigQueryConnection.class); + mockConnection2 = Mockito.mock(BigQueryConnection.class); + } + + @AfterEach + public void tearDown() { + BigQueryJdbcMdc.clear(); + } + + @Test + public void testRegisterAndRetrieveConnectionId() { + BigQueryJdbcMdc.registerInstance(mockConnection1, "123"); + + assertEquals("JdbcConnection-123", BigQueryJdbcMdc.getConnectionId(mockConnection1)); + assertEquals("JdbcConnection-123", BigQueryJdbcMdc.getConnectionId()); + } + + @Test + public void testClearContext() { + BigQueryJdbcMdc.registerInstance(mockConnection1, "456"); + assertEquals("JdbcConnection-456", BigQueryJdbcMdc.getConnectionId(mockConnection1)); + + BigQueryJdbcMdc.clear(); + + assertNull(BigQueryJdbcMdc.getConnectionId()); + } + + @Test + public void testMultipleConnectionsSameThread() { + BigQueryJdbcMdc.registerInstance(mockConnection1, "111"); + BigQueryJdbcMdc.registerInstance(mockConnection2, "222"); + + assertEquals("JdbcConnection-111", BigQueryJdbcMdc.getConnectionId(mockConnection1)); + assertEquals("JdbcConnection-222", BigQueryJdbcMdc.getConnectionId(mockConnection2)); + } +}