Skip to content

Commit 467dfc0

Browse files
laurittrask
laurit
and
trask
authored
Instrument jdbc batch queries (#12797)
Co-authored-by: Trask Stalnaker <[email protected]>
1 parent aaed8ac commit 467dfc0

File tree

20 files changed

+1016
-94
lines changed

20 files changed

+1016
-94
lines changed

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

+35-2
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
package io.opentelemetry.instrumentation.api.incubator.semconv.db;
77

88
import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
9+
import io.opentelemetry.instrumentation.api.internal.SemconvStability;
10+
import java.util.Collection;
911

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

@@ -94,9 +96,40 @@ private SqlClientSpanNameExtractor(SqlClientAttributesGetter<REQUEST> getter) {
9496
@Override
9597
public String extract(REQUEST request) {
9698
String namespace = getter.getDbNamespace(request);
97-
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(getter.getRawQueryText(request));
99+
Collection<String> rawQueryTexts = getter.getRawQueryTexts(request);
100+
101+
if (rawQueryTexts.isEmpty()) {
102+
return computeSpanName(namespace, null, null);
103+
}
104+
105+
if (!SemconvStability.emitStableDatabaseSemconv()) {
106+
if (rawQueryTexts.size() > 1) { // for backcompat(?)
107+
return computeSpanName(namespace, null, null);
108+
}
109+
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(rawQueryTexts.iterator().next());
110+
return computeSpanName(
111+
namespace, sanitizedStatement.getOperation(), sanitizedStatement.getMainIdentifier());
112+
}
113+
114+
if (rawQueryTexts.size() == 1) {
115+
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(rawQueryTexts.iterator().next());
116+
String operation = sanitizedStatement.getOperation();
117+
if (isBatch(request)) {
118+
operation = "BATCH " + operation;
119+
}
120+
return computeSpanName(namespace, operation, sanitizedStatement.getMainIdentifier());
121+
}
122+
123+
MultiQuery multiQuery = MultiQuery.analyze(rawQueryTexts, false);
98124
return computeSpanName(
99-
namespace, sanitizedStatement.getOperation(), sanitizedStatement.getMainIdentifier());
125+
namespace,
126+
multiQuery.getOperation() != null ? "BATCH " + multiQuery.getOperation() : "BATCH",
127+
multiQuery.getMainIdentifier());
128+
}
129+
130+
private boolean isBatch(REQUEST request) {
131+
Long batchSize = getter.getBatchSize(request);
132+
return batchSize != null && batchSize > 1;
100133
}
101134
}
102135
}

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

+75
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

+65-20
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import io.opentelemetry.context.Context;
1313
import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
1414
import io.opentelemetry.instrumentation.api.internal.SemconvStability;
15+
import java.util.Collection;
1516

1617
/**
1718
* Extractor of <a
@@ -35,6 +36,8 @@ public final class SqlClientAttributesExtractor<REQUEST, RESPONSE>
3536
private static final AttributeKey<String> DB_QUERY_TEXT = AttributeKey.stringKey("db.query.text");
3637
static final AttributeKey<String> DB_COLLECTION_NAME =
3738
AttributeKey.stringKey("db.collection.name");
39+
private static final AttributeKey<Long> DB_OPERATION_BATCH_SIZE =
40+
AttributeKey.longKey("db.operation.batch.size");
3841

3942
/** Creates the SQL client attributes extractor with default configuration. */
4043
public static <REQUEST, RESPONSE> AttributesExtractor<REQUEST, RESPONSE> create(
@@ -52,7 +55,7 @@ public static <REQUEST, RESPONSE> SqlClientAttributesExtractorBuilder<REQUEST, R
5255
}
5356

5457
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
58+
// sanitizer is also used to extract operation and table name, so we have it always enabled here
5659
private static final SqlStatementSanitizer sanitizer = SqlStatementSanitizer.create(true);
5760

