Instrument jdbc batch queries

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

@@ -5,7 +5,9 @@
 
 package io.opentelemetry.instrumentation.api.incubator.semconv.db;
 
+import io.opentelemetry.instrumentation.api.incubator.semconv.db.internal.MultiQuerySqlClientAttributesGetter;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import io.opentelemetry.instrumentation.api.internal.SemconvStability;
 
 public abstract class DbClientSpanNameExtractor<REQUEST> implements SpanNameExtractor<REQUEST> {
 
@@ -93,10 +95,39 @@ private SqlClientSpanNameExtractor(SqlClientAttributesGetter<REQUEST> getter) {
 
     @Override
     public String extract(REQUEST request) {
+      boolean isMultiQuery = false;
+      MultiQuerySqlClientAttributesGetter<REQUEST> multiGetter = null;
+      if (getter instanceof MultiQuerySqlClientAttributesGetter) {
+        multiGetter = (MultiQuerySqlClientAttributesGetter<REQUEST>) getter;
+        isMultiQuery = multiGetter.getRawQueryTexts(request).size() > 1;
+      }
+
+      Long batchSize = getter.getBatchSize(request);
+      boolean isBatch = batchSize != null && batchSize > 1;
+
       String namespace = getter.getDbNamespace(request);
-      SqlStatementInfo sanitizedStatement = sanitizer.sanitize(getter.getRawQueryText(request));
-      return computeSpanName(
-          namespace, sanitizedStatement.getOperation(), sanitizedStatement.getMainIdentifier());
+      if (!isBatch || (!SemconvStability.emitStableDatabaseSemconv() && !isMultiQuery)) {
+        SqlStatementInfo sanitizedStatement = sanitizer.sanitize(getter.getRawQueryText(request));
+        return computeSpanName(
+            namespace, sanitizedStatement.getOperation(), sanitizedStatement.getMainIdentifier());
+      } else if (SemconvStability.emitStableDatabaseSemconv()) {
+        if (!isMultiQuery) { // batch query with single unique query
+          SqlStatementInfo sanitizedStatement = sanitizer.sanitize(getter.getRawQueryText(request));
+          return computeSpanName(
+              namespace,
+              "BATCH " + sanitizedStatement.getOperation(),
+              sanitizedStatement.getMainIdentifier());
+        } else { // batch query with multiple unique queries
+          MultiQuery multiQuery = MultiQuery.analyze(multiGetter.getRawQueryTexts(request), false);
+
+          return computeSpanName(
+              namespace,
+              multiQuery.getOperation() != null ? "BATCH " + multiQuery.getOperation() : "BATCH",
+              multiQuery.getMainIdentifier());
+        }
+      } else {
+        return computeSpanName(namespace, null, null);
+      }
     }
   }
 }

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

