From 45563b94a15057073e69b9d2eaaa070e3a07d3bc Mon Sep 17 00:00:00 2001 From: Patrick Schlindwein Date: Sat, 20 Dec 2025 13:30:06 +0100 Subject: [PATCH] ClickHouse extract queryId attribute --- .../common/ClickHouseAttributesExtractor.java | 37 +++++++++++++++++++ .../common/ClickHouseDbRequest.java | 11 +++++- .../common/ClickHouseInstrumenterFactory.java | 1 + .../ClickHouseClientV1Instrumentation.java | 1 + .../clientv1/v0_5/ClickHouseClientV1Test.java | 14 ++++++- .../ClickHouseClientV2Instrumentation.java | 7 +++- .../clientv2/v0_8/ClickHouseClientV2Test.java | 14 ++++++- 7 files changed, 78 insertions(+), 7 deletions(-) create mode 100644 instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseAttributesExtractor.java diff --git a/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseAttributesExtractor.java b/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseAttributesExtractor.java new file mode 100644 index 000000000000..8c014a81f1b8 --- /dev/null +++ b/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseAttributesExtractor.java @@ -0,0 +1,37 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.clickhouse.common; + +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.AttributesBuilder; +import io.opentelemetry.context.Context; +import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor; +import org.jspecify.annotations.Nullable; + +class ClickHouseAttributesExtractor implements AttributesExtractor { + + private static final AttributeKey QUERY_ID = + AttributeKey.stringKey("clickhouse.query_id"); + + @Override + public void onStart( + AttributesBuilder attributes, + Context parentContext, + ClickHouseDbRequest clickHouseDbRequest) { + String queryId = clickHouseDbRequest.getQueryId(); + if (queryId != null) { + attributes.put(QUERY_ID, clickHouseDbRequest.getQueryId()); + } + } + + @Override + public void onEnd( + AttributesBuilder attributes, + Context context, + ClickHouseDbRequest clickHouseDbRequest, + @Nullable Void unused, + @Nullable Throwable error) {} +} diff --git a/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseDbRequest.java b/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseDbRequest.java index 35753981cb9b..6c1b0d54ded2 100644 --- a/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseDbRequest.java +++ b/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseDbRequest.java @@ -18,8 +18,12 @@ public abstract class ClickHouseDbRequest { SqlStatementSanitizer.create(AgentCommonConfig.get().isStatementSanitizationEnabled()); public static ClickHouseDbRequest create( - @Nullable String host, @Nullable Integer port, @Nullable String dbName, String sql) { - return new AutoValue_ClickHouseDbRequest(host, port, dbName, sanitizer.sanitize(sql)); + @Nullable String host, + @Nullable Integer port, + @Nullable String dbName, + @Nullable String queryId, + String sql) { + return new AutoValue_ClickHouseDbRequest(host, port, dbName, queryId, sanitizer.sanitize(sql)); } @Nullable @@ -31,5 +35,8 @@ public static ClickHouseDbRequest create( @Nullable public abstract String getDbName(); + @Nullable + public abstract String getQueryId(); + public abstract SqlStatementInfo getSqlStatementInfo(); } diff --git a/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseInstrumenterFactory.java b/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseInstrumenterFactory.java index 7de945c17102..01d0e037a6dc 100644 --- a/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseInstrumenterFactory.java +++ b/instrumentation/clickhouse/clickhouse-client-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/common/ClickHouseInstrumenterFactory.java @@ -28,6 +28,7 @@ public static Instrumenter createInstrumenter( .addAttributesExtractor(DbClientAttributesExtractor.create(dbAttributesGetter)) .addAttributesExtractor( ServerAttributesExtractor.create(new ClickHouseNetworkAttributesGetter())) + .addAttributesExtractor(new ClickHouseAttributesExtractor()) .addOperationMetrics(DbClientMetrics.get()) .buildInstrumenter(SpanKindExtractor.alwaysClient()); } diff --git a/instrumentation/clickhouse/clickhouse-client-v1-0.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv1/v0_5/ClickHouseClientV1Instrumentation.java b/instrumentation/clickhouse/clickhouse-client-v1-0.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv1/v0_5/ClickHouseClientV1Instrumentation.java index 78a3c697f186..18db359efa17 100644 --- a/instrumentation/clickhouse/clickhouse-client-v1-0.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv1/v0_5/ClickHouseClientV1Instrumentation.java +++ b/instrumentation/clickhouse/clickhouse-client-v1-0.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv1/v0_5/ClickHouseClientV1Instrumentation.java @@ -60,6 +60,7 @@ public static ClickHouseScope onEnter( .getServer() .getDatabase() .orElse(ClickHouseDefaults.DATABASE.getDefaultValue().toString()), + clickHouseRequest.getQueryId().orElse(null), ClickHouseRequestAccess.getQuery(clickHouseRequest)); return ClickHouseScope.start(instrumenter(), currentContext(), request); diff --git a/instrumentation/clickhouse/clickhouse-client-v1-0.5/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv1/v0_5/ClickHouseClientV1Test.java b/instrumentation/clickhouse/clickhouse-client-v1-0.5/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv1/v0_5/ClickHouseClientV1Test.java index 53920b0b64f8..8814a68249de 100644 --- a/instrumentation/clickhouse/clickhouse-client-v1-0.5/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv1/v0_5/ClickHouseClientV1Test.java +++ b/instrumentation/clickhouse/clickhouse-client-v1-0.5/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv1/v0_5/ClickHouseClientV1Test.java @@ -20,6 +20,7 @@ import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_STATEMENT; import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_SYSTEM; import static java.util.Arrays.asList; +import static java.util.stream.Collectors.toList; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.catchThrowable; @@ -32,6 +33,7 @@ import com.clickhouse.client.ClickHouseResponseSummary; import com.clickhouse.data.ClickHouseFormat; import com.google.common.collect.ImmutableMap; +import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.instrumentation.api.internal.SemconvStability; @@ -44,6 +46,7 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.stream.Stream; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -65,6 +68,9 @@ class ClickHouseClientV1Test { private static ClickHouseNode server; private static ClickHouseClient client; + private static final AttributeKey queryIdKey = + AttributeKey.stringKey("clickhouse.query_id"); + @BeforeAll static void setup() throws ClickHouseException { clickhouseServer.start(); @@ -182,6 +188,11 @@ void testExecuteAndWaitWithStringQueryAndId() throws ClickHouseException { response.close(); }); + List attributeAssertions = + Stream.concat( + attributeAssertions("select * from " + tableName, "SELECT").stream(), + Stream.of(equalTo(queryIdKey, "test_query_id"))) + .collect(toList()); testing.waitAndAssertTraces( trace -> trace.hasSpansSatisfyingExactly( @@ -190,8 +201,7 @@ void testExecuteAndWaitWithStringQueryAndId() throws ClickHouseException { span.hasName("SELECT " + dbName) .hasKind(SpanKind.CLIENT) .hasParent(trace.getSpan(0)) - .hasAttributesSatisfyingExactly( - attributeAssertions("select * from " + tableName, "SELECT")))); + .hasAttributesSatisfyingExactly(attributeAssertions))); } @Test diff --git a/instrumentation/clickhouse/clickhouse-client-v2-0.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv2/v0_8/ClickHouseClientV2Instrumentation.java b/instrumentation/clickhouse/clickhouse-client-v2-0.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv2/v0_8/ClickHouseClientV2Instrumentation.java index 3255589460fe..85ea9188f20f 100644 --- a/instrumentation/clickhouse/clickhouse-client-v2-0.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv2/v0_8/ClickHouseClientV2Instrumentation.java +++ b/instrumentation/clickhouse/clickhouse-client-v2-0.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv2/v0_8/ClickHouseClientV2Instrumentation.java @@ -68,10 +68,15 @@ public static ClickHouseScope onEnter( } String database = client.getConfiguration().get("database"); + String queryId = null; + if (querySettings != null) { + queryId = querySettings.getQueryId(); + } + Context parentContext = currentContext(); ClickHouseDbRequest request = ClickHouseDbRequest.create( - addressAndPort.getAddress(), addressAndPort.getPort(), database, sqlQuery); + addressAndPort.getAddress(), addressAndPort.getPort(), database, queryId, sqlQuery); return ClickHouseScope.start(instrumenter(), parentContext, request); } diff --git a/instrumentation/clickhouse/clickhouse-client-v2-0.8/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv2/v0_8/ClickHouseClientV2Test.java b/instrumentation/clickhouse/clickhouse-client-v2-0.8/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv2/v0_8/ClickHouseClientV2Test.java index 93b396d3ccb8..779f0727d71e 100644 --- a/instrumentation/clickhouse/clickhouse-client-v2-0.8/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv2/v0_8/ClickHouseClientV2Test.java +++ b/instrumentation/clickhouse/clickhouse-client-v2-0.8/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/clickhouse/clientv2/v0_8/ClickHouseClientV2Test.java @@ -20,6 +20,7 @@ import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_STATEMENT; import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_SYSTEM; import static java.util.Arrays.asList; +import static java.util.stream.Collectors.toList; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.catchThrowable; @@ -31,6 +32,7 @@ import com.clickhouse.client.api.query.QueryResponse; import com.clickhouse.client.api.query.QuerySettings; import com.clickhouse.client.api.query.Records; +import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.instrumentation.api.internal.SemconvStability; @@ -45,6 +47,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.stream.Stream; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -67,6 +70,9 @@ class ClickHouseClientV2Test { private static final String username = "default"; private static final String password = ""; + private static final AttributeKey queryIdKey = + AttributeKey.stringKey("clickhouse.query_id"); + @BeforeAll static void setup() throws Exception { clickhouseServer.start(); @@ -178,6 +184,11 @@ void testQueryWithStringQueryAndId() throws Exception { response.close(); }); + List attributeAssertions = + Stream.concat( + attributeAssertions("select * from " + tableName, "SELECT").stream(), + Stream.of(equalTo(queryIdKey, "test_query_id"))) + .collect(toList()); testing.waitAndAssertTraces( trace -> trace.hasSpansSatisfyingExactly( @@ -186,8 +197,7 @@ void testQueryWithStringQueryAndId() throws Exception { span.hasName("SELECT " + dbName) .hasKind(SpanKind.CLIENT) .hasParent(trace.getSpan(0)) - .hasAttributesSatisfyingExactly( - attributeAssertions("select * from " + tableName, "SELECT")))); + .hasAttributesSatisfyingExactly(attributeAssertions))); } @Test