Skip to content

Commit be24264

Browse files
lauritlaurit
committed
Instrument jdbc batch queries
1 parent 057ba16 commit be24264

File tree

16 files changed

+998
-85
lines changed

16 files changed

+998
-85
lines changed

‎instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/db/DbClientSpanNameExtractor.java‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,9 @@
55

66
package io.opentelemetry.instrumentation.api.incubator.semconv.db;
77

8+
import io.opentelemetry.instrumentation.api.incubator.semconv.db.internal.MultiQuerySqlClientAttributesGetter;
89
import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
10+
import io.opentelemetry.instrumentation.api.internal.SemconvStability;
911

1012
public abstract class DbClientSpanNameExtractor<REQUEST> implements SpanNameExtractor<REQUEST> {
1113

@@ -93,10 +95,39 @@ private SqlClientSpanNameExtractor(SqlClientAttributesGetter<REQUEST> getter) {
9395

9496
@Override
9597
public String extract(REQUEST request) {
98+
boolean isMultiQuery = false;
99+
MultiQuerySqlClientAttributesGetter<REQUEST> multiGetter = null;
100+
if (getter instanceof MultiQuerySqlClientAttributesGetter) {
101+
multiGetter = (MultiQuerySqlClientAttributesGetter<REQUEST>) getter;
102+
isMultiQuery = multiGetter.getRawQueryTexts(request).size() > 1;
103+
}
104+
105+
Long batchSize = getter.getBatchSize(request);
106+
boolean isBatch = batchSize != null && batchSize > 1;
107+
96108
String namespace = getter.getDbNamespace(request);
97-
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(getter.getRawQueryText(request));
98-
return computeSpanName(
99-
namespace, sanitizedStatement.getOperation(), sanitizedStatement.getMainIdentifier());
109+
if (!isBatch || (!SemconvStability.emitStableDatabaseSemconv() && !isMultiQuery)) {
110+
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(getter.getRawQueryText(request));
111+
return computeSpanName(
112+
namespace, sanitizedStatement.getOperation(), sanitizedStatement.getMainIdentifier());
113+
} else if (SemconvStability.emitStableDatabaseSemconv()) {
114+
if (!isMultiQuery) { // batch query with single unique query
115+
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(getter.getRawQueryText(request));
116+
return computeSpanName(
117+
namespace,
118+
"BATCH " + sanitizedStatement.getOperation(),
119+
sanitizedStatement.getMainIdentifier());
120+
} else { // batch query with multiple unique queries
121+
MultiQuery multiQuery = MultiQuery.analyze(multiGetter.getRawQueryTexts(request), false);
122+
123+
return computeSpanName(
124+
namespace,
125+
multiQuery.getOperation() != null ? "BATCH " + multiQuery.getOperation() : "BATCH",
126+
multiQuery.getMainIdentifier());
127+
}
128+
} else {
129+
return computeSpanName(namespace, null, null);
130+
}
100131
}
101132
}
102133
}

