module/apmsql: report rows affected

CHANGELOG.asciidoc

@@ -24,6 +24,7 @@ endif::[]
 https://github.com/elastic/apm-agent-go/compare/v1.6.0...master[View commits]
 
  - Add support for API Key auth {pull}698[(#698)]
+ - module/apmsql: report rows affected {pull}700[(#700)]
 
 [[release-notes-1.x]]
 === Go Agent version 1.x

model/model.go

@@ -314,6 +314,10 @@ type DatabaseSpanContext struct {
 	// Statement holds the database statement (e.g. query).
 	Statement string `json:"statement,omitempty"`
 
+	// RowsAffected holds the number of rows affected by the
+	// database operation.
+	RowsAffected *int64 `json:"rows_affected,omitempty"`
+
 	// Type holds the database type. For any SQL database,
 	// this should be "sql"; for others, the lower-cased
 	// database category, e.g. "cassandra", "hbase", "redis".

module/apmsql/apmsql_bench_test.go

@@ -78,3 +78,50 @@ func benchmarkQueries(b *testing.B, ctx context.Context, stmt *sql.Stmt) {
 		rows.Close()
 	}
 }
+
+func BenchmarkStmtExecContext(b *testing.B) {
+	db, err := apmsql.Open("sqlite3", ":memory:")
+	require.NoError(b, err)
+	defer db.Close()
+
+	_, err = db.Exec("CREATE TABLE foo (bar INT)")
+	require.NoError(b, err)
+
+	stmt, err := db.Prepare("DELETE FROM foo")
+	require.NoError(b, err)
+	defer stmt.Close()
+
+	b.Run("baseline", func(b *testing.B) {
+		benchmarkExec(b, context.Background(), stmt)
+	})
+	b.Run("instrumented", func(b *testing.B) {
+		invalidServerURL, err := url.Parse("http://testing.invalid:8200")
+		if err != nil {
+			panic(err)
+		}
+		httpTransport, err := transport.NewHTTPTransport()
+		require.NoError(b, err)
+		httpTransport.SetServerURL(invalidServerURL)
+
+		tracer, err := apm.NewTracerOptions(apm.TracerOptions{
+			ServiceName:    "apmhttp_test",
+			ServiceVersion: "0.1",
+			Transport:      httpTransport,
+		})
+		require.NoError(b, err)
+		defer tracer.Close()
+
+		tracer.SetMaxSpans(b.N)
+		tx := tracer.StartTransaction("name", "type")
+		ctx := apm.ContextWithTransaction(context.Background(), tx)
+		benchmarkExec(b, ctx, stmt)
+	})
+}
+
+func benchmarkExec(b *testing.B, ctx context.Context, stmt *sql.Stmt) {
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := stmt.ExecContext(ctx)
+		require.NoError(b, err)
+	}
+}

module/apmsql/apmsql_test.go

@@ -60,15 +60,67 @@ func TestExecContext(t *testing.T) {
 	defer db.Close()
 
 	db.Ping() // connect
+	const N = 5
 	_, spans, _ := apmtest.WithTransaction(func(ctx context.Context) {
 		_, err := db.ExecContext(ctx, "CREATE TABLE foo (bar INT)")
 		require.NoError(t, err)
+		for i := 0; i < N; i++ {
+			result, err := db.ExecContext(ctx, "INSERT INTO foo VALUES (?)", i)
+			require.NoError(t, err)
+
+			rowsAffected, err := result.RowsAffected()
+			require.NoError(t, err)
+			assert.Equal(t, int64(1), rowsAffected)
+		}
+		result, err := db.ExecContext(ctx, "DELETE FROM foo")
+		require.NoError(t, err)
+		rowsAffected, err := result.RowsAffected()
+		require.NoError(t, err)
+		assert.Equal(t, int64(N), rowsAffected)
 	})
-	require.Len(t, spans, 1)
+	require.Len(t, spans, 2+N)
+
+	int64ptr := func(n int64) *int64 { return &n }
+
 	assert.Equal(t, "CREATE", spans[0].Name)
 	assert.Equal(t, "db", spans[0].Type)
 	assert.Equal(t, "sqlite3", spans[0].Subtype)
 	assert.Equal(t, "exec", spans[0].Action)
+	assert.Equal(t, &model.SpanContext{
+		Database: &model.DatabaseSpanContext{
+			Instance: ":memory:",
+			// Ideally RowsAffected would not be returned for DDL
+			// statements, but this is on the driver; it should
+			// be returning database/sql/driver.ResultNoRows for
+			// DDL statements, in which case we'll omit this.
+			RowsAffected: int64ptr(0),
+			Statement:    "CREATE TABLE foo (bar INT)",
+			Type:         "sql",
+		},
+	}, spans[0].Context)
+
+	for i := 0; i < N; i++ {
+		span := spans[i+1]
+		assert.Equal(t, "INSERT INTO foo", span.Name)
+		assert.Equal(t, &model.SpanContext{
+			Database: &model.DatabaseSpanContext{
+				Instance:     ":memory:",
+				RowsAffected: int64ptr(1),
+				Statement:    "INSERT INTO foo VALUES (?)",
+				Type:         "sql",
+			},
+		}, span.Context)
+	}
+
+	assert.Equal(t, "DELETE FROM foo", spans[N+1].Name)
+	assert.Equal(t, &model.SpanContext{
+		Database: &model.DatabaseSpanContext{
+			Instance:     ":memory:",
+			RowsAffected: int64ptr(N),
+			Statement:    "DELETE FROM foo",
+			Type:         "sql",
+		},
+	}, spans[N+1].Context)
 }
 
 func TestQueryContext(t *testing.T) {

module/apmsql/conn.go

@@ -83,7 +83,7 @@ func (c *conn) startSpan(ctx context.Context, name, spanType, stmt string) (*apm
 	return span, ctx
 }
 
-func (c *conn) finishSpan(ctx context.Context, span *apm.Span, resultError *error) {
+func (c *conn) finishSpan(ctx context.Context, span *apm.Span, result *driver.Result, resultError *error) {
 	if *resultError == driver.ErrSkip {
 		// TODO(axw) mark span as abandoned,
 		// so it's not sent and not counted
@@ -93,7 +93,14 @@ func (c *conn) finishSpan(ctx context.Context, span *apm.Span, resultError *erro
 		return
 	}
 	switch *resultError {
-	case nil, driver.ErrBadConn, context.Canceled:
+	case nil:
+		if !span.Dropped() && result != nil && *result != nil && *result != driver.ResultNoRows {
+			rowsAffected, err := (*result).RowsAffected()
+			if err == nil && rowsAffected >= 0 {
+				span.Context.SetDatabaseRowsAffected(rowsAffected)
+			}
+		}
+	case driver.ErrBadConn, context.Canceled:
 		// ErrBadConn is used by the connection pooling
 		// logic in database/sql, and so is expected and
 		// should not be reported.
@@ -113,7 +120,7 @@ func (c *conn) Ping(ctx context.Context) (resultError error) {
 		return nil
 	}
 	span, ctx := c.startSpan(ctx, "ping", c.driver.pingSpanType, "")
-	defer c.finishSpan(ctx, span, &resultError)
+	defer c.finishSpan(ctx, span, nil, &resultError)
 	return c.pinger.Ping(ctx)
 }
 
@@ -122,7 +129,7 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 		return nil, driver.ErrSkip
 	}
 	span, ctx := c.startStmtSpan(ctx, query, c.driver.querySpanType)
-	defer c.finishSpan(ctx, span, &resultError)
+	defer c.finishSpan(ctx, span, nil, &resultError)
 
 	if c.queryerContext != nil {
 		return c.queryerContext.QueryContext(ctx, query, args)
@@ -145,7 +152,7 @@ func (*conn) Query(query string, args []driver.Value) (driver.Rows, error) {
 
 func (c *conn) PrepareContext(ctx context.Context, query string) (_ driver.Stmt, resultError error) {
 	span, ctx := c.startStmtSpan(ctx, query, c.driver.prepareSpanType)
-	defer c.finishSpan(ctx, span, &resultError)
+	defer c.finishSpan(ctx, span, nil, &resultError)
 	var stmt driver.Stmt
 	var err error
 	if c.connPrepareContext != nil {
@@ -167,12 +174,12 @@ func (c *conn) PrepareContext(ctx context.Context, query string) (_ driver.Stmt,
 	return stmt, err
 }
 
-func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (_ driver.Result, resultError error) {
+func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (result driver.Result, resultError error) {
 	if c.execerContext == nil && c.execer == nil {
 		return nil, driver.ErrSkip
 	}
 	span, ctx := c.startStmtSpan(ctx, query, c.driver.execSpanType)
-	defer c.finishSpan(ctx, span, &resultError)
+	defer c.finishSpan(ctx, span, &result, &resultError)
 
 	if c.execerContext != nil {
 		return c.execerContext.ExecContext(ctx, query, args)

module/apmsql/stmt.go

@@ -64,9 +64,9 @@ func (s *stmt) ColumnConverter(idx int) driver.ValueConverter {
 	return driver.DefaultParameterConverter
 }
 
-func (s *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (_ driver.Result, resultError error) {
+func (s *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (result driver.Result, resultError error) {
 	span, ctx := s.startSpan(ctx, s.conn.driver.execSpanType)
-	defer s.conn.finishSpan(ctx, span, &resultError)
+	defer s.conn.finishSpan(ctx, span, &result, &resultError)
 	if s.stmtExecContext != nil {
 		return s.stmtExecContext.ExecContext(ctx, args)
 	}
@@ -84,7 +84,7 @@ func (s *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (_ dri
 
 func (s *stmt) QueryContext(ctx context.Context, args []driver.NamedValue) (_ driver.Rows, resultError error) {
 	span, ctx := s.startSpan(ctx, s.conn.driver.querySpanType)
-	defer s.conn.finishSpan(ctx, span, &resultError)
+	defer s.conn.finishSpan(ctx, span, nil, &resultError)
 	if s.stmtQueryContext != nil {
 		return s.stmtQueryContext.QueryContext(ctx, args)
 	}

spancontext.go

@@ -29,11 +29,12 @@ import (
 
 // SpanContext provides methods for setting span context.
 type SpanContext struct {
-	model              model.SpanContext
-	destination        model.DestinationSpanContext
-	destinationService model.DestinationServiceSpanContext
-	database           model.DatabaseSpanContext
-	http               model.HTTPSpanContext
+	model                model.SpanContext
+	destination          model.DestinationSpanContext
+	destinationService   model.DestinationServiceSpanContext
+	databaseRowsAffected int64
+	database             model.DatabaseSpanContext
+	http                 model.HTTPSpanContext
 }
 
 // DatabaseSpanContext holds database span context.
@@ -120,6 +121,13 @@ func (c *SpanContext) SetDatabase(db DatabaseSpanContext) {
 	c.model.Database = &c.database
 }
 
+// SetDatabaseRowsAffected records the number of rows affected by
+// a database operation.
+func (c *SpanContext) SetDatabaseRowsAffected(n int64) {
+	c.databaseRowsAffected = n
+	c.database.RowsAffected = &c.databaseRowsAffected
+}
+
 // SetHTTPRequest sets the details of the HTTP request in the context.
 //
 // This function relates to client requests. If the request URL contains