Capture row counts for plain statements and batch (#1333)

Previously, only row counts for executeUpdate on prepared statements was
included as sql.affected_rows. With this change, row counts are also
included for executeUpdate on plain statements, and for executeBatch on
prepared and plain statements.

Capturing row counts on plain execute is a lot trickier, since it´s not
safe to call getUpdateCount behind the application´s back, as it is not
idempotent - we´d unexpectedly change state.

Co-authored-by: Magnus Reftel <Magnus.Reftel@skatteetaten.no>

instrumentation/p6spy/README.md

@@ -38,8 +38,10 @@ excludebinary=true
 
 ### Affected rows count
 When the `includeAffectedRowsCount` option is set to `true`, the tag `sql.affected_rows` of traces
-for SQL insert and update commands will include the number of rows that were inserted/updated, if
-the database and driver supports that. No row count is included for select statements.
+for SQL insert and update statements performed using `executeUpdate` and `executeBatch` (but not
+`execute`, as it does not report the update count) will include the number of rows that were
+inserted/updated (comma separated in the case of `executebatch`), if the database and driver
+supports that. No row count is included for calls to `executeQuery` (e.g. select statements).
 
 ## Service name as URL query parameter
 `spy.properties` applies globally to any instrumented jdbc connection. To override this, add the

instrumentation/p6spy/src/main/java/brave/p6spy/TracingJdbcEventListener.java

@@ -18,7 +18,6 @@
 import brave.propagation.ThreadLocalSpan;
 
 import com.p6spy.engine.common.PreparedStatementInformation;
-import com.p6spy.engine.common.ResultSetInformation;
 import com.p6spy.engine.common.StatementInformation;
 import com.p6spy.engine.event.SimpleJdbcEventListener;
 import com.p6spy.engine.logging.P6LogLoadableOptions;
@@ -87,6 +86,33 @@ final class TracingJdbcEventListener extends SimpleJdbcEventListener {
     finishSpan(span, e);
   }
 
+  @Override public void onAfterExecuteUpdate(
+    StatementInformation statementInformation, long timeElapsedNanos, String sql, int rowCount, SQLException e
+  ) {
+    Span span = ThreadLocalSpan.CURRENT_TRACER.remove();
+    if (span == null || span.isNoop()) return;
+    if (includeAffectedRowsCount) {
+      span.tag("sql.affected_rows", String.valueOf(rowCount));
+    }
+    finishSpan(span, e);
+  }
+
+  @Override
+  public void onAfterExecuteBatch(StatementInformation statementInformation, long timeElapsedNanos, int[] updateCounts, SQLException e) {
+    Span span = ThreadLocalSpan.CURRENT_TRACER.remove();
+    if (span == null || span.isNoop()) return;
+    if (includeAffectedRowsCount && updateCounts.length > 0) {
+      StringBuilder sb = new StringBuilder();
+      sb.append(updateCounts[0]);
+      for (int i = 1; i < updateCounts.length; i++) {
+        sb.append(',');
+        sb.append(updateCounts[i]);
+      }
+      span.tag("sql.affected_rows", sb.toString());
+    }
+    finishSpan(span, e);
+  }
+
   @Override public void onAfterAnyExecute(StatementInformation info, long elapsed, SQLException e) {
     Span span = ThreadLocalSpan.CURRENT_TRACER.remove();
     if (span == null || span.isNoop()) return;

instrumentation/p6spy/src/test/java/brave/p6spy/ITTracingP6Factory.java

@@ -113,14 +113,100 @@ public void addsQueryTag() throws Exception {
   }
 
   @Test
-  public void addsRowsWrittenTag() throws Exception {
+  public void addsAffectedRowsTagToPreparedUpdateStatements() throws Exception {
     prepareExecuteUpdate("update t set c='x' where i=2");
 
     assertThat(spans)
       .flatExtracting(s -> s.tags().entrySet())
       .contains(entry("sql.affected_rows", "2"));
   }
 
+  @Test
+  public void addsAffectedRowsTagToPlainUpdateStatements() throws Exception {
+    executeUpdate("update t set c='x' where i=2");
+
+    assertThat(spans)
+      .flatExtracting(s -> s.tags().entrySet())
+      .contains(entry("sql.affected_rows", "2"));
+  }
+
+  @Test
+  public void addsAffectedRowsTagToPlainBatchUpdateStatements() throws Exception {
+    executeBatch("update t set c='x' where i=2", "update t set c='y' where i=1");
+
+    assertThat(spans)
+      .flatExtracting(s -> s.tags().entrySet())
+      .contains(entry("sql.affected_rows", "2,1"));
+  }
+
+  @Test
+  public void doesNotProduceAnySpansForEmptyPlainBatchUpdates() throws Exception {
+    // No SQL at all means no span is started in onBeforeAnyExecute due to there not being any loggable SQL
+    // (see isLoggable)
+    executeBatch();
+
+    assertThat(spans).isEmpty();
+  }
+
+  @Test
+  public void addsAffectedRowsTagToPreparedBatchUpdateStatementsWithOneBatch() throws Exception {
+    prepareExecuteBatchWithInts("update t set c='x' where i=?", 2);
+
+    assertThat(spans)
+      .flatExtracting(s -> s.tags().entrySet())
+      .contains(entry("sql.affected_rows", "2"));
+  }
+
+  @Test
+  public void addsAffectedRowsTagToPreparedBatchUpdateStatementsWithOneBatchWithZeroUpdates() throws Exception {
+    prepareExecuteBatchWithInts("update t set c='x' where i=?", 0);
+
+    assertThat(spans)
+      .flatExtracting(s -> s.tags().entrySet())
+      .contains(entry("sql.affected_rows", "0"));
+  }
+
+  @Test
+  public void addsAffectedRowsTagToPreparedBatchUpdateStatementsWithMoreThanOneBatch() throws Exception {
+    prepareExecuteBatchWithInts("update t set c='x' where i=?", 2, 1);
+
+    assertThat(spans)
+      .flatExtracting(s -> s.tags().entrySet())
+      .contains(entry("sql.affected_rows", "2,1"));
+  }
+
+  @Test
+  public void addsAffectedRowsTagToPreparedBatchUpdateStatementsWithMoreThanOneBatchWhereOneBatcheHasZeroUpdates() throws Exception {
+    prepareExecuteBatchWithInts("update t set c='x' where i=?", 2, 0, 1);
+
+    assertThat(spans)
+      .flatExtracting(s -> s.tags().entrySet())
+      .contains(entry("sql.affected_rows", "2,0,1"));
+  }
+
+  @Test
+  public void addsAffectedRowsTagToPreparedBatchUpdateStatementsWithMoreThanOneBatchWhereBatchesHaveZeroUpdates() throws Exception {
+    prepareExecuteBatchWithInts("update t set c='x' where i=?", 2, 0, 3);
+
+    assertThat(spans)
+      .flatExtracting(s -> s.tags().entrySet())
+      .contains(entry("sql.affected_rows", "2,0,0"));
+  }
+
+  @Test
+  public void addsEmptyAffectedRowsTagToEmptyPreparedBatchUpdates() throws Exception {
+    // In contrast to the plain statement case, this does produce loggable SQL, so a span is started. Since there are
+    // no entries in the batch, no updates are made, so there are no update counts. Therefore, the span does not have
+    // any sql.affected_rows tag.
+    prepareExecuteBatchWithInts("update t set c='x' where i=?");
+
+    assertThat(spans).anySatisfy(span -> {
+      assertThat(span.tags())
+        .contains(entry("sql.query", "update t set c='x' where i=?"))
+        .doesNotContainKey("sql.affected_rows");
+    });
+  }
+
   @Test
   public void reportsServerAddress() throws Exception {
     prepareExecuteSelect(QUERY);
@@ -145,4 +231,29 @@ void prepareExecuteUpdate(String sql) throws SQLException {
       ps.executeUpdate();
     }
   }
+
+  void executeUpdate(String sql) throws SQLException {
+    try (Statement s = connection.createStatement()) {
+      s.executeUpdate(sql);
+    }
+  }
+
+  void executeBatch(String... sqls) throws SQLException {
+    try (Statement s = connection.createStatement()) {
+      for (String sql : sqls) {
+        s.addBatch(sql);
+      }
+      s.executeBatch();
+    }
+  }
+
+  void prepareExecuteBatchWithInts(String sql, int... ints) throws SQLException {
+    try (PreparedStatement s = connection.prepareStatement(sql)) {
+      for (int i : ints) {
+        s.setInt(1, i);
+        s.addBatch();
+      }
+      s.executeBatch();
+    }
+  }
 }