‎instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/db/MultiQuery.java‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.api.incubator.semconv.db;
7+
8+
import java.util.Collection;
9+
import java.util.LinkedHashSet;
10+
import java.util.Set;
11+
12+
class MultiQuery {
13+
private static final SqlStatementSanitizer sanitizer = SqlStatementSanitizer.create(true);
14+
15+
private final String mainIdentifier;
16+
private final String operation;
17+
private final Set<String> statements;
18+
19+
private MultiQuery(String mainIdentifier, String operation, Set<String> statements) {
20+
this.mainIdentifier = mainIdentifier;
21+
this.operation = operation;
22+
this.statements = statements;
23+
}
24+
25+
static MultiQuery analyze(
26+
Collection<String> rawQueryTexts, boolean statementSanitizationEnabled) {
27+
UniqueValue uniqueMainIdentifier = new UniqueValue();
28+
UniqueValue uniqueOperation = new UniqueValue();
29+
Set<String> uniqueStatements = new LinkedHashSet<>();
30+
for (String rawQueryText : rawQueryTexts) {
31+
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(rawQueryText);
32+
String mainIdentifier = sanitizedStatement.getMainIdentifier();
33+
uniqueMainIdentifier.set(mainIdentifier);
34+
String operation = sanitizedStatement.getOperation();
35+
uniqueOperation.set(operation);
36+
uniqueStatements.add(
37+
statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
38+
}
39+
40+
return new MultiQuery(
41+
uniqueMainIdentifier.getValue(), uniqueOperation.getValue(), uniqueStatements);
42+
}
43+
44+
public String getMainIdentifier() {
45+
return mainIdentifier;
46+
}
47+
48+
public String getOperation() {
49+
return operation;
50+
}
51+
52+
public Set<String> getStatements() {
53+
return statements;
54+
}
55+
56+
private static class UniqueValue {
57+
private String value;
58+
private boolean valid = true;
59+
60+
void set(String value) {
61+
if (!valid) {
62+
return;
63+
}
64+
if (this.value == null) {
65+
this.value = value;
66+
} else if (!this.value.equals(value)) {
67+
valid = false;
68+
}
69+
}
70+
71+
String getValue() {
72+
return valid ? value : null;
73+
}
74+
}
75+
}

‎instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/db/SqlClientAttributesExtractor.java‎

Lines changed: 67 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,10 @@
1010
import io.opentelemetry.api.common.AttributeKey;
1111
import io.opentelemetry.api.common.AttributesBuilder;
1212
import io.opentelemetry.context.Context;
13+
import io.opentelemetry.instrumentation.api.incubator.semconv.db.internal.MultiQuerySqlClientAttributesGetter;
1314
import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
1415
import io.opentelemetry.instrumentation.api.internal.SemconvStability;
16+
import java.util.Collection;
1517

1618
/**
1719
* Extractor of <a
@@ -35,6 +37,8 @@ public final class SqlClientAttributesExtractor<REQUEST, RESPONSE>
3537
private static final AttributeKey<String> DB_QUERY_TEXT = AttributeKey.stringKey("db.query.text");
3638
private static final AttributeKey<String> DB_COLLECTION_NAME =
3739
AttributeKey.stringKey("db.collection.name");
40+
private static final AttributeKey<Long> DB_OPERATION_BATCH_SIZE =
41+
AttributeKey.longKey("db.operation.batch.size");
3842

3943
/** Creates the SQL client attributes extractor with default configuration. */
4044
public static <REQUEST, RESPONSE> AttributesExtractor<REQUEST, RESPONSE> create(
@@ -52,7 +56,7 @@ public static <REQUEST, RESPONSE> SqlClientAttributesExtractorBuilder<REQUEST, R
5256
}
5357

5458
private static final String SQL_CALL = "CALL";
55-
// sanitizer is also used to extract operation and table name, so we have it always enable here
59+
// sanitizer is also used to extract operation and table name, so we have it always enabled here
5660
private static final SqlStatementSanitizer sanitizer = SqlStatementSanitizer.create(true);
5761

5862
private final AttributeKey<String> oldSemconvTableAttribute;
@@ -71,30 +75,73 @@ public static <REQUEST, RESPONSE> SqlClientAttributesExtractorBuilder<REQUEST, R
7175
public void onStart(AttributesBuilder attributes, Context parentContext, REQUEST request) {
7276
super.onStart(attributes, parentContext, request);
7377

74-
String rawQueryText = getter.getRawQueryText(request);
75-
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(rawQueryText);
76-
String operation = sanitizedStatement.getOperation();
77-
if (SemconvStability.emitStableDatabaseSemconv()) {
78-
internalSet(
79-
attributes,
80-
DB_QUERY_TEXT,
81-
statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
82-
internalSet(attributes, DB_OPERATION_NAME, operation);
83-
}
84-
if (SemconvStability.emitOldDatabaseSemconv()) {
85-
internalSet(
86-
attributes,
87-
DB_STATEMENT,
88-
statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
89-
internalSet(attributes, DB_OPERATION, operation);
78+
boolean isMultiQuery = false;
79+
MultiQuerySqlClientAttributesGetter<REQUEST> multiGetter = null;
80+
if (getter instanceof MultiQuerySqlClientAttributesGetter) {
81+
multiGetter = (MultiQuerySqlClientAttributesGetter<REQUEST>) getter;
82+
isMultiQuery = multiGetter.getRawQueryTexts(request).size() > 1;
9083
}
91-
if (!SQL_CALL.equals(operation)) {
84+
85+
Long batchSize = getter.getBatchSize(request);
86+
boolean isBatch = batchSize != null && batchSize > 1;
87+
88+
if (!isMultiQuery) {
89+
String rawQueryText = getter.getRawQueryText(request);
90+
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(rawQueryText);
91+
String operation = sanitizedStatement.getOperation();
9292
if (SemconvStability.emitStableDatabaseSemconv()) {
93-
internalSet(attributes, DB_COLLECTION_NAME, sanitizedStatement.getMainIdentifier());
93+
internalSet(
94+
attributes,
95+
DB_QUERY_TEXT,
96+
statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
97+
if (operation != null) {
98+
internalSet(attributes, DB_OPERATION_NAME, (isBatch ? "BATCH " : "") + operation);
99+
}
94100
}
95101
if (SemconvStability.emitOldDatabaseSemconv()) {
96-
internalSet(attributes, oldSemconvTableAttribute, sanitizedStatement.getMainIdentifier());
102+
internalSet(
103+
attributes,
104+
DB_STATEMENT,
105+
statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
106+
internalSet(attributes, DB_OPERATION, operation);
107+
}
108+
if (!SQL_CALL.equals(operation)) {
109+
if (SemconvStability.emitStableDatabaseSemconv()) {
110+
internalSet(attributes, DB_COLLECTION_NAME, sanitizedStatement.getMainIdentifier());
111+
}
112+
if (SemconvStability.emitOldDatabaseSemconv()) {
113+
internalSet(attributes, oldSemconvTableAttribute, sanitizedStatement.getMainIdentifier());
114+
}
115+
}
116+
if (SemconvStability.emitStableDatabaseSemconv() && isBatch) {
117+
internalSet(attributes, DB_OPERATION_BATCH_SIZE, batchSize);
118+
}
119+
} else if (SemconvStability.emitStableDatabaseSemconv()) {
120+
MultiQuery multiQuery =
121+
MultiQuery.analyze(multiGetter.getRawQueryTexts(request), statementSanitizationEnabled);
122+
123+
internalSet(attributes, DB_QUERY_TEXT, join(";", multiQuery.getStatements()));
124+
String operation =
125+
multiQuery.getOperation() != null ? "BATCH " + multiQuery.getOperation() : "BATCH";
126+
internalSet(attributes, DB_OPERATION_NAME, operation);
127+
128+
if (multiQuery.getMainIdentifier() != null
129+
&& (multiQuery.getOperation() == null || !SQL_CALL.equals(multiQuery.getOperation()))) {
130+
internalSet(attributes, DB_COLLECTION_NAME, multiQuery.getMainIdentifier());
131+
}
132+
internalSet(attributes, DB_OPERATION_BATCH_SIZE, batchSize);
133+
}
134+
}
135+
136+
// String.join is not available on android
137+
private static String join(String delimiter, Collection<String> collection) {
138+
StringBuilder builder = new StringBuilder();
139+
for (String string : collection) {
140+
if (builder.length() != 0) {
141+
builder.append(delimiter);
97142
}
143+
builder.append(string);
98144
}
145+
return builder.toString();
99146
}
100147
}

‎instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/db/SqlClientAttributesGetter.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,4 +38,9 @@ default String getRawStatement(REQUEST request) {
3838
default String getRawQueryText(REQUEST request) {
3939
return getRawStatement(request);
4040
}
41+
42+
// TODO: make this required to implement
43+
default Long getBatchSize(REQUEST request) {
44+
return null;
45+
}
4146
}

‎instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/db/internal/MultiQuerySqlClientAttributesGetter.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.api.incubator.semconv.db.internal;
7+
8+
import io.opentelemetry.instrumentation.api.incubator.semconv.db.SqlClientAttributesExtractor;
9+
import io.opentelemetry.instrumentation.api.incubator.semconv.db.SqlClientAttributesGetter;
10+
import java.util.Collection;
11+
12+
/**
13+
* An extended version of {@link SqlClientAttributesGetter} for getting SQL database client
14+
* attributes for operations that run multiple distinct queries.
15+
*
16+
* <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
17+
* at any time.
18+
*/
19+
public interface MultiQuerySqlClientAttributesGetter<REQUEST>
20+
extends SqlClientAttributesGetter<REQUEST> {
21+
22+
/**
23+
* Get the raw SQL statements. The value returned by this method is later sanitized by the {@link
24+
* SqlClientAttributesExtractor} before being set as span attribute.
25+
*/
26+
Collection<String> getRawQueryTexts(REQUEST request);
27+
}

‎instrumentation-api-incubator/src/test/java/io/opentelemetry/instrumentation/api/incubator/semconv/db/DbClientSpanNameExtractorTest.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,11 @@
88
import static org.junit.jupiter.api.Assertions.assertEquals;
99
import static org.mockito.Mockito.when;
1010

11+
import io.opentelemetry.instrumentation.api.incubator.semconv.db.internal.MultiQuerySqlClientAttributesGetter;
1112
import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
13+
import io.opentelemetry.instrumentation.api.internal.SemconvStability;
14+
import java.util.Arrays;
15+
import java.util.Collections;
1216
import org.junit.jupiter.api.Test;
1317
import org.junit.jupiter.api.extension.ExtendWith;
1418
import org.mockito.Mock;
@@ -18,6 +22,7 @@
1822
class DbClientSpanNameExtractorTest {
1923
@Mock DbClientAttributesGetter<DbRequest> dbAttributesGetter;
2024
@Mock SqlClientAttributesGetter<DbRequest> sqlAttributesGetter;
25+
@Mock MultiQuerySqlClientAttributesGetter<DbRequest> multiQuerySqlClientAttributesGetter;
2126

2227
@Test
2328
void shouldExtractFullSpanName() {
@@ -132,5 +137,53 @@ void shouldFallBackToDefaultSpanName() {
132137
assertEquals("DB Query", spanName);
133138
}
134139

140+
@Test
141+
void shouldExtractFullSpanNameForBatch() {
142+
// given
143+
DbRequest dbRequest = new DbRequest();
144+
145+
when(multiQuerySqlClientAttributesGetter.getRawQueryTexts(dbRequest))
146+
.thenReturn(Arrays.asList("INSERT INTO table VALUES(1)", "INSERT INTO table VALUES(2)"));
147+
when(multiQuerySqlClientAttributesGetter.getDbNamespace(dbRequest)).thenReturn("database");
148+
when(multiQuerySqlClientAttributesGetter.getBatchSize(dbRequest)).thenReturn(2L);
149+
150+
SpanNameExtractor<DbRequest> underTest =
151+
DbClientSpanNameExtractor.create(multiQuerySqlClientAttributesGetter);
152+
153+
// when
154+
String spanName = underTest.extract(dbRequest);
155+
156+
// then
157+
assertEquals(
158+
SemconvStability.emitStableDatabaseSemconv() ? "BATCH INSERT database.table" : "database",
159+
spanName);
160+
}
161+
162+
@Test
163+
void shouldExtractFullSpanNameForSingleQueryBatch() {
164+
// given
165+
DbRequest dbRequest = new DbRequest();
166+
167+
when(multiQuerySqlClientAttributesGetter.getRawQueryTexts(dbRequest))
168+
.thenReturn(Collections.singletonList("INSERT INTO table VALUES(?)"));
169+
if (SemconvStability.emitStableDatabaseSemconv()) {
170+
when(multiQuerySqlClientAttributesGetter.getRawQueryText(dbRequest))
171+
.thenReturn("INSERT INTO table VALUES(?)");
172+
}
173+
when(multiQuerySqlClientAttributesGetter.getDbNamespace(dbRequest)).thenReturn("database");
174+
when(multiQuerySqlClientAttributesGetter.getBatchSize(dbRequest)).thenReturn(2L);
175+
176+
SpanNameExtractor<DbRequest> underTest =
177+
DbClientSpanNameExtractor.create(multiQuerySqlClientAttributesGetter);
178+
179+
// when
180+
String spanName = underTest.extract(dbRequest);
181+
182+
// then
183+
assertEquals(
184+
SemconvStability.emitStableDatabaseSemconv() ? "BATCH INSERT database.table" : "database",
185+
spanName);
186+
}
187+
135188
static class DbRequest {}
136189
}

0 commit comments

Comments
 (0)