From 9044bd3682694743a9f0937af67939bafb95d795 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 17 Apr 2026 12:40:33 -0400 Subject: [PATCH 01/10] add BigQueryJdbcMdc --- .../cloud/bigquery/jdbc/BigQueryJdbcMdc.java | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java 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..8d829a572a6e --- /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(); + } +} \ No newline at end of file From 066f36b991de877de202301a44b200c971018628 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 17 Apr 2026 12:41:38 -0400 Subject: [PATCH 02/10] Use DateTimeFormatter instead of SimpleDateFormat for better performance --- .../bigquery/jdbc/BigQueryJdbcRootLogger.java | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) 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..affd47bd41d7 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) @@ -196,4 +200,4 @@ static void setPath(String logPath) { logger.warning("Log File warning : " + ex); } } -} +} \ No newline at end of file From f7ac0ebd2695ace5ec0aef89443945c9003962a4 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 17 Apr 2026 12:42:42 -0400 Subject: [PATCH 03/10] Add BigQueryJdbcMdcTest --- .../bigquery/jdbc/BigQueryJdbcMdcTest.java | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdcTest.java 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..f624c5206d81 --- /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)); + } +} \ No newline at end of file From 93d7b5d25a7397213e21636787f3943fe4ebb292 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 17 Apr 2026 12:43:23 -0400 Subject: [PATCH 04/10] reformat --- .../java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java | 2 +- .../com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java | 2 +- .../com/google/cloud/bigquery/jdbc/BigQueryJdbcMdcTest.java | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) 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 8d829a572a6e..d0632bf79421 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 @@ -94,4 +94,4 @@ public static void clear() { instanceLocals.clear(); instanceIds.clear(); } -} \ No newline at end of file +} 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 affd47bd41d7..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 @@ -200,4 +200,4 @@ static void setPath(String logPath) { logger.warning("Log File warning : " + ex); } } -} \ No newline at end of file +} 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 index f624c5206d81..4016f7da8105 100644 --- 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 @@ -66,4 +66,4 @@ public void testMultipleConnectionsSameThread() { assertEquals("JdbcConnection-111", BigQueryJdbcMdc.getConnectionId(mockConnection1)); assertEquals("JdbcConnection-222", BigQueryJdbcMdc.getConnectionId(mockConnection2)); } -} \ No newline at end of file +} From 17920d434c4372774704f10f4f18cebfb9009d93 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Mon, 20 Apr 2026 14:12:10 -0400 Subject: [PATCH 05/10] fix clear method --- .../java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java | 2 -- 1 file changed, 2 deletions(-) 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 d0632bf79421..7caa1b0e38d3 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 @@ -91,7 +91,5 @@ public static void clear() { for (InheritableThreadLocal local : instanceLocals.values()) { local.remove(); } - instanceLocals.clear(); - instanceIds.clear(); } } From 0d5aa137cc29321144b6c2493e92f0d3f72f9794 Mon Sep 17 00:00:00 2001 From: cloud-java-bot Date: Mon, 20 Apr 2026 18:18:34 +0000 Subject: [PATCH 06/10] chore: generate libraries at Mon Apr 20 18:16:36 UTC 2026 --- java-iam/.repo-metadata.json | 1 - java-iam/README.md | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/java-iam/.repo-metadata.json b/java-iam/.repo-metadata.json index fa9ab8c76a75..d35667f774d9 100644 --- a/java-iam/.repo-metadata.json +++ b/java-iam/.repo-metadata.json @@ -10,7 +10,6 @@ "repo": "googleapis/google-cloud-java", "repo_short": "java-iam", "distribution_name": "com.google.cloud:google-iam-policy", - "api_id": "iam.googleapis.com", "library_type": "GAPIC_AUTO", "requires_billing": true, "excluded_dependencies": "google-iam-policy", diff --git a/java-iam/README.md b/java-iam/README.md index a31d56fecfdc..b5f33684f9cf 100644 --- a/java-iam/README.md +++ b/java-iam/README.md @@ -188,7 +188,7 @@ Java is a registered trademark of Oracle and/or its affiliates. [code-of-conduct]: https://github.com/googleapis/google-cloud-java/blob/main/CODE_OF_CONDUCT.md#contributor-code-of-conduct [license]: https://github.com/googleapis/google-cloud-java/blob/main/LICENSE [enable-billing]: https://cloud.google.com/apis/docs/getting-started#enabling_billing -[enable-api]: https://console.cloud.google.com/flows/enableapi?apiid=iam.googleapis.com + [libraries-bom]: https://github.com/GoogleCloudPlatform/cloud-opensource-java/wiki/The-Google-Cloud-Platform-Libraries-BOM [shell_img]: https://gstatic.com/cloudssh/images/open-btn.png From d77b0a700b4b22eed74cee1f18af5998a6acd733 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Mon, 20 Apr 2026 17:03:36 -0400 Subject: [PATCH 07/10] revert unknown file changes --- java-iam-policy/.repo-metadata.json | 1 + java-iam-policy/README.md | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/java-iam-policy/.repo-metadata.json b/java-iam-policy/.repo-metadata.json index bf65bd5c8655..0d797eebf786 100644 --- a/java-iam-policy/.repo-metadata.json +++ b/java-iam-policy/.repo-metadata.json @@ -10,6 +10,7 @@ "repo": "googleapis/google-cloud-java", "repo_short": "java-iam-policy", "distribution_name": "com.google.cloud:google-iam-policy", + "api_id": "iam.googleapis.com", "library_type": "GAPIC_AUTO", "requires_billing": true, "excluded_dependencies": "google-iam-policy", diff --git a/java-iam-policy/README.md b/java-iam-policy/README.md index b5f33684f9cf..a31d56fecfdc 100644 --- a/java-iam-policy/README.md +++ b/java-iam-policy/README.md @@ -188,7 +188,7 @@ Java is a registered trademark of Oracle and/or its affiliates. [code-of-conduct]: https://github.com/googleapis/google-cloud-java/blob/main/CODE_OF_CONDUCT.md#contributor-code-of-conduct [license]: https://github.com/googleapis/google-cloud-java/blob/main/LICENSE [enable-billing]: https://cloud.google.com/apis/docs/getting-started#enabling_billing - +[enable-api]: https://console.cloud.google.com/flows/enableapi?apiid=iam.googleapis.com [libraries-bom]: https://github.com/GoogleCloudPlatform/cloud-opensource-java/wiki/The-Google-Cloud-Platform-Libraries-BOM [shell_img]: https://gstatic.com/cloudssh/images/open-btn.png From eadc09eef1ec07871a100a71c53175c845986102 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Wed, 22 Apr 2026 09:23:12 -0400 Subject: [PATCH 08/10] revert changes to BigQueryJdbcRootLogger --- .../cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) 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 973857364ebb..a6723441550e 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,9 +22,8 @@ import java.nio.file.Path; import java.nio.file.Paths; import java.nio.file.StandardCopyOption; -import java.time.Instant; -import java.time.ZoneId; -import java.time.format.DateTimeFormatter; +import java.text.SimpleDateFormat; +import java.util.Date; import java.util.Optional; import java.util.logging.ConsoleHandler; import java.util.logging.FileHandler; @@ -63,8 +62,7 @@ class BigQueryJdbcRootLogger { public static Formatter getFormatter() { return new Formatter() { - private final DateTimeFormatter dateFormatter = - DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS").withZone(ZoneId.systemDefault()); + private static final String PATTERN = "yyyy-MM-dd HH:mm:ss.SSS"; 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; @@ -83,9 +81,7 @@ Optional getThread(long threadId) { @Override public String format(LogRecord record) { - - String date = dateFormatter.format(Instant.ofEpochMilli(record.getMillis())); - + String date = new SimpleDateFormat(PATTERN).format(new Date(record.getMillis())); String threadName = getThread(record.getThreadID()) .map(Thread::getName) From 9febba06f552b009a1754c2983401f596781803b Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Wed, 22 Apr 2026 09:24:35 -0400 Subject: [PATCH 09/10] remove getConnectionID method --- .../cloud/bigquery/jdbc/BigQueryJdbcMdc.java | 15 --------------- 1 file changed, 15 deletions(-) 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 7caa1b0e38d3..016b3ca470aa 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 @@ -53,21 +53,6 @@ public static void registerInstance(BigQueryConnection connection, String id) { } } - /** 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. */ From 39e64edea1a8b9cfb27f061255cf328e13554940 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Wed, 22 Apr 2026 10:07:24 -0400 Subject: [PATCH 10/10] refactor test --- .../bigquery/jdbc/BigQueryJdbcMdcTest.java | 101 ++++++++++++++++-- 1 file changed, 92 insertions(+), 9 deletions(-) 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 index 4016f7da8105..764b51f38d49 100644 --- 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 @@ -18,7 +18,11 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -43,27 +47,106 @@ public void tearDown() { @Test public void testRegisterAndRetrieveConnectionId() { BigQueryJdbcMdc.registerInstance(mockConnection1, "123"); - - assertEquals("JdbcConnection-123", BigQueryJdbcMdc.getConnectionId(mockConnection1)); assertEquals("JdbcConnection-123", BigQueryJdbcMdc.getConnectionId()); } + @Test + public void testRemoveInstance() { + BigQueryJdbcMdc.registerInstance(mockConnection1, "1"); + assertEquals("JdbcConnection-1", BigQueryJdbcMdc.getConnectionId()); + + BigQueryJdbcMdc.removeInstance(mockConnection1); + // Note: removeInstance does not clear currentConnectionId on the current thread + // based on current implementation. + assertEquals("JdbcConnection-1", BigQueryJdbcMdc.getConnectionId()); + + BigQueryJdbcMdc.clear(); + assertNull(BigQueryJdbcMdc.getConnectionId()); + } + @Test public void testClearContext() { BigQueryJdbcMdc.registerInstance(mockConnection1, "456"); - assertEquals("JdbcConnection-456", BigQueryJdbcMdc.getConnectionId(mockConnection1)); + assertEquals("JdbcConnection-456", BigQueryJdbcMdc.getConnectionId()); BigQueryJdbcMdc.clear(); - assertNull(BigQueryJdbcMdc.getConnectionId()); } @Test - public void testMultipleConnectionsSameThread() { - BigQueryJdbcMdc.registerInstance(mockConnection1, "111"); - BigQueryJdbcMdc.registerInstance(mockConnection2, "222"); + public void testThreadInheritance() throws InterruptedException { + BigQueryJdbcMdc.registerInstance(mockConnection1, "parent"); + assertEquals("JdbcConnection-parent", BigQueryJdbcMdc.getConnectionId()); + + AtomicReference childConnectionId = new AtomicReference<>(); + CountDownLatch latch = new CountDownLatch(1); + + Thread childThread = + new Thread( + () -> { + childConnectionId.set(BigQueryJdbcMdc.getConnectionId()); + latch.countDown(); + }); + childThread.start(); + assertTrue(latch.await(5, TimeUnit.SECONDS)); + + assertEquals("JdbcConnection-parent", childConnectionId.get()); + } + + @Test + public void testThreadIsolation() throws InterruptedException { + CountDownLatch threadARegistered = new CountDownLatch(1); + CountDownLatch threadBChecked = new CountDownLatch(1); + CountDownLatch threadBRegistered = new CountDownLatch(1); + CountDownLatch testFinished = new CountDownLatch(2); + + AtomicReference threadAIdBeforeB = new AtomicReference<>(); + AtomicReference threadAIdAfterB = new AtomicReference<>(); + AtomicReference threadBIdBeforeRegister = new AtomicReference<>(); + AtomicReference threadBIdAfterRegister = new AtomicReference<>(); + + Thread threadA = + new Thread( + () -> { + try { + BigQueryJdbcMdc.registerInstance(mockConnection1, "A"); + threadAIdBeforeB.set(BigQueryJdbcMdc.getConnectionId()); + threadARegistered.countDown(); + + threadBRegistered.await(); + threadAIdAfterB.set(BigQueryJdbcMdc.getConnectionId()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + testFinished.countDown(); + } + }); + + Thread threadB = + new Thread( + () -> { + try { + threadARegistered.await(); + threadBIdBeforeRegister.set(BigQueryJdbcMdc.getConnectionId()); + + BigQueryJdbcMdc.registerInstance(mockConnection2, "B"); + threadBIdAfterRegister.set(BigQueryJdbcMdc.getConnectionId()); + threadBRegistered.countDown(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + testFinished.countDown(); + } + }); + + threadA.start(); + threadB.start(); + + assertTrue(testFinished.await(5, TimeUnit.SECONDS)); - assertEquals("JdbcConnection-111", BigQueryJdbcMdc.getConnectionId(mockConnection1)); - assertEquals("JdbcConnection-222", BigQueryJdbcMdc.getConnectionId(mockConnection2)); + assertEquals("JdbcConnection-A", threadAIdBeforeB.get()); + assertNull(threadBIdBeforeRegister.get()); + assertEquals("JdbcConnection-B", threadBIdAfterRegister.get()); + assertEquals("JdbcConnection-A", threadAIdAfterB.get()); } }