Unify database batch tests into parameterized scenario tests

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

@@ -94,7 +94,10 @@ static <REQUEST, RESPONSE> void onStartCommon(
       REQUEST request,
       boolean captureQueryParameters) {
     Long batchSize = getter.getDbOperationBatchSize(request);
-    boolean isBatch = batchSize != null && batchSize > 1;
+    // db.operation.batch.size is captured for every batch execution (including an empty batch with
+    // size 0); it is only omitted for a single-statement batch, which is reported as a non-batch
+    boolean emitBatchSize = batchSize != null && batchSize != 1;
+    boolean isBatch = emitBatchSize;
 
     if (emitStableDatabaseSemconv()) {
       attributes.put(
@@ -104,7 +107,7 @@ static <REQUEST, RESPONSE> void onStartCommon(
       attributes.put(DB_QUERY_TEXT, getter.getDbQueryText(request));
       attributes.put(DB_OPERATION_NAME, getter.getDbOperationName(request));
       attributes.put(DB_QUERY_SUMMARY, getter.getDbQuerySummary(request));
-      if (isBatch) {
+      if (emitBatchSize) {
         attributes.put(DB_OPERATION_BATCH_SIZE, batchSize);
       }
     }

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

@@ -193,6 +193,9 @@ public String extract(REQUEST request) {
 
       if (rawQueryTexts.isEmpty()) {
         if (emitStableDatabaseSemconv()) {
+          if (isBatch(request)) {
+            return "BATCH";
+          }
           return computeSpanNameStable(getter, request, null, null, null);
         }
         String dbName = getter.getDbName(request);
@@ -240,7 +243,7 @@ public String extract(REQUEST request) {
 
     private boolean isBatch(REQUEST request) {
       Long batchSize = getter.getDbOperationBatchSize(request);
-      return batchSize != null && batchSize > 1;
+      return batchSize != null && batchSize != 1;
     }
   }
 

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

@@ -15,12 +15,20 @@ class MultiQuery {
   @Nullable private final String storedProcedureName;
   private final Set<String> queryTexts;
   @Nullable private final String querySummary;
+  @Nullable private final String operationName;
+  @Nullable private final String collectionName;
 
   private MultiQuery(
-      @Nullable String storedProcedureName, Set<String> queryTexts, @Nullable String querySummary) {
+      @Nullable String storedProcedureName,
+      Set<String> queryTexts,
+      @Nullable String querySummary,
+      @Nullable String operationName,
+      @Nullable String collectionName) {
     this.storedProcedureName = storedProcedureName;
     this.queryTexts = queryTexts;
     this.querySummary = querySummary;
+    this.operationName = operationName;
+    this.collectionName = collectionName;
   }
 
   static MultiQuery analyzeWithSummary(Collection<String> rawQueryTexts, SqlDialect dialect) {
@@ -47,6 +55,16 @@ public String getQuerySummary() {
     return querySummary;
   }
 
+  @Nullable
+  public String getOperationName() {
+    return operationName;
+  }
+
+  @Nullable
+  public String getCollectionName() {
+    return collectionName;
+  }
+
   public Set<String> getQueryTexts() {
     return queryTexts;
   }
@@ -55,19 +73,28 @@ static class Builder {
     private final UniqueValue uniqueStoredProcedureName = new UniqueValue();
     private final Set<String> uniqueQueryTexts = new LinkedHashSet<>();
     private final UniqueValue uniqueQuerySummary = new UniqueValue();
+    private final UniqueValue uniqueOperationName = new UniqueValue();
+    private final UniqueValue uniqueCollectionName = new UniqueValue();
 
+    @SuppressWarnings(
+        "deprecation") // getOperationName()/getCollectionName() package-private in 3.0
     void add(SqlQuery analyzedQuery, @Nullable String queryText) {
       uniqueStoredProcedureName.set(analyzedQuery.getStoredProcedureName());
       uniqueQueryTexts.add(queryText);
       uniqueQuerySummary.set(analyzedQuery.getQuerySummary());
+      uniqueOperationName.set(analyzedQuery.getOperationName());
+      uniqueCollectionName.set(analyzedQuery.getCollectionName());
     }
 
     MultiQuery build() {
       String querySummary = uniqueQuerySummary.getValue();
+      String operationName = uniqueOperationName.getValue();
       return new MultiQuery(
           uniqueStoredProcedureName.getValue(),
           uniqueQueryTexts,
-          querySummary == null ? "BATCH" : "BATCH " + querySummary);
+          querySummary == null ? "BATCH" : "BATCH " + querySummary,
+          operationName == null ? "BATCH" : "BATCH " + operationName,
+          uniqueCollectionName.getValue());
     }
   }
 

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

@@ -101,7 +101,10 @@ public void onStart(AttributesBuilder attributes, Context parentContext, REQUEST
     SqlDialect dialect = getter.getSqlDialect(request);
 
     Long batchSize = getter.getDbOperationBatchSize(request);
-    boolean isBatch = batchSize != null && batchSize > 1;
+    // db.operation.batch.size is captured for every batch execution (including an empty batch with
+    // size 0); it is only omitted for a single-statement batch, which is reported as a non-batch
+    boolean emitBatchSize = batchSize != null && batchSize != 1;
+    boolean isBatch = emitBatchSize;
 
     if (emitOldDatabaseSemconv()) {
       if (rawQueryTexts.size() == 1) { // for backcompat(?)
@@ -118,7 +121,7 @@ public void onStart(AttributesBuilder attributes, Context parentContext, REQUEST
     }
 
     if (emitStableDatabaseSemconv()) {
-      if (isBatch) {
+      if (emitBatchSize) {
         attributes.put(DB_OPERATION_BATCH_SIZE, batchSize);
       }
       if (rawQueryTexts.size() == 1) {
@@ -149,6 +152,10 @@ public void onStart(AttributesBuilder attributes, Context parentContext, REQUEST
         MultiQuery multiQuery = builder.build();
         attributes.put(DB_QUERY_TEXT, join("; ", multiQuery.getQueryTexts()));
         attributes.put(DB_QUERY_SUMMARY, multiQuery.getQuerySummary());
+        if (singleOperationAndCollection) {
+          attributes.put(DB_OPERATION_NAME, multiQuery.getOperationName());
+          attributes.put(DB_COLLECTION_NAME, multiQuery.getCollectionName());
+        }
         attributes.put(DB_STORED_PROCEDURE_NAME, multiQuery.getStoredProcedureName());
       }
     }

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

@@ -35,6 +35,7 @@ void setUp() {
     lenient()
         .when(sqlAttributesGetter.getSqlDialect(any()))
         .thenReturn(DOUBLE_QUOTES_ARE_STRING_LITERALS);
+    lenient().when(sqlAttributesGetter.getDbOperationBatchSize(any())).thenReturn(null);
   }
 
   @Test
@@ -266,6 +267,30 @@ void shouldExtractFullSpanNameForSingleQueryBatch() {
         .isEqualTo(emitStableDatabaseSemconv() ? "BATCH INSERT table" : "INSERT database.table");
   }
 
+  @Test
+  void shouldExtractFullSpanNameForSingleQueryEmptyBatch() {
+    // given
+    DbRequest dbRequest = new DbRequest();
+
+    when(sqlAttributesGetter.getRawQueryTexts(dbRequest))
+        .thenReturn(singleton("INSERT INTO table VALUES(?)"));
+    if (emitOldDatabaseSemconv() && !emitStableDatabaseSemconv()) {
+      when(sqlAttributesGetter.getDbName(dbRequest)).thenReturn("database");
+    }
+    if (emitStableDatabaseSemconv()) {
+      when(sqlAttributesGetter.getDbOperationBatchSize(dbRequest)).thenReturn(0L);
+    }
+
+    SpanNameExtractor<DbRequest> underTest = DbClientSpanNameExtractor.create(sqlAttributesGetter);
+
+    // when
+    String spanName = underTest.extract(dbRequest);
+
+    // then
+    assertThat(spanName)
+        .isEqualTo(emitStableDatabaseSemconv() ? "BATCH INSERT table" : "INSERT database.table");
+  }
+
   @Test
   void shouldFallBackToNamespaceForEmptySqlQuery() {
     // given
@@ -288,6 +313,28 @@ void shouldFallBackToNamespaceForEmptySqlQuery() {
     assertThat(spanName).isEqualTo("mydb");
   }
 
+  @Test
+  void shouldExtractBatchSpanNameForEmptySqlQueryBatch() {
+    // given
+    DbRequest dbRequest = new DbRequest();
+
+    when(sqlAttributesGetter.getRawQueryTexts(dbRequest)).thenReturn(emptyList());
+    if (emitStableDatabaseSemconv()) {
+      when(sqlAttributesGetter.getDbOperationBatchSize(dbRequest)).thenReturn(0L);
+    }
+    if (emitOldDatabaseSemconv() && !emitStableDatabaseSemconv()) {
+      when(sqlAttributesGetter.getDbName(dbRequest)).thenReturn("mydb");
+    }
+
+    SpanNameExtractor<DbRequest> underTest = DbClientSpanNameExtractor.create(sqlAttributesGetter);
+
+    // when
+    String spanName = underTest.extract(dbRequest);
+
+    // then
+    assertThat(spanName).isEqualTo(emitStableDatabaseSemconv() ? "BATCH" : "mydb");
+  }
+
   @Test
   @SuppressWarnings("deprecation") // testing deprecated method
   void shouldPreserveOldSemconvSpanNameForMigration() {
@@ -335,5 +382,29 @@ void shouldFallBackToNamespaceForEmptySqlQueryInMigration() {
     assertThat(spanName).isEqualTo("mydb");
   }
 
+  @Test
+  @SuppressWarnings("deprecation") // testing deprecated method
+  void shouldExtractBatchSpanNameForEmptySqlQueryBatchInMigration() {
+    // given
+    DbRequest dbRequest = new DbRequest();
+
+    when(sqlAttributesGetter.getRawQueryTexts(dbRequest)).thenReturn(emptyList());
+    if (emitStableDatabaseSemconv()) {
+      when(sqlAttributesGetter.getDbOperationBatchSize(dbRequest)).thenReturn(0L);
+    }
+    if (emitOldDatabaseSemconv() && !emitStableDatabaseSemconv()) {
+      when(sqlAttributesGetter.getDbName(dbRequest)).thenReturn("mydb");
+    }
+
+    SpanNameExtractor<DbRequest> underTest =
+        DbClientSpanNameExtractor.createWithGenericOldSpanName(sqlAttributesGetter);
+
+    // when
+    String spanName = underTest.extract(dbRequest);
+
+    // then
+    assertThat(spanName).isEqualTo(emitStableDatabaseSemconv() ? "BATCH" : "mydb");
+  }
+
   static class DbRequest {}
 }

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

@@ -9,8 +9,10 @@
 import static io.opentelemetry.instrumentation.api.internal.SemconvStability.emitOldDatabaseSemconv;
 import static io.opentelemetry.instrumentation.api.internal.SemconvStability.emitStableDatabaseSemconv;
 import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.assertThat;
+import static io.opentelemetry.semconv.DbAttributes.DB_COLLECTION_NAME;
 import static io.opentelemetry.semconv.DbAttributes.DB_NAMESPACE;
 import static io.opentelemetry.semconv.DbAttributes.DB_OPERATION_BATCH_SIZE;
+import static io.opentelemetry.semconv.DbAttributes.DB_OPERATION_NAME;
 import static io.opentelemetry.semconv.DbAttributes.DB_QUERY_SUMMARY;
 import static io.opentelemetry.semconv.DbAttributes.DB_QUERY_TEXT;
 import static io.opentelemetry.semconv.DbAttributes.DB_SYSTEM_NAME;
@@ -318,6 +320,57 @@ void shouldExtractSingleQueryBatchAttributes() {
     assertThat(endAttributes.build().isEmpty()).isTrue();
   }
 
+  @Test
+  void shouldExtractSingleQueryEmptyBatchAttributes() {
+    // given
+    Map<String, Object> request = new HashMap<>();
+    request.put("db.namespace", "potatoes");
+    request.put("db.query.texts", singleton("INSERT INTO potato VALUES(?)"));
+    request.put(DB_OPERATION_BATCH_SIZE.getKey(), 0L);
+
+    Context context = Context.root();
+
+    AttributesExtractor<Map<String, Object>, Void> underTest =
+        SqlClientAttributesExtractor.create(new TestMultiAttributesGetter());
+
+    // when
+    AttributesBuilder startAttributes = Attributes.builder();
+    underTest.onStart(startAttributes, context, request);
+
+    AttributesBuilder endAttributes = Attributes.builder();
+    underTest.onEnd(endAttributes, context, request, null, null);
+
+    // then
+    if (emitStableDatabaseSemconv() && emitOldDatabaseSemconv()) {
+      assertThat(startAttributes.build())
+          .containsOnly(
+              entry(DB_NAME, "potatoes"),
+              entry(DB_STATEMENT, "INSERT INTO potato VALUES(?)"),
+              entry(DB_OPERATION, "INSERT"),
+              entry(DB_SQL_TABLE, "potato"),
+              entry(DB_NAMESPACE, "potatoes"),
+              entry(DB_QUERY_TEXT, "INSERT INTO potato VALUES(?)"),
+              entry(DB_QUERY_SUMMARY, "BATCH INSERT potato"),
+              entry(DB_OPERATION_BATCH_SIZE, 0L));
+    } else if (emitOldDatabaseSemconv()) {
+      assertThat(startAttributes.build())
+          .containsOnly(
+              entry(DB_NAME, "potatoes"),
+              entry(DB_STATEMENT, "INSERT INTO potato VALUES(?)"),
+              entry(DB_OPERATION, "INSERT"),
+              entry(DB_SQL_TABLE, "potato"));
+    } else if (emitStableDatabaseSemconv()) {
+      assertThat(startAttributes.build())
+          .containsOnly(
+              entry(DB_NAMESPACE, "potatoes"),
+              entry(DB_QUERY_TEXT, "INSERT INTO potato VALUES(?)"),
+              entry(DB_QUERY_SUMMARY, "BATCH INSERT potato"),
+              entry(DB_OPERATION_BATCH_SIZE, 0L));
+    }
+
+    assertThat(endAttributes.build().isEmpty()).isTrue();
+  }
+
   @Test
   void shouldExtractMultiQueryBatchAttributes() {
     // given
@@ -412,6 +465,66 @@ void shouldExtractMixedParameterizedMultiQueryBatchAttributes() {
     assertThat(endAttributes.build().isEmpty()).isTrue();
   }
 
+  @Test
+  void shouldExtractMultiQueryBatchOperationNameWhenSingleOperationAndCollection() {
+    // given
+    Map<String, Object> sameOperation = new HashMap<>();
+    sameOperation.put("db.namespace", "potatoes");
+    sameOperation.put(
+        "db.query.texts", asList("INSERT INTO potato VALUES(1)", "INSERT INTO potato VALUES(2)"));
+    sameOperation.put(DB_OPERATION_BATCH_SIZE.getKey(), 2L);
+
+    Map<String, Object> mixedOperations = new HashMap<>();
+    mixedOperations.put("db.namespace", "potatoes");
+    mixedOperations.put(
+        "db.query.texts",
+        asList("INSERT INTO potato VALUES(1)", "UPDATE potato SET name='bob' WHERE id=1"));
+    mixedOperations.put(DB_OPERATION_BATCH_SIZE.getKey(), 2L);
+
+    Map<String, Object> mixedCollections = new HashMap<>();
+    mixedCollections.put("db.namespace", "potatoes");
+    mixedCollections.put(
+        "db.query.texts", asList("INSERT INTO potato VALUES(1)", "INSERT INTO tomato VALUES(2)"));
+    mixedCollections.put(DB_OPERATION_BATCH_SIZE.getKey(), 2L);
+
+    Context context = Context.root();
+
+    AttributesExtractor<Map<String, Object>, Void> underTest =
+        SqlClientAttributesExtractor.builder(new TestMultiAttributesGetter())
+            .setSingleOperationAndCollection(true)
+            .build();
+
+    // when
+    AttributesBuilder sameOperationAttributes = Attributes.builder();
+    underTest.onStart(sameOperationAttributes, context, sameOperation);
+
+    AttributesBuilder mixedOperationsAttributes = Attributes.builder();
+    underTest.onStart(mixedOperationsAttributes, context, mixedOperations);
+
+    AttributesBuilder mixedCollectionsAttributes = Attributes.builder();
+    underTest.onStart(mixedCollectionsAttributes, context, mixedCollections);
+
+    // then
+    if (emitStableDatabaseSemconv()) {
+      assertThat(sameOperationAttributes.build())
+          .containsEntry(DB_OPERATION_NAME, "BATCH INSERT")
+          .containsEntry(DB_COLLECTION_NAME, "potato")
+          .containsEntry(DB_QUERY_SUMMARY, "BATCH INSERT potato");
+      assertThat(mixedOperationsAttributes.build())
+          .containsEntry(DB_OPERATION_NAME, "BATCH")
+          .containsEntry(DB_COLLECTION_NAME, "potato")
+          .containsEntry(DB_QUERY_SUMMARY, "BATCH");
+      // different collections -> db.collection.name is omitted, db.operation.name is BATCH INSERT
+      assertThat(mixedCollectionsAttributes.build())
+          .containsEntry(DB_OPERATION_NAME, "BATCH INSERT")
+          .doesNotContainKey(DB_COLLECTION_NAME);
+    } else {
+      assertThat(sameOperationAttributes.build().get(DB_OPERATION_NAME)).isNull();
+      assertThat(mixedOperationsAttributes.build().get(DB_OPERATION_NAME)).isNull();
+      assertThat(mixedCollectionsAttributes.build().get(DB_OPERATION_NAME)).isNull();
+    }
+  }
+
   @Test
   void shouldIgnoreBatchSizeOne() {
     // given