Repository: gobuffalo/pop
Pull request #599 (2020-10-12 15:37:02)

Support SQL instrumentation with Open Census, Open Tracing, AWS Xray, Google Stack Driver
Adds options

```
	// UseInstrumentedDriver if set to true uses a wrapper for the underlying driver which exposes tracing
	// information in the Open Tracing, Open Census, Google, and AWS Xray format. This is useful when using
	// tracing with Jaeger, DataDog, Zipkin, or other tracing software.
	UseInstrumentedDriver bool
	// InstrumentedDriverOptions sets the options for the instrumented driver. These options are empty by default meaning
	// that instrumentation is disabled.
	//
	// For more information check out the docs at https://github.com/luna-duclos/instrumentedsql. If you use Open Tracing, these options
	// could looks as follows:
	//
	//		InstrumentedDriverOptions: []instrumentedsql.Opt{instrumentedsql.WithTracer(opentracing.NewTracer(true))}
	//
	// It is also recommended to include `instrumentedsql.WithOmitArgs()` which prevents SQL arguments (e.g. passwords)
	// from being traced or logged.
	InstrumentedDriverOptions []instrumentedsql.Opt
```

to `pop.ConnectionDetails`, allowing to instrument SQL queries using Open Tracing, Open Census, AWS XRay or Google Stack Driver.


aeneasr: @stanislas-m I'm looking into some weird behavior with pgx in combination with instrumentation - please stay tuned before merging :)

Patch:
diff --git a/connection.go b/connection.go
index a6b08f3..df116fb 100644
--- a/connection.go
+++ b/connection.go
@@ -7,7 +7,6 @@ import (
 
 	"github.com/gobuffalo/pop/v5/internal/defaults"
 	"github.com/gobuffalo/pop/v5/internal/randx"
-	"github.com/jmoiron/sqlx"
 	"github.com/pkg/errors"
 )
 
@@ -95,15 +94,12 @@ func (c *Connection) Open() error {
 		return errors.New("invalid connection instance")
 	}
 	details := c.Dialect.Details()
-	driver := c.Dialect.DefaultDriver()
-	if details.Driver != "" {
-		driver = details.Driver
-	}
 
-	db, err := sqlx.Open(driver, c.Dialect.URL())
+	db, err := openPotentiallyInstrumentedConnection(c.Dialect, c.Dialect.URL())
 	if err != nil {
-		return errors.Wrap(err, "could not open database connection")
+		return err
 	}