5861
private final AttributeKey<String> oldSemconvTableAttribute;
@@ -71,30 +74,72 @@ public static <REQUEST, RESPONSE> SqlClientAttributesExtractorBuilder<REQUEST, R
7174
public void onStart(AttributesBuilder attributes, Context parentContext, REQUEST request) {
7275
super.onStart(attributes, parentContext, request);
7376

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);
77+
Collection<String> rawQueryTexts = getter.getRawQueryTexts(request);
78+
79+
if (rawQueryTexts.isEmpty()) {
80+
return;
8381
}
82+
8483
if (SemconvStability.emitOldDatabaseSemconv()) {
85-
internalSet(
86-
attributes,
87-
DB_STATEMENT,
88-
statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
89-
internalSet(attributes, DB_OPERATION, operation);
84+
if (rawQueryTexts.size() == 1) { // for backcompat(?)
85+
String rawQueryText = rawQueryTexts.iterator().next();
86+
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(rawQueryText);
87+
String operation = sanitizedStatement.getOperation();
88+
internalSet(
89+
attributes,
90+
DB_STATEMENT,
91+
statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
92+
internalSet(attributes, DB_OPERATION, operation);
93+
if (!SQL_CALL.equals(operation)) {
94+
internalSet(attributes, oldSemconvTableAttribute, sanitizedStatement.getMainIdentifier());
95+
}
96+
}
9097
}
91-
if (!SQL_CALL.equals(operation)) {
92-
if (SemconvStability.emitStableDatabaseSemconv()) {
93-
internalSet(attributes, DB_COLLECTION_NAME, sanitizedStatement.getMainIdentifier());
98+
99+
if (SemconvStability.emitStableDatabaseSemconv()) {
100+
Long batchSize = getter.getBatchSize(request);
101+
boolean isBatch = batchSize != null && batchSize > 1;
102+
if (isBatch) {
103+
internalSet(attributes, DB_OPERATION_BATCH_SIZE, batchSize);
94104
}
95-
if (SemconvStability.emitOldDatabaseSemconv()) {
96-
internalSet(attributes, oldSemconvTableAttribute, sanitizedStatement.getMainIdentifier());
105+
if (rawQueryTexts.size() == 1) {
106+
String rawQueryText = rawQueryTexts.iterator().next();
107+
SqlStatementInfo sanitizedStatement = sanitizer.sanitize(rawQueryText);
108+
String operation = sanitizedStatement.getOperation();
109+
internalSet(
110+
attributes,
111+
DB_QUERY_TEXT,
112+
statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
113+
internalSet(attributes, DB_OPERATION_NAME, isBatch ? "BATCH " + operation : operation);
114+
if (!SQL_CALL.equals(operation)) {
115+
internalSet(attributes, DB_COLLECTION_NAME, sanitizedStatement.getMainIdentifier());
116+
}
117+
} else {
118+
MultiQuery multiQuery =
119+
MultiQuery.analyze(getter.getRawQueryTexts(request), statementSanitizationEnabled);
120+
internalSet(attributes, DB_QUERY_TEXT, join("; ", multiQuery.getStatements()));
121+
122+
String operation =
123+
multiQuery.getOperation() != null ? "BATCH " + multiQuery.getOperation() : "BATCH";
124+
internalSet(attributes, DB_OPERATION_NAME, operation);
125+
126+
if (multiQuery.getMainIdentifier() != null
127+
&& (multiQuery.getOperation() == null || !SQL_CALL.equals(multiQuery.getOperation()))) {
128+
internalSet(attributes, DB_COLLECTION_NAME, multiQuery.getMainIdentifier());
129+
}
130+
}
131+
}
132+
}
133+
134+
// String.join is not available on android
135+
private static String join(String delimiter, Collection<String> collection) {
136+
StringBuilder builder = new StringBuilder();
137+
for (String string : collection) {
138+
if (builder.length() != 0) {
139+
builder.append(delimiter);
97140
}
141+
builder.append(string);
98142
}
143+
return builder.toString();
99144
}
100145
}

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

+29-1
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,10 @@
55

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

8+
import static java.util.Collections.emptySet;
9+
import static java.util.Collections.singleton;
10+
11+
import java.util.Collection;
812
import javax.annotation.Nullable;
913

1014
/**
@@ -33,9 +37,33 @@ default String getRawStatement(REQUEST request) {
3337
return null;
3438
}
3539

36-
// TODO: make this required to implement
40+
/**
41+
* Get the raw SQL query text. The value returned by this method is later sanitized by the {@link
42+
* SqlClientAttributesExtractor} before being set as span attribute.
43+
*
44+
* @deprecated Use {@link #getRawQueryTexts(Object)} instead.
45+
*/
46+
@Deprecated
3747
@Nullable
3848
default String getRawQueryText(REQUEST request) {
3949
return getRawStatement(request);
4050
}
51+
52+
/**
53+
* Get the raw SQL query texts. The values returned by this method is later sanitized by the
54+
* {@link SqlClientAttributesExtractor} before being set as span attribute.
55+
*
56+
* <p>If {@code request} is not a batch query, then this method should return a collection with a
57+
* single element.
58+
*/
59+
// TODO: make this required to implement
60+
default Collection<String> getRawQueryTexts(REQUEST request) {
61+
String rawQueryText = getRawQueryText(request);
62+
return rawQueryText == null ? emptySet() : singleton(rawQueryText);
63+
}
64+
65+
// TODO: make this required to implement
66+
default Long getBatchSize(REQUEST request) {
67+
return null;
68+
}
4169
}

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

+54-3
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,13 @@
55

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

8+
import static java.util.Collections.singleton;
89
import static org.junit.jupiter.api.Assertions.assertEquals;
910
import static org.mockito.Mockito.when;
1011

1112
import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
13+
import io.opentelemetry.instrumentation.api.internal.SemconvStability;
14+
import java.util.Arrays;
1215
import org.junit.jupiter.api.Test;
1316
import org.junit.jupiter.api.extension.ExtendWith;
1417
import org.mockito.Mock;
@@ -24,7 +27,8 @@ void shouldExtractFullSpanName() {
2427
// given
2528
DbRequest dbRequest = new DbRequest();
2629

27-
when(sqlAttributesGetter.getRawQueryText(dbRequest)).thenReturn("SELECT * from table");
30+
when(sqlAttributesGetter.getRawQueryTexts(dbRequest))
31+
.thenReturn(singleton("SELECT * from table"));
2832
when(sqlAttributesGetter.getDbNamespace(dbRequest)).thenReturn("database");
2933

3034
SpanNameExtractor<DbRequest> underTest = DbClientSpanNameExtractor.create(sqlAttributesGetter);
@@ -41,7 +45,8 @@ void shouldSkipDbNameIfTableAlreadyHasDbNamePrefix() {
4145
// given
4246
DbRequest dbRequest = new DbRequest();
4347

44-
when(sqlAttributesGetter.getRawQueryText(dbRequest)).thenReturn("SELECT * from another.table");
48+
when(sqlAttributesGetter.getRawQueryTexts(dbRequest))
49+
.thenReturn(singleton("SELECT * from another.table"));
4550
when(sqlAttributesGetter.getDbNamespace(dbRequest)).thenReturn("database");
4651

4752
SpanNameExtractor<DbRequest> underTest = DbClientSpanNameExtractor.create(sqlAttributesGetter);
@@ -58,7 +63,8 @@ void shouldExtractOperationAndTable() {
5863
// given
5964
DbRequest dbRequest = new DbRequest();
6065

61-
when(sqlAttributesGetter.getRawQueryText(dbRequest)).thenReturn("SELECT * from table");
66+
when(sqlAttributesGetter.getRawQueryTexts(dbRequest))
67+
.thenReturn(singleton("SELECT * from table"));
6268

6369
SpanNameExtractor<DbRequest> underTest = DbClientSpanNameExtractor.create(sqlAttributesGetter);
6470

@@ -132,5 +138,50 @@ void shouldFallBackToDefaultSpanName() {
132138
assertEquals("DB Query", spanName);
133139
}
134140

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

0 commit comments

Comments
 (0)