module/apmgorm: introduce GORM instrumentation

Package apmgorm provides a means of instrumenting [gorm](http://gorm.io)
database operations.

By using apmgorm.Open instead of gorm.Open, we obtain a *gorm.DB which
can be used with apmgorm.WithContext. The WithContext function can be
used to propagate context (i.e. containing a transaction) to callbacks
which report spans.

We provide "dialects" packages which import the gorm/dialects namesake
packages, as well as register the apmsql drivers. The latter is required
in order to parse DSNs.

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.

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{}) {}

module/apmgorm/context.go

@@ -0,0 +1,130 @@
+package apmgorm
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/jinzhu/gorm"
+
+	"github.com/elastic/apm-agent-go"
+	"github.com/elastic/apm-agent-go/internal/sqlutil"
+	"github.com/elastic/apm-agent-go/module/apmsql"
+)
+
+const (
+	apmContextKey = "elasticapm:context"
+)
+
+// WithContext returns a copy of db with ctx recorded for use by
+// the callbacks registered via RegisterCallbacks.
+func WithContext(ctx context.Context, db *gorm.DB) *gorm.DB {
+	return db.Set(apmContextKey, ctx)
+}
+
+func scopeContext(scope *gorm.Scope) (context.Context, bool) {
+	value, ok := scope.Get(apmContextKey)
+	if !ok {
+		return nil, false
+	}
+	ctx, _ := value.(context.Context)
+	return ctx, ctx != nil
+}
+
+// RegisterCallbacks registers callbacks on db for reporting spans
+// to Elastic APM. This is called automatically by apmgorm.Open;
+// it is provided for cases where a *gorm.DB is acquired by other
+// means.
+func RegisterCallbacks(db *gorm.DB) {
+	registerCallbacks(db, apmsql.DSNInfo{})
+}
+
+func registerCallbacks(db *gorm.DB, dsnInfo apmsql.DSNInfo) {
+	driverName := db.Dialect().GetName()
+	switch driverName {
+	case "postgres":
+		driverName = "postgresql"
+	}
+	spanTypePrefix := fmt.Sprintf("db.%s.", driverName)
+	querySpanType := spanTypePrefix + "query"
+	execSpanType := spanTypePrefix + "exec"
+
+	type params struct {
+		spanType  string
+		processor func() *gorm.CallbackProcessor
+	}
+	callbacks := map[string]params{
+		"gorm:create": {
+			spanType:  execSpanType,
+			processor: func() *gorm.CallbackProcessor { return db.Callback().Create() },
+		},
+		"gorm:delete": {
+			spanType:  execSpanType,
+			processor: func() *gorm.CallbackProcessor { return db.Callback().Delete() },
+		},
+		"gorm:query": {
+			spanType:  querySpanType,
+			processor: func() *gorm.CallbackProcessor { return db.Callback().Query() },
+		},
+		"gorm:update": {
+			spanType:  execSpanType,
+			processor: func() *gorm.CallbackProcessor { return db.Callback().Update() },
+		},
+	}
+	for name, params := range callbacks {
+		const callbackPrefix = "elasticapm"
+		params.processor().Before(name).Register(
+			fmt.Sprintf("%s:before:%s", callbackPrefix, name),
+			newBeforeCallback(params.spanType),
+		)
+		params.processor().After(name).Register(
+			fmt.Sprintf("%s:after:%s", callbackPrefix, name),
+			newAfterCallback(dsnInfo),
+		)
+	}
+}
+
+func newBeforeCallback(spanType string) func(*gorm.Scope) {
+	return func(scope *gorm.Scope) {
+		ctx, ok := scopeContext(scope)
+		if !ok {
+			return
+		}
+		span, ctx := elasticapm.StartSpan(ctx, "", spanType)
+		if span.Dropped() {
+			span.End()
+			ctx = nil
+		}
+		scope.Set(apmContextKey, ctx)
+	}
+}
+
+func newAfterCallback(dsnInfo apmsql.DSNInfo) func(*gorm.Scope) {
+	return func(scope *gorm.Scope) {
+		ctx, ok := scopeContext(scope)
+		if !ok {
+			return
+		}
+		span := elasticapm.SpanFromContext(ctx)
+		if span == nil {
+			return
+		}
+		span.Name = sqlutil.QuerySignature(scope.SQL)
+		span.Context.SetDatabase(elasticapm.DatabaseSpanContext{
+			Instance:  dsnInfo.Database,
+			Statement: scope.SQL,
+			Type:      "sql",
+			User:      dsnInfo.User,
+		})
+		span.End()
+
+		// Capture errors, except for "record not found", which may be expected.
+		for _, err := range scope.DB().GetErrors() {
+			if gorm.IsRecordNotFoundError(err) {
+				continue
+			}
+			if e := elasticapm.CaptureError(ctx, err); e != nil {
+				e.Send()
+			}
+		}
+	}
+}

module/apmgorm/dialects/mysql/init.go

@@ -0,0 +1,9 @@
+// Package apmgormmysql imports the gorm mysql dialect package,
+// and also registers the mysql driver with apmsql.
+package apmgormmysql
+
+import (
+	_ "github.com/jinzhu/gorm/dialects/mysql" // import the mysql dialect
+
+	_ "github.com/elastic/apm-agent-go/module/apmsql/mysql" // register mysql with apmsql
+)

module/apmgorm/dialects/postgres/init.go

@@ -0,0 +1,9 @@
+// Package apmgormpostgres imports the gorm postgres dialect package,
+// and also registers the lib/pq driver with apmsql.
+package apmgormpostgres
+
+import (
+	_ "github.com/jinzhu/gorm/dialects/postgres" // import the postgres dialect
+
+	_ "github.com/elastic/apm-agent-go/module/apmsql/pq" // register lib/pq with apmsql
+)

module/apmgorm/dialects/sqlite/init.go

@@ -0,0 +1,9 @@
+// Package apmgormsqlite imports the gorm sqlite dialect package,
+// and also registers the sqlite3 driver with apmsql.
+package apmgormsqlite
+
+import (
+	_ "github.com/jinzhu/gorm/dialects/sqlite" // import the sqlite dialect
+
+	_ "github.com/elastic/apm-agent-go/module/apmsql/sqlite3" // register sqlite3 with apmsql
+)

module/apmgorm/doc.go

@@ -0,0 +1,2 @@
+// Package apmgorm provides wrappers for tracing GORM operations.
+package apmgorm