+
 	db.SetMaxOpenConns(details.Pool)
 	if details.IdlePool != 0 {
 		db.SetMaxIdleConns(details.IdlePool)
diff --git a/connection_details.go b/connection_details.go
index 6a132e1..76c165b 100644
--- a/connection_details.go
+++ b/connection_details.go
@@ -2,6 +2,7 @@ package pop
 
 import (
 	"fmt"
+	"github.com/luna-duclos/instrumentedsql"
 	"net/url"
 	"regexp"
 	"strconv"
@@ -47,6 +48,21 @@ type ConnectionDetails struct {
 	Options map[string]string
 	// Query string encoded options from URL. Example: "sslmode=disable"
 	RawOptions string
+	// UseInstrumentedDriver if set to true uses a wrapper for the underlying driver which exposes tracing
+	// information in the Open Tracing, Open Census, Google, and AWS Xray format. This is useful when using
+	// tracing with Jaeger, DataDog, Zipkin, or other tracing software.
+	UseInstrumentedDriver bool
+	// InstrumentedDriverOptions sets the options for the instrumented driver. These options are empty by default meaning
+	// that instrumentation is disabled.
+	//
+	// For more information check out the docs at https://github.com/luna-duclos/instrumentedsql. If you use Open Tracing, these options
+	// could looks as follows:
+	//
+	//		InstrumentedDriverOptions: []instrumentedsql.Opt{instrumentedsql.WithTracer(opentracing.NewTracer(true))}
+	//
+	// It is also recommended to include `instrumentedsql.WithOmitArgs()` which prevents SQL arguments (e.g. passwords)
+	// from being traced or logged.
+	InstrumentedDriverOptions []instrumentedsql.Opt
 }
 
 var dialectX = regexp.MustCompile(`\S+://`)
diff --git a/connection_instrumented.go b/connection_instrumented.go
new file mode 100644
index 0000000..e27fee2
--- /dev/null
+++ b/connection_instrumented.go
@@ -0,0 +1,84 @@
+package pop
+
+import (
+	"database/sql"
+	"database/sql/driver"
+	mysqld "github.com/go-sql-driver/mysql"
+	"github.com/gobuffalo/pop/v5/logging"
+	pgx "github.com/jackc/pgx/v4/stdlib"
+	"github.com/jmoiron/sqlx"
+	"github.com/luna-duclos/instrumentedsql"
+	"github.com/mattn/go-sqlite3"
+	"github.com/pkg/errors"
+)
+
+const instrumentedDriverName = "instrumented-sql-driver"
+
+func instrumentDriver(deets *ConnectionDetails, defaultDriverName string) (driverName, dialect string) {
+	driverName = defaultDriverName
+	if deets.Driver != "" {
+		driverName = deets.Driver
+	}
+	dialect = driverName
+
+	if !deets.UseInstrumentedDriver {
+		if len(deets.InstrumentedDriverOptions) > 0 {
+			log(logging.Warn, "SQL driver instrumentation is disabled but `ConnectionDetails.InstrumentedDriverOptions` is not empty. Please double-check if this is a error.")
+		}
+
+		// If instrumentation is disabled, we just return the driver name we got (e.g. "pgx").
+		return driverName, dialect
+	}
+
+	if len(deets.InstrumentedDriverOptions) == 0 {
+		log(logging.Warn, "SQL driver instrumentation was enabled but no options have been passed to `ConnectionDetails.InstrumentedDriverOptions`. Instrumentation will therefore not result in any output.")
+	}
+
+	var dr driver.Driver
+	var newDriverName string
+	switch normalizeSynonyms(driverName) {
+	case nameCockroach:
+		fallthrough
+	case namePostgreSQL:
+		dr = new(pgx.Driver)
+		newDriverName = instrumentedDriverName + "-" + namePostgreSQL
+	case nameMariaDB:
+		fallthrough
+	case nameMySQL:
+		dr = mysqld.MySQLDriver{}
+		newDriverName = instrumentedDriverName + "-" + nameMySQL
+	case nameSQLite3:
+		dr = new(sqlite3.SQLiteDriver)
+		newDriverName = instrumentedDriverName + "-" + nameSQLite3
+	}
+
+	var found bool
+	for _, n := range sql.Drivers() {
+		if n == newDriverName {
+			found = true
+			break
+		}
+	}
+
+	if !found {
+		sql.Register(newDriverName, instrumentedsql.WrapDriver(dr, deets.InstrumentedDriverOptions...))
+	}
+
+	return newDriverName, dialect
+}
+
+// openPotentiallyInstrumentedConnection first opens a raw SQL connection and then wraps it with `sqlx`.
+//
+// We do this because `sqlx` needs the database type in order to properly
+// translate arguments (e.g. `?` to `$1`) in SQL queries. Because we use
+// a custom driver name when using instrumentation, this detection would fail
+// otherwise.
+func openPotentiallyInstrumentedConnection(c dialect, dsn string) (*sqlx.DB, error) {
+	driverName, dialect := instrumentDriver(c.Details(), c.DefaultDriver())
+	con, err := sql.Open(driverName, dsn)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not open database connection")
+	}
+
+	return sqlx.NewDb(con, dialect), nil
+}
diff --git a/dialect_cockroach.go b/dialect_cockroach.go
index 61dd8ef..6109476 100644
--- a/dialect_cockroach.go
+++ b/dialect_cockroach.go
@@ -1,7 +1,6 @@
 package pop
 
 import (
-	"database/sql"
 	"fmt"
 	"io"
 	"os"
@@ -123,13 +122,7 @@ func (p *cockroach) CreateDB() error {
 	// createdb -h db -p 5432 -U cockroach enterprise_development
 	deets := p.ConnectionDetails
 
-	// Overwrite dialect to match pgx driver for sql.Open
-	driver := p.DefaultDriver()
-	if p.ConnectionDetails.Driver != "" {
-		driver = p.ConnectionDetails.Driver
-	}
-
-	db, err := sql.Open(driver, p.urlWithoutDb())
+	db, err := openPotentiallyInstrumentedConnection(p, p.urlWithoutDb())
 	if err != nil {
 		return errors.Wrapf(err, "error creating Cockroach database %s", deets.Database)
 	}
@@ -149,13 +142,7 @@ func (p *cockroach) CreateDB() error {
 func (p *cockroach) DropDB() error {
 	deets := p.ConnectionDetails
 
-	// Overwrite dialect to match pgx driver for sql.Open
-	driver := p.DefaultDriver()
-	if p.ConnectionDetails.Driver != "" {
-		driver = p.ConnectionDetails.Driver
-	}
-
-	db, err := sql.Open(driver, p.urlWithoutDb())
+	db, err := openPotentiallyInstrumentedConnection(p, p.urlWithoutDb())
 	if err != nil {
 		return errors.Wrapf(err, "error dropping Cockroach database %s", deets.Database)
 	}
@@ -219,7 +206,7 @@ func (p *cockroach) DumpSchema(w io.Writer) error {
 }
 
 func (p *cockroach) LoadSchema(r io.Reader) error {
-	return genericLoadSchema(p.ConnectionDetails, p.DefaultDriver(), p.MigrationURL(), r)
+	return genericLoadSchema(p, r)
 }
 
 func (p *cockroach) TruncateAll(tx *Connection) error {
diff --git a/dialect_common.go b/dialect_common.go
index 16177dd..b34caae 100644
--- a/dialect_common.go
+++ b/dialect_common.go
@@ -14,7 +14,6 @@ import (
 	"github.com/gobuffalo/pop/v5/columns"
 	"github.com/gobuffalo/pop/v5/logging"
 	"github.com/gofrs/uuid"
-	"github.com/jmoiron/sqlx"
 	"github.com/pkg/errors"
 )
 
@@ -144,13 +143,11 @@ func genericSelectMany(s store, models *Model, query Query) error {
 	return nil
 }
 
-func genericLoadSchema(deets *ConnectionDetails, defaultDriver, migrationURL string, r io.Reader) error {
+func genericLoadSchema(d dialect, r io.Reader) error {
+	deets := d.Details()
+
 	// Open DB connection on the target DB
-	driver := defaultDriver
-	if deets.Driver != "" {
-		driver = deets.Driver
-	}
-	db, err := sqlx.Open(driver, migrationURL)
+	db, err := openPotentiallyInstrumentedConnection(d, d.MigrationURL())
 	if err != nil {
 		return errors.WithMessage(err, fmt.Sprintf("unable to load schema for %s", deets.Database))
 	}
diff --git a/dialect_mysql.go b/dialect_mysql.go
index 510f050..63fb900 100644
--- a/dialect_mysql.go
+++ b/dialect_mysql.go
@@ -2,7 +2,6 @@ package pop
 
 import (
 	"bytes"
-	"database/sql"
 	"fmt"
 	"io"
 	"os/exec"
@@ -106,7 +105,7 @@ func (m *mysql) SelectMany(s store, models *Model, query Query) error {
 // CreateDB creates a new database, from the given connection credentials
 func (m *mysql) CreateDB() error {
 	deets := m.ConnectionDetails
-	db, err := sql.Open(deets.Dialect, m.urlWithoutDb())
+	db, err := openPotentiallyInstrumentedConnection(m, m.urlWithoutDb())
 	if err != nil {
 		return errors.Wrapf(err, "error creating MySQL database %s", deets.Database)
 	}
@@ -128,7 +127,7 @@ func (m *mysql) CreateDB() error {
 // DropDB drops an existing database, from the given connection credentials
 func (m *mysql) DropDB() error {
 	deets := m.ConnectionDetails
-	db, err := sql.Open(deets.Dialect, m.urlWithoutDb())
+	db, err := openPotentiallyInstrumentedConnection(m, m.urlWithoutDb())
 	if err != nil {
 		return errors.Wrapf(err, "error dropping MySQL database %s", deets.Database)
 	}
@@ -165,7 +164,7 @@ func (m *mysql) DumpSchema(w io.Writer) error {
 
 // LoadSchema executes a schema sql file against the configured database.
 func (m *mysql) LoadSchema(r io.Reader) error {
-	return genericLoadSchema(m.ConnectionDetails, m.DefaultDriver(), m.MigrationURL(), r)
+	return genericLoadSchema(m, r)
 }
 
 // TruncateAll truncates all tables for the given connection.
diff --git a/dialect_postgresql.go b/dialect_postgresql.go
index c4ecd2e..3409d4b 100644
--- a/dialect_postgresql.go
+++ b/dialect_postgresql.go
@@ -1,7 +1,6 @@
 package pop
 
 import (
-	"database/sql"
 	"fmt"
 	"io"
 	"os/exec"
@@ -113,12 +112,7 @@ func (p *postgresql) CreateDB() error {
 	// createdb -h db -p 5432 -U postgres enterprise_development
 	deets := p.ConnectionDetails
 
-	driver := p.DefaultDriver()
-	if p.ConnectionDetails.Driver != "" {
-		driver = p.ConnectionDetails.Driver
-	}
-
-	db, err := sql.Open(driver, p.urlWithoutDb())
+	db, err := openPotentiallyInstrumentedConnection(p, p.urlWithoutDb())
 	if err != nil {
 		return errors.Wrapf(err, "error creating PostgreSQL database %s", deets.Database)
 	}
@@ -138,13 +132,7 @@ func (p *postgresql) CreateDB() error {
 func (p *postgresql) DropDB() error {
 	deets := p.ConnectionDetails
 
-	// Overwrite dialect to match pgx driver for sql.Open
-	driver := p.DefaultDriver()
-	if p.ConnectionDetails.Driver != "" {
-		driver = p.ConnectionDetails.Driver
-	}
-
-	db, err := sql.Open(driver, p.urlWithoutDb())
+	db, err := openPotentiallyInstrumentedConnection(p, p.urlWithoutDb())
 	if err != nil {
 		return errors.Wrapf(err, "error dropping PostgreSQL database %s", deets.Database)
 	}
@@ -208,7 +196,7 @@ func (p *postgresql) DumpSchema(w io.Writer) error {
 
 // LoadSchema executes a schema sql file against the configured database.
 func (p *postgresql) LoadSchema(r io.Reader) error {
-	return genericLoadSchema(p.ConnectionDetails, p.DefaultDriver(), p.MigrationURL(), r)
+	return genericLoadSchema(p, r)
 }
 
 // TruncateAll truncates all tables for the given connection.
diff --git a/go.mod b/go.mod
index 9db639e..225a6ac 100644
--- a/go.mod
+++ b/go.mod
@@ -25,6 +25,7 @@ require (
 	github.com/jmoiron/sqlx v1.2.0
 	github.com/karrick/godirwalk v1.16.1 // indirect
 	github.com/lib/pq v1.3.0
+	github.com/luna-duclos/instrumentedsql v1.1.3
 	github.com/mattn/go-sqlite3 v2.0.3+incompatible
 	github.com/pkg/errors v0.9.1
 	github.com/spf13/cobra v0.0.6
diff --git a/go.sum b/go.sum
index 2732087..22100ac 100644
--- a/go.sum
+++ b/go.sum
@@ -175,6 +175,8 @@ github.com/lib/pq v1.1.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/lib/pq v1.2.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/lib/pq v1.3.0 h1:/qkRGz8zljWiDcFvgpwUpwIAPu3r07TDvs3Rws+o/pU=
 github.com/lib/pq v1.3.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
+github.com/luna-duclos/instrumentedsql v1.1.3 h1:t7mvC0z1jUt5A0UQ6I/0H31ryymuQRnJcWCiqV3lSAA=
+github.com/luna-duclos/instrumentedsql v1.1.3/go.mod h1:9J1njvFds+zN7y85EDhN9XNQLANWwZt2ULeIC8yMNYs=
 github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
 github.com/markbates/errx v1.1.0 h1:QDFeR+UP95dO12JgW+tgi2UVfo0V8YBHiUIOaeBPiEI=
 github.com/markbates/errx v1.1.0/go.mod h1:PLa46Oex9KNbVDZhKel8v1OT7hD5JZ2eI7AHhA0wswc=


Test patch:
diff --git a/connection_instrumented_test.go b/connection_instrumented_test.go
new file mode 100644
index 0000000..d512563
--- /dev/null
+++ b/connection_instrumented_test.go
@@ -0,0 +1,100 @@
+package pop
+
+import (
+	"context"
+	"fmt"
+	"github.com/luna-duclos/instrumentedsql"
+	"github.com/stretchr/testify/suite"
+	"os"
+	"strings"
+	"sync"
+	"time"
+)
+
+func testInstrumentedDriver(p *suite.Suite) {
+	r := p.Require()
+	deets := *Connections[os.Getenv("SODA_DIALECT")].Dialect.Details()
+
+	ctx, cancel := context.WithTimeout(context.TODO(), time.Second*5)
+	defer cancel()
+
+	// The WaitGroup and channel ensures that the logger is properly called. This can only happen
+	// when the instrumented driver is working as expected and returns the expected query.
+	var (
+		queryMySQL = "SELECT 1 FROM DUAL WHERE 1=?"
+		queryOther = "SELECT 1 WHERE 1=?"
+		mc         = make(chan string)
+		wg         sync.WaitGroup
+		expected   = []string{
+			"SELECT 1 FROM DUAL WHERE 1=?",
+			"SELECT 1 FROM DUAL WHERE 1=$1",
+			"SELECT 1 WHERE 1=?",
+			"SELECT 1 WHERE 1=$1",
+		}
+	)
+
+	query := queryOther
+	if os.Getenv("SODA_DIALECT") == "mysql" {
+		query = queryMySQL
+	}
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		var messages []string
+		var found bool
+		for {
+			select {
+			case m := <-mc:
+				p.T().Logf("Received message: %s", m)
+				messages = append(messages, m)
+				for _, e := range expected {
+					if strings.Contains(m, e) {
+						p.T().Logf("Found part %s in %s", e, m)
+						found = true
+						break
+					}
+				}
+			case <-ctx.Done():
+				if !found {
+					r.FailNow(fmt.Sprintf("Expected tracer to return the \"%s\" query but only the following messages have been received:\n\n\t%s", query, strings.Join(messages, "\n\t")))
+					return
+				}
+				return
+			}
+		}
+	}()
+
+	var checker = instrumentedsql.LoggerFunc(func(ctx context.Context, msg string, keyvals ...interface{}) {
+		p.T().Logf("Instrumentation received message: %s - %+v", msg, keyvals)
+		mc <- fmt.Sprintf("%s - %+v", msg, keyvals)
+	})
+
+	deets.UseInstrumentedDriver = true
+	deets.InstrumentedDriverOptions = []instrumentedsql.Opt{instrumentedsql.WithLogger(checker)}
+
+	c, err := NewConnection(&deets)
+	r.NoError(err)
+	r.NoError(c.Open())
+
+	err = c.WithContext(context.TODO()).RawQuery(query, 1).Exec()
+	r.NoError(err)
+
+	wg.Wait()
+}
+
+func (s *PostgreSQLSuite) Test_Instrumentation() {
+	testInstrumentedDriver(&s.Suite)
+}
+
+func (s *MySQLSuite) Test_Instrumentation() {
+	testInstrumentedDriver(&s.Suite)
+}
+
+func (s *SQLiteSuite) Test_Instrumentation() {
+	testInstrumentedDriver(&s.Suite)
+}
+
+func (s *CockroachSuite) Test_Instrumentation() {
+	testInstrumentedDriver(&s.Suite)
+}
diff --git a/pop_test.go b/pop_test.go
index a7dab49..33efcbe 100644
--- a/pop_test.go
+++ b/pop_test.go
@@ -29,6 +29,10 @@ type SQLiteSuite struct {
 	suite.Suite
 }
 
+type CockroachSuite struct {
+	suite.Suite
+}
+
 func TestSpecificSuites(t *testing.T) {
 	switch os.Getenv("SODA_DIALECT") {
 	case "postgres":
@@ -37,6 +41,8 @@ func TestSpecificSuites(t *testing.T) {
 		suite.Run(t, &MySQLSuite{})
 	case "sqlite":
 		suite.Run(t, &SQLiteSuite{})
+	case "cockroach":
+		suite.Run(t, &CockroachSuite{})
 	}
 }
 
diff --git a/slices_test.go b/slices_test.go
index e9b343b..4668e4e 100644
--- a/slices_test.go
+++ b/slices_test.go
@@ -15,9 +15,9 @@ type Cake struct {
 	UpdatedAt time.Time     `json:"updated_at" db:"updated_at"`
 }
 
-func (p *PostgreSQLSuite) Test_String() {
+func (s *PostgreSQLSuite) Test_String() {
 	transaction(func(tx *Connection) {
-		r := p.Require()
+		r := s.Require()
 
 		c := &Cake{
 			String: slices.String{"a", "b", "c"},
@@ -31,9 +31,9 @@ func (p *PostgreSQLSuite) Test_String() {
 	})
 }
 
-func (p *PostgreSQLSuite) Test_Int() {
+func (s *PostgreSQLSuite) Test_Int() {
 	transaction(func(tx *Connection) {
-		r := p.Require()
+		r := s.Require()
 
 		c := &Cake{
 			Int: slices.Int{1, 2, 3},
@@ -47,9 +47,9 @@ func (p *PostgreSQLSuite) Test_Int() {
 	})
 }
 
-func (p *PostgreSQLSuite) Test_Float() {
+func (s *PostgreSQLSuite) Test_Float() {
 	transaction(func(tx *Connection) {
-		r := p.Require()
+		r := s.Require()
 
 		c := &Cake{
 			Float: slices.Float{1.0, 2.1, 3.2},