Merge pull request #170 from axw/apmgorm-take2

 module/apmgorm: introduce GORM instrumentation

CHANGELOG.md

@@ -15,6 +15,7 @@
  - Add `ELASTIC_APM_IGNORE_URLS` config (#158)
  - module/apmsql: fix a bug preventing errors from being captured (#160)
  - Introduce `Tracer.StartTransactionOptions`, drop variadic args from `Tracer.StartTransaction` (#165)
+ - module/apmgorm: introduce GORM instrumentation module (#169, #170)
 
 ## [v0.4.0](https://github.com/elastic/apm-agent-go/releases/tag/v0.4.0)
 

docs/instrumenting.asciidoc

@@ -264,6 +264,31 @@ func main() {
 Spans will be created for queries and other statement executions if the context methods are
 used, and the context includes a transaction.
 
+===== module/apmgorm
+Package apmgorm provides a means of instrumenting [gorm](http://gorm.io) database operations.
+
+To trace GORM operations, import the appropriate `apmgorm/dialects` package (instead of the
+`gorm/dialects` package), and use `apmgorm.Open` (instead of `gorm.Open`). The parameters are
+exactly the same.
+
+Once you have a `*gorm.DB` from `apmgorm.Open`, you can call `apmgorm.WithContext` to
+propagate a context containing a transaction to the operations:
+
+[source,go]
+----
+import (
+	"github.com/elastic/apm-agent-go/module/apmgorm"
+	_ "github.com/elastic/apm-agent-go/module/apmgorm/dialects/postgres"
+)
+
+func main() {
+	db, err := apmgorm.Open("postgres", "")
+	...
+	db = apmgorm.WithContext(ctx, db)
+	db.Find(...) // creates a "SELECT FROM <foo>" span
+}
+----
+
 ===== module/apmgocql
 Package apmgocql provides a means of instrumenting https://github.com/gocql/gocql[gocql] so
 that queries are reported as spans within the current transaction.

internal/sqlutil/doc.go

@@ -0,0 +1,2 @@
+// Package sqlutil provides utilities to SQL-related instrumentation modules.
+package sqlutil

internal/sqlutil/drivername.go

@@ -0,0 +1,32 @@
+package sqlutil
+
+import (
+	"database/sql/driver"
+	"reflect"
+	"strings"
+)
+
+// DriverName returns the name of the driver, based on its type.
+// If the driver name cannot be deduced, DriverName will return
+// "generic".
+func DriverName(d driver.Driver) string {
+	t := reflect.TypeOf(d)
+	if t.Kind() == reflect.Ptr {
+		t = t.Elem()
+	}
+	switch t.Name() {
+	case "SQLiteDriver":
+		return "sqlite3"
+	case "MySQLDriver":
+		return "mysql"
+	case "Driver":
+		// Check suffix in case of vendoring.
+		if strings.HasSuffix(t.PkgPath(), "github.com/lib/pq") {
+			return "postgresql"
+		}
+	}
+	// TODO include the package path of the driver in context
+	// so we can easily determine how the rules above should
+	// be updated.
+	return "generic"
+}

module/apmsql/signature.go → internal/sqlutil/signature.go

@@ -1,12 +1,12 @@
-package apmsql
+package sqlutil
 
 import (
 	"strings"
 
 	"github.com/elastic/apm-agent-go/internal/sqlscanner"
 )
 
-// genericQuerySignature returns the "signature" for a query:
+// QuerySignature returns the "signature" for a query:
 // a high level description of the operation.
 //
 // For DDL statements (CREATE, DROP, ALTER, etc.), we we only
@@ -15,7 +15,7 @@ import (
 // an application. For SELECT, INSERT, and UPDATE, and DELETE,
 // we attempt to extract the first table name. If we are unable
 // to identify the table name, we simply omit it.
-func genericQuerySignature(query string) string {
+func QuerySignature(query string) string {
 	s := sqlscanner.NewScanner(query)
 	for s.Scan() {
 		if s.Token() != sqlscanner.COMMENT {

module/apmsql/signature_test.go → internal/sqlutil/signature_test.go

@@ -1,14 +1,16 @@
-package apmsql
+package sqlutil_test
 
 import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"github.com/elastic/apm-agent-go/internal/sqlutil"
 )
 
 func TestQuerySignature(t *testing.T) {
 	assertSignatureEqual := func(expect, stmt string) {
-		out := genericQuerySignature(stmt)
+		out := sqlutil.QuerySignature(stmt)
 		assert.Equal(t, expect, out, "%s", stmt)
 	}
 
@@ -55,7 +57,7 @@ func TestQuerySignature(t *testing.T) {
 func BenchmarkQuerySignature(b *testing.B) {
 	sql := "SELECT *,(SELECT COUNT(*) FROM table2 WHERE table2.field1 = table1.id) AS count FROM table1 WHERE table1.field1 = 'value'"
 	for i := 0; i < b.N; i++ {
-		signature := genericQuerySignature(sql)
+		signature := sqlutil.QuerySignature(sql)
 		if signature != "SELECT FROM table1" {
 			panic("unexpected result: " + signature)
 		}

module/apmgorm/apmgorm_test.go

@@ -0,0 +1,182 @@
+package apmgorm_test
+
+import (
+	"context"
+	"database/sql"
+	"os"
+	"testing"
+
+	"github.com/jinzhu/gorm"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/elastic/apm-agent-go/apmtest"
+	"github.com/elastic/apm-agent-go/module/apmgorm"
+	_ "github.com/elastic/apm-agent-go/module/apmgorm/dialects/mysql"
+	_ "github.com/elastic/apm-agent-go/module/apmgorm/dialects/postgres"
+	_ "github.com/elastic/apm-agent-go/module/apmgorm/dialects/sqlite"
+	"github.com/elastic/apm-agent-go/module/apmsql"
+)
+
+type Product struct {
+	gorm.Model
+	Code  string
+	Price uint
+}
+
+func TestWithContext(t *testing.T) {
+	t.Run("sqlite3", func(t *testing.T) {
+		testWithContext(t,
+			apmsql.DSNInfo{Database: ":memory:"},
+			"sqlite3", ":memory:",
+		)
+	})
+
+	if os.Getenv("PGHOST") == "" {
+		t.Logf("PGHOST not specified, skipping")
+	} else {
+		t.Run("postgres", func(t *testing.T) {
+			testWithContext(t,
+				apmsql.DSNInfo{Database: "test_db", User: "postgres"},
+				"postgres", "user=postgres password=hunter2 dbname=test_db sslmode=disable",
+			)
+		})
+	}
+
+	if mysqlHost := os.Getenv("MYSQL_HOST"); mysqlHost == "" {
+		t.Logf("MYSQL_HOST not specified, skipping")
+	} else {
+		t.Run("mysql", func(t *testing.T) {
+			testWithContext(t,
+				apmsql.DSNInfo{Database: "test_db", User: "root"},
+				"mysql", "root:hunter2@tcp("+mysqlHost+")/test_db?parseTime=true",
+			)
+		})
+	}
+}
+
+func testWithContext(t *testing.T, dsnInfo apmsql.DSNInfo, dialect string, args ...interface{}) {
+	tx, errors := apmtest.WithTransaction(func(ctx context.Context) {
+		db, err := apmgorm.Open(dialect, args...)
+		require.NoError(t, err)
+		defer db.Close()
+		db = apmgorm.WithContext(ctx, db)
+
+		db.AutoMigrate(&Product{})
+		db.Create(&Product{Code: "L1212", Price: 1000})
+
+		var product Product
+		assert.NoError(t, db.First(&product, "code = ?", "L1212").Error)
+		assert.NoError(t, db.Model(&product).Update("Price", 2000).Error)
+		assert.NoError(t, db.Delete(&product).Error)            // soft
+		assert.NoError(t, db.Unscoped().Delete(&product).Error) // hard
+	})
+	require.NotEmpty(t, tx.Spans)
+	assert.Empty(t, errors)
+
+	spanNames := make([]string, len(tx.Spans))
+	for i, span := range tx.Spans {
+		spanNames[i] = span.Name
+		require.NotNil(t, span.Context)
+		require.NotNil(t, span.Context.Database)
+		assert.Equal(t, dsnInfo.Database, span.Context.Database.Instance)
+		assert.NotEmpty(t, span.Context.Database.Statement)
+		assert.Equal(t, "sql", span.Context.Database.Type)
+		assert.Equal(t, dsnInfo.User, span.Context.Database.User)
+	}
+	assert.Equal(t, []string{
+		"INSERT INTO products",
+		"SELECT FROM products",
+		"UPDATE products",
+		"UPDATE products", // soft delete
+		"DELETE FROM products",
+	}, spanNames)
+}
+
+// TestWithContextNoTransaction checks that using WithContext without
+// a transaction won't cause any issues.
+func TestWithContextNoTransaction(t *testing.T) {
+	db, err := apmgorm.Open("sqlite3", ":memory:")
+	require.NoError(t, err)
+	defer db.Close()
+	db = apmgorm.WithContext(context.Background(), db)
+
+	db.AutoMigrate(&Product{})
+	db.Create(&Product{Code: "L1212", Price: 1000})
+
+	var product Product
+	assert.NoError(t, db.Where("code=?", "L1212").First(&product).Error)
+}
+
+func TestWithContextNonSampled(t *testing.T) {
+	os.Setenv("ELASTIC_APM_TRANSACTION_SAMPLE_RATE", "0")
+	defer os.Unsetenv("ELASTIC_APM_TRANSACTION_SAMPLE_RATE")
+
+	db, err := apmgorm.Open("sqlite3", ":memory:")
+	require.NoError(t, err)
+	defer db.Close()
+	db.AutoMigrate(&Product{})
+
+	tx, _ := apmtest.WithTransaction(func(ctx context.Context) {
+		db = apmgorm.WithContext(ctx, db)
+		db.Create(&Product{Code: "L1212", Price: 1000})
+	})
+	require.Empty(t, tx.Spans)
+}
+
+func TestCaptureErrors(t *testing.T) {
+	db, err := apmgorm.Open("sqlite3", ":memory:")
+	require.NoError(t, err)
+	defer db.Close()
+	db.SetLogger(nopLogger{})
+	db.AutoMigrate(&Product{})
+
+	tx, errors := apmtest.WithTransaction(func(ctx context.Context) {
+		db = apmgorm.WithContext(ctx, db)
+
+		// record not found should not cause an error
+		db.Where("code=?", "L1212").First(&Product{})
+
+		// invalid SQL should
+		db.Where("bananas").First(&Product{})
+	})
+	assert.Len(t, tx.Spans, 2)
+	require.Len(t, errors, 1)
+	assert.Regexp(t, "no such column: bananas", errors[0].Exception.Message)
+}
+
+func TestOpenWithDriver(t *testing.T) {
+	db, err := apmgorm.Open("sqlite3", "sqlite3", ":memory:")
+	require.NoError(t, err)
+	defer db.Close()
+	db.AutoMigrate(&Product{})
+
+	tx, _ := apmtest.WithTransaction(func(ctx context.Context) {
+		db = apmgorm.WithContext(ctx, db)
+		db.Create(&Product{Code: "L1212", Price: 1000})
+	})
+	require.Len(t, tx.Spans, 1)
+	assert.Equal(t, ":memory:", tx.Spans[0].Context.Database.Instance)
+}
+
+func TestOpenWithDB(t *testing.T) {
+	sqldb, err := sql.Open("sqlite3", ":memory:")
+	require.NoError(t, err)
+	defer sqldb.Close()
+
+	db, err := apmgorm.Open("sqlite3", sqldb)
+	require.NoError(t, err)
+	defer db.Close()
+	db.AutoMigrate(&Product{})
+
+	tx, _ := apmtest.WithTransaction(func(ctx context.Context) {
+		db = apmgorm.WithContext(ctx, db)
+		db.Create(&Product{Code: "L1212", Price: 1000})
+	})
+	require.Len(t, tx.Spans, 1)
+	assert.Empty(t, tx.Spans[0].Context.Database.Instance) // no DSN info
+}
+
+type nopLogger struct{}
+
+func (nopLogger) Print(v ...interface{}) {}