Skip to content
Merged
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
/*
* 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);
Comment thread
Neenu1995 marked this conversation as resolved.
private static final ConcurrentHashMap<BigQueryConnection, InheritableThreadLocal<String>>
instanceLocals = new ConcurrentHashMap<>();
private static final ConcurrentHashMap<BigQueryConnection, String> instanceIds =
new ConcurrentHashMap<>();
Comment thread
Neenu1995 marked this conversation as resolved.

/** Allocates an exclusive InheritableThreadLocal and registers the connection mapping. */
private static final InheritableThreadLocal<String> 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());
Comment thread
Neenu1995 marked this conversation as resolved.
return "JdbcConnection-" + suffix;
});

currentConnectionId.set(cleanId);
InheritableThreadLocal<String> threadLocal =
instanceLocals.computeIfAbsent(connection, k -> new InheritableThreadLocal<>());
Comment thread
Neenu1995 marked this conversation as resolved.
threadLocal.set(cleanId);
}
}

/** Retrieves the connection ID mapped to a specific BigQueryConnection instance. */
public static String getConnectionId(BigQueryConnection connection) {
Comment thread
Neenu1995 marked this conversation as resolved.
Outdated
if (connection != null) {
InheritableThreadLocal<String> 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<String> local = instanceLocals.remove(connection);
if (local != null) {
local.remove();
}
instanceIds.remove(connection);
}
}
Comment thread
Neenu1995 marked this conversation as resolved.

public static void clear() {
Comment thread
keshavdandeva marked this conversation as resolved.
currentConnectionId.remove();
for (InheritableThreadLocal<String> local : instanceLocals.values()) {
local.remove();
}
}
Comment thread
Neenu1995 marked this conversation as resolved.
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -81,7 +83,9 @@ Optional<Thread> 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)
Expand Down
Original file line number Diff line number Diff line change
@@ -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));
}
}
Loading