contrib/jinzhu/gorm: add gorm integration (#268)

contrib/jinzhu/gorm/example_test.go

@@ -0,0 +1,31 @@
+package gorm_test
+
+import (
+	"log"
+
+	"github.com/jinzhu/gorm"
+	"github.com/lib/pq"
+
+	sqltrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
+	gormtrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/jinzhu/gorm"
+)
+
+func ExampleOpen() {
+	// Register augments the provided driver with tracing, enabling it to be loaded by gormtrace.Open.
+	sqltrace.Register("postgres", &pq.Driver{}, sqltrace.WithServiceName("my-service"))
+
+	// Open the registered driver, allowing all uses of the returned *gorm.DB to be traced.
+	db, err := gormtrace.Open("postgres", "postgres://pqgotest:password@localhost/pqgotest?sslmode=disable")
+	defer db.Close()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	user := struct {
+		gorm.Model
+		Name string
+	}{}
+
+	// All calls through gorm.DB are now traced.
+	db.Where("name = ?", "jinzhu").First(&user)
+}

contrib/jinzhu/gorm/gorm.go

@@ -0,0 +1,18 @@
+// Package gorm provides helper functions for tracing the jinzhu/gorm package (https://github.com/jinzhu/gorm).
+package gorm
+
+import (
+	sqltraced "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
+
+	"github.com/jinzhu/gorm"
+)
+
+// Open opens a new (traced) database connection. The used dialect must be formerly registered
+// using (gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql).Register.
+func Open(dialect, source string) (*gorm.DB, error) {
+	db, err := sqltraced.Open(dialect, source)
+	if err != nil {
+		return nil, err
+	}
+	return gorm.Open(dialect, db)
+}

contrib/jinzhu/gorm/gorm_test.go

@@ -0,0 +1,78 @@
+package gorm
+
+import (
+	"fmt"
+	"log"
+	"os"
+	"testing"
+
+	sqltrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
+	"gopkg.in/DataDog/dd-trace-go.v1/contrib/internal/sqltest"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+
+	"github.com/go-sql-driver/mysql"
+	"github.com/lib/pq"
+)
+
+// tableName holds the SQL table that these tests will be run against. It must be unique cross-repo.
+const tableName = "testgorm"
+
+func TestMain(m *testing.M) {
+	_, ok := os.LookupEnv("INTEGRATION")
+	if !ok {
+		fmt.Println("--- SKIP: to enable integration test, set the INTEGRATION environment variable")
+		os.Exit(0)
+	}
+	defer sqltest.Prepare(tableName)()
+	os.Exit(m.Run())
+}
+
+func TestMySQL(t *testing.T) {
+	sqltrace.Register("mysql", &mysql.MySQLDriver{}, sqltrace.WithServiceName("mysql-test"))
+	db, err := Open("mysql", "test:test@tcp(127.0.0.1:3306)/test")
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer db.Close()
+
+	testConfig := &sqltest.Config{
+		DB:         db.DB(),
+		DriverName: "mysql",
+		TableName:  tableName,
+		ExpectName: "mysql.query",
+		ExpectTags: map[string]interface{}{
+			ext.ServiceName: "mysql-test",
+			ext.SpanType:    ext.SpanTypeSQL,
+			ext.TargetHost:  "127.0.0.1",
+			ext.TargetPort:  "3306",
+			"db.user":       "test",
+			"db.name":       "test",
+		},
+	}
+	sqltest.RunAll(t, testConfig)
+}
+
+func TestPostgres(t *testing.T) {
+	sqltrace.Register("postgres", &pq.Driver{})
+	db, err := Open("postgres", "postgres://postgres:postgres@127.0.0.1:5432/postgres?sslmode=disable")
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer db.Close()
+
+	testConfig := &sqltest.Config{
+		DB:         db.DB(),
+		DriverName: "postgres",
+		TableName:  tableName,
+		ExpectName: "postgres.query",
+		ExpectTags: map[string]interface{}{
+			ext.ServiceName: "postgres.db",
+			ext.SpanType:    ext.SpanTypeSQL,
+			ext.TargetHost:  "127.0.0.1",
+			ext.TargetPort:  "5432",
+			"db.user":       "postgres",
+			"db.name":       "postgres",
+		},
+	}
+	sqltest.RunAll(t, testConfig)
+}