@@ -0,0 +1,75 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.incubator.semconv.db;
+
+import java.util.Collection;
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+class MultiQuery {
+  private static final SqlStatementSanitizer sanitizer = SqlStatementSanitizer.create(true);
+
+  private final String mainIdentifier;
+  private final String operation;
+  private final Set<String> statements;
+
+  private MultiQuery(String mainIdentifier, String operation, Set<String> statements) {
+    this.mainIdentifier = mainIdentifier;
+    this.operation = operation;
+    this.statements = statements;
+  }
+
+  static MultiQuery analyze(
+      Collection<String> rawQueryTexts, boolean statementSanitizationEnabled) {
+    UniqueValue uniqueMainIdentifier = new UniqueValue();
+    UniqueValue uniqueOperation = new UniqueValue();
+    Set<String> uniqueStatements = new LinkedHashSet<>();
+    for (String rawQueryText : rawQueryTexts) {
+      SqlStatementInfo sanitizedStatement = sanitizer.sanitize(rawQueryText);
+      String mainIdentifier = sanitizedStatement.getMainIdentifier();
+      uniqueMainIdentifier.set(mainIdentifier);
+      String operation = sanitizedStatement.getOperation();
+      uniqueOperation.set(operation);
+      uniqueStatements.add(
+          statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
+    }
+
+    return new MultiQuery(
+        uniqueMainIdentifier.getValue(), uniqueOperation.getValue(), uniqueStatements);
+  }
+
+  public String getMainIdentifier() {
+    return mainIdentifier;
+  }
+
+  public String getOperation() {
+    return operation;
+  }
+
+  public Set<String> getStatements() {
+    return statements;
+  }
+
+  private static class UniqueValue {
+    private String value;
+    private boolean valid = true;
+
+    void set(String value) {
+      if (!valid) {
+        return;
+      }
+      if (this.value == null) {
+        this.value = value;
+      } else if (!this.value.equals(value)) {
+        valid = false;
+      }
+    }
+
+    String getValue() {
+      return valid ? value : null;
+    }
+  }
+}

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

@@ -10,8 +10,10 @@
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.incubator.semconv.db.internal.MultiQuerySqlClientAttributesGetter;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
 import io.opentelemetry.instrumentation.api.internal.SemconvStability;
+import java.util.Collection;
 
 /**
  * Extractor of <a
@@ -35,6 +37,8 @@ public final class SqlClientAttributesExtractor<REQUEST, RESPONSE>
   private static final AttributeKey<String> DB_QUERY_TEXT = AttributeKey.stringKey("db.query.text");
   static final AttributeKey<String> DB_COLLECTION_NAME =
       AttributeKey.stringKey("db.collection.name");
+  private static final AttributeKey<Long> DB_OPERATION_BATCH_SIZE =
+      AttributeKey.longKey("db.operation.batch.size");
 
   /** Creates the SQL client attributes extractor with default configuration. */
   public static <REQUEST, RESPONSE> AttributesExtractor<REQUEST, RESPONSE> create(
@@ -52,7 +56,7 @@ public static <REQUEST, RESPONSE> SqlClientAttributesExtractorBuilder<REQUEST, R
   }
 
   private static final String SQL_CALL = "CALL";
-  // sanitizer is also used to extract operation and table name, so we have it always enable here
+  // sanitizer is also used to extract operation and table name, so we have it always enabled here
   private static final SqlStatementSanitizer sanitizer = SqlStatementSanitizer.create(true);
 
   private final AttributeKey<String> oldSemconvTableAttribute;
@@ -71,30 +75,73 @@ public static <REQUEST, RESPONSE> SqlClientAttributesExtractorBuilder<REQUEST, R
   public void onStart(AttributesBuilder attributes, Context parentContext, REQUEST request) {
     super.onStart(attributes, parentContext, request);
 
-    String rawQueryText = getter.getRawQueryText(request);
-    SqlStatementInfo sanitizedStatement = sanitizer.sanitize(rawQueryText);
-    String operation = sanitizedStatement.getOperation();
-    if (SemconvStability.emitStableDatabaseSemconv()) {
-      internalSet(
-          attributes,
-          DB_QUERY_TEXT,
-          statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
-      internalSet(attributes, DB_OPERATION_NAME, operation);
-    }
-    if (SemconvStability.emitOldDatabaseSemconv()) {
-      internalSet(
-          attributes,
-          DB_STATEMENT,
-          statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
-      internalSet(attributes, DB_OPERATION, operation);
+    boolean isMultiQuery = false;
+    MultiQuerySqlClientAttributesGetter<REQUEST> multiGetter = null;
+    if (getter instanceof MultiQuerySqlClientAttributesGetter) {
+      multiGetter = (MultiQuerySqlClientAttributesGetter<REQUEST>) getter;
+      isMultiQuery = multiGetter.getRawQueryTexts(request).size() > 1;
     }
-    if (!SQL_CALL.equals(operation)) {
+
+    Long batchSize = getter.getBatchSize(request);
+    boolean isBatch = batchSize != null && batchSize > 1;
+
+    if (!isMultiQuery) {
+      String rawQueryText = getter.getRawQueryText(request);
+      SqlStatementInfo sanitizedStatement = sanitizer.sanitize(rawQueryText);
+      String operation = sanitizedStatement.getOperation();
       if (SemconvStability.emitStableDatabaseSemconv()) {
-        internalSet(attributes, DB_COLLECTION_NAME, sanitizedStatement.getMainIdentifier());
+        internalSet(
+            attributes,
+            DB_QUERY_TEXT,
+            statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
+        if (operation != null) {
+          internalSet(attributes, DB_OPERATION_NAME, (isBatch ? "BATCH " : "") + operation);
+        }
       }
       if (SemconvStability.emitOldDatabaseSemconv()) {
-        internalSet(attributes, oldSemconvTableAttribute, sanitizedStatement.getMainIdentifier());
+        internalSet(
+            attributes,
+            DB_STATEMENT,
+            statementSanitizationEnabled ? sanitizedStatement.getFullStatement() : rawQueryText);
+        internalSet(attributes, DB_OPERATION, operation);
+      }
+      if (!SQL_CALL.equals(operation)) {
+        if (SemconvStability.emitStableDatabaseSemconv()) {
+          internalSet(attributes, DB_COLLECTION_NAME, sanitizedStatement.getMainIdentifier());
+        }
+        if (SemconvStability.emitOldDatabaseSemconv()) {
+          internalSet(attributes, oldSemconvTableAttribute, sanitizedStatement.getMainIdentifier());
+        }
+      }
+      if (SemconvStability.emitStableDatabaseSemconv() && isBatch) {
+        internalSet(attributes, DB_OPERATION_BATCH_SIZE, batchSize);
+      }
+    } else if (SemconvStability.emitStableDatabaseSemconv()) {
+      MultiQuery multiQuery =
+          MultiQuery.analyze(multiGetter.getRawQueryTexts(request), statementSanitizationEnabled);
+
+      internalSet(attributes, DB_QUERY_TEXT, join(";", multiQuery.getStatements()));
+      String operation =
+          multiQuery.getOperation() != null ? "BATCH " + multiQuery.getOperation() : "BATCH";
+      internalSet(attributes, DB_OPERATION_NAME, operation);
+
+      if (multiQuery.getMainIdentifier() != null
+          && (multiQuery.getOperation() == null || !SQL_CALL.equals(multiQuery.getOperation()))) {
+        internalSet(attributes, DB_COLLECTION_NAME, multiQuery.getMainIdentifier());
+      }
+      internalSet(attributes, DB_OPERATION_BATCH_SIZE, batchSize);
+    }
+  }
+
+  // String.join is not available on android
+  private static String join(String delimiter, Collection<String> collection) {
+    StringBuilder builder = new StringBuilder();
+    for (String string : collection) {
+      if (builder.length() != 0) {
+        builder.append(delimiter);
       }
+      builder.append(string);
     }
+    return builder.toString();
   }
 }

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

@@ -38,4 +38,9 @@ default String getRawStatement(REQUEST request) {
   default String getRawQueryText(REQUEST request) {
     return getRawStatement(request);
   }
+
+  // TODO: make this required to implement
+  default Long getBatchSize(REQUEST request) {
+    return null;
+  }
 }

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.incubator.semconv.db.internal;
+
+import io.opentelemetry.instrumentation.api.incubator.semconv.db.SqlClientAttributesExtractor;
+import io.opentelemetry.instrumentation.api.incubator.semconv.db.SqlClientAttributesGetter;
+import java.util.Collection;
+
+/**
+ * An extended version of {@link SqlClientAttributesGetter} for getting SQL database client
+ * attributes for operations that run multiple distinct queries.
+ *
+ * <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
+ * at any time.
+ */
+public interface MultiQuerySqlClientAttributesGetter<REQUEST>
+    extends SqlClientAttributesGetter<REQUEST> {
+
+  /**
+   * Get the raw SQL statements. The value returned by this method is later sanitized by the {@link
+   * SqlClientAttributesExtractor} before being set as span attribute.
+   */
+  Collection<String> getRawQueryTexts(REQUEST request);
+}

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

@@ -8,7 +8,11 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.Mockito.when;
 
+import io.opentelemetry.instrumentation.api.incubator.semconv.db.internal.MultiQuerySqlClientAttributesGetter;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import io.opentelemetry.instrumentation.api.internal.SemconvStability;
+import java.util.Arrays;
+import java.util.Collections;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
@@ -18,6 +22,7 @@
 class DbClientSpanNameExtractorTest {
   @Mock DbClientAttributesGetter<DbRequest> dbAttributesGetter;
   @Mock SqlClientAttributesGetter<DbRequest> sqlAttributesGetter;
+  @Mock MultiQuerySqlClientAttributesGetter<DbRequest> multiQuerySqlClientAttributesGetter;
 
   @Test
   void shouldExtractFullSpanName() {
@@ -132,5 +137,53 @@ void shouldFallBackToDefaultSpanName() {
     assertEquals("DB Query", spanName);
   }
 
+  @Test
+  void shouldExtractFullSpanNameForBatch() {
+    // given
+    DbRequest dbRequest = new DbRequest();
+
+    when(multiQuerySqlClientAttributesGetter.getRawQueryTexts(dbRequest))
+        .thenReturn(Arrays.asList("INSERT INTO table VALUES(1)", "INSERT INTO table VALUES(2)"));
+    when(multiQuerySqlClientAttributesGetter.getDbNamespace(dbRequest)).thenReturn("database");
+    when(multiQuerySqlClientAttributesGetter.getBatchSize(dbRequest)).thenReturn(2L);
+
+    SpanNameExtractor<DbRequest> underTest =
+        DbClientSpanNameExtractor.create(multiQuerySqlClientAttributesGetter);
+
+    // when
+    String spanName = underTest.extract(dbRequest);
+
+    // then
+    assertEquals(
+        SemconvStability.emitStableDatabaseSemconv() ? "BATCH INSERT database.table" : "database",
+        spanName);
+  }
+
+  @Test
+  void shouldExtractFullSpanNameForSingleQueryBatch() {
+    // given
+    DbRequest dbRequest = new DbRequest();
+
+    when(multiQuerySqlClientAttributesGetter.getRawQueryTexts(dbRequest))
+        .thenReturn(Collections.singletonList("INSERT INTO table VALUES(?)"));
+    if (SemconvStability.emitStableDatabaseSemconv()) {
+      when(multiQuerySqlClientAttributesGetter.getRawQueryText(dbRequest))
+          .thenReturn("INSERT INTO table VALUES(?)");
+    }
+    when(multiQuerySqlClientAttributesGetter.getDbNamespace(dbRequest)).thenReturn("database");
+    when(multiQuerySqlClientAttributesGetter.getBatchSize(dbRequest)).thenReturn(2L);
+
+    SpanNameExtractor<DbRequest> underTest =
+        DbClientSpanNameExtractor.create(multiQuerySqlClientAttributesGetter);
+
+    // when
+    String spanName = underTest.extract(dbRequest);
+
+    // then
+    assertEquals(
+        SemconvStability.emitStableDatabaseSemconv() ? "BATCH INSERT database.table" : "database",
+        spanName);
+  }
+
   static class DbRequest {}
 }