Support gormio

.gitignore

@@ -3,3 +3,4 @@
 *.out
 docs/html
 build
+.idea

module/apmgormv2/README.md

@@ -0,0 +1,20 @@
+# apmgormv2
+
+Package apmgormv2 provides drivers to gorm.io/gorm
+for tracing database operations as spans of a transaction traced
+by Elastic APM.
+
+## Usage
+
+Swap `gorm.io/driver/*` to `go.elastic.co/apm/module/apmgormv2/driver/*`
+
+Example :-
+
+```golang
+import (
+    mysql "go.elastic.co/apm/module/apmgormv2/driver/mysql"
+    "gorm.io/gorm"
+)	
+
+db, err := gorm.Open(mysql.Open("dsn"), &gorm.Config{})
+```

module/apmgormv2/apmgorm_test.go

@@ -0,0 +1,229 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// +build go1.14
+
+package apmgormv2_test
+
+import (
+	"context"
+	"os"
+	"strings"
+	"testing"
+
+	"gorm.io/gorm/logger"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+
+	"go.elastic.co/apm/apmtest"
+	mysql "go.elastic.co/apm/module/apmgormv2/driver/mysql"
+	sqlite "go.elastic.co/apm/module/apmgormv2/driver/sqlite"
+	"go.elastic.co/apm/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:"},
+			sqlite.Open(":memory:"), &gorm.Config{},
+		)
+	})
+
+	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{
+					Address:  mysqlHost,
+					Port:     3306,
+					Database: "test_db",
+					User:     "root",
+				},
+				mysql.Open("root:hunter2@tcp("+mysqlHost+")/test_db?parseTime=true"), &gorm.Config{},
+			)
+		})
+	}
+}
+
+func testWithContext(t *testing.T, dsnInfo apmsql.DSNInfo, dialect gorm.Dialector, config *gorm.Config) {
+	_, spans, errors := apmtest.WithTransaction(func(ctx context.Context) {
+		db, err := gorm.Open(dialect, config)
+		require.NoError(t, err)
+		ddb, _ := db.DB()
+		defer ddb.Close()
+
+		db = db.WithContext(ctx)
+
+		if db.Migrator().HasTable(&Product{}) {
+			db.Migrator().DropTable(&Product{})
+		}
+		db.AutoMigrate(&Product{})
+		db.Create(&Product{Code: "L1212", Price: 1000})
+
+		var product Product
+		var count int64
+		assert.NoError(t, db.Model(&product).Count(&count).Error)
+		assert.Equal(t, int64(1), count)
+		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, spans)
+	assert.Empty(t, errors)
+
+	var spanNames []string
+	for _, span := range spans {
+		if strings.Contains(span.Name, "products") && span.Action != "prepare" {
+			spanNames = append(spanNames, 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)
+		if dsnInfo.Address == "" {
+			assert.Nil(t, span.Context.Destination)
+		} else {
+			assert.Equal(t, dsnInfo.Address, span.Context.Destination.Address)
+			assert.Equal(t, dsnInfo.Port, span.Context.Destination.Port)
+		}
+	}
+	assert.Equal(t, []string{
+		"INSERT INTO products",
+		"SELECT FROM products", // count
+		"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 := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+	require.NoError(t, err)
+	ddb, _ := db.DB()
+	defer ddb.Close()
+
+	db = db.WithContext(context.Background())
+	if db.Migrator().HasTable(&Product{}) {
+		db.Migrator().DropTable(&Product{})
+	}
+	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 := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+
+	require.NoError(t, err)
+
+	ddb, _ := db.DB()
+	defer ddb.Close()
+
+	if db.Migrator().HasTable(&Product{}) {
+		db.Migrator().DropTable(&Product{})
+	}
+	db.AutoMigrate(&Product{})
+
+	_, spans, _ := apmtest.WithTransaction(func(ctx context.Context) {
+		db = db.WithContext(ctx)
+		db.Create(&Product{Code: "L1212", Price: 1000})
+	})
+	require.Empty(t, spans)
+}
+
+func TestCaptureErrors(t *testing.T) {
+	t.Run("sqlite3", func(t *testing.T) {
+		db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+		require.NoError(t, err)
+
+		ddb, _ := db.DB()
+		defer ddb.Close()
+
+		testCaptureErrors(t, db)
+	})
+}
+
+func testCaptureErrors(t *testing.T, db *gorm.DB) {
+	db.Config.Logger = logger.Default.LogMode(logger.Silent)
+	if db.Migrator().HasTable(&Product{}) {
+		db.Migrator().DropTable(&Product{})
+	}
+	db.AutoMigrate(&Product{})
+
+	_, spans, errors := apmtest.WithTransaction(func(ctx context.Context) {
+		db = db.WithContext(ctx)
+
+		// gorm.ErrRecordNotFound should not cause an error
+		db.Where("code=?", "L1212").First(&Product{})
+
+		product := Product{
+			Model: gorm.Model{
+				ID: 1001,
+			},
+			Code:  "1001",
+			Price: 1001,
+		}
+		require.NoError(t, db.Create(&product).Error)
+
+		// invalid SQL should
+		db.Where("bananas").First(&Product{})
+	})
+	assert.Len(t, spans, 3)
+	require.Len(t, errors, 1)
+	assert.Regexp(t, `.*bananas.*`, errors[0].Exception.Message)
+}
+
+func TestOpenWithDriver(t *testing.T) {
+	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+
+	require.NoError(t, err)
+
+	ddb, _ := db.DB()
+	defer ddb.Close()
+
+	if db.Migrator().HasTable(&Product{}) {
+		db.Migrator().DropTable(&Product{})
+	}
+	db.AutoMigrate(&Product{})
+
+	_, spans, _ := apmtest.WithTransaction(func(ctx context.Context) {
+		db = db.WithContext(ctx)
+		db.Create(&Product{Code: "L1212", Price: 1000})
+	})
+	require.Len(t, spans, 1)
+	assert.Equal(t, ":memory:", spans[0].Context.Database.Instance)
+}

module/apmgormv2/doc.go

@@ -0,0 +1,19 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Package apmgormv2 provides wrappers for tracing GORM operations.
+package apmgormv2

module/apmgormv2/driver/mysql/init.go

@@ -0,0 +1,44 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// +build go1.14
+
+// Package apmmysql imports the gorm mysql dialect package,
+// and also registers the mysql driver with apmsql.
+package apmmysql
+
+import (
+	"gorm.io/driver/mysql"
+	"gorm.io/gorm"
+
+	"go.elastic.co/apm/module/apmsql"
+
+	_ "go.elastic.co/apm/module/apmsql/mysql" // register mysql with apmsql
+)
+
+// Open creates a dialect with apmsql
+func Open(dsn string) gorm.Dialector {
+	driverName := mysql.Dialector{}.Name()
+	dialect := &mysql.Dialector{
+		Config: &mysql.Config{
+			DriverName: apmsql.DriverPrefix + driverName,
+			DSN:        dsn,
+		},
+	}
+
+	return dialect
+}

module/apmgormv2/driver/sqlite/init.go

@@ -0,0 +1,42 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// +build go1.14
+
+// Package apmsqlite imports the gorm sqlite dialect package,
+// and also registers the sqlite3 driver with apmsql.
+package apmsqlite
+
+import (
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+
+	"go.elastic.co/apm/module/apmsql"
+
+	_ "go.elastic.co/apm/module/apmsql/sqlite3" // register sqlite3 with apmsql
+)
+
+// Open creates a dialect with apmsql
+func Open(dsn string) gorm.Dialector {
+
+	dialect := &sqlite.Dialector{
+		DSN:        dsn,
+		DriverName: apmsql.DriverPrefix + sqlite.DriverName,
+	}
+
+	return dialect
+}

module/apmgormv2/go.mod

@@ -0,0 +1,17 @@
+module go.elastic.co/apm/module/apmgormv2
+
+require (
+	github.com/stretchr/testify v1.5.1
+	go.elastic.co/apm v1.8.0
+	go.elastic.co/apm/module/apmsql v1.8.0
+	golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae // indirect
+	gorm.io/driver/mysql v1.0.2
+	gorm.io/driver/sqlite v1.1.4-0.20200928065301-698e250a3b0d
+	gorm.io/gorm v1.20.2
+)
+
+replace go.elastic.co/apm => ../..
+
+replace go.elastic.co/apm/module/apmsql => ../apmsql
+
+go 1.13