trace/sql: Wrap and return standard library sql.DB (beatlabs#178)

Signed-off-by: Stratos Neiros <[email protected]>

trace/sql/sql.go

@@ -140,6 +140,18 @@ func OpenDB(c driver.Connector) *DB {
 	return &DB{db: db}
 }
 
+// FromDB wraps an opened db. This allows to support libraries that provide
+// *sql.DB like sqlmock.
+func FromDB(db *sql.DB) *DB {
+	return &DB{db: db}
+}
+
+// DB returns the underlying db. This is useful for SQL code that does not
+// require tracing.
+func (db *DB) DB() *sql.DB {
+	return db.db
+}
+
 // BeginTx starts a transaction.
 func (db *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*Tx, error) {
 	sp, _ := db.startSpan(ctx, "db.BeginTx", "")

trace/sql/sql_test.go

@@ -2,6 +2,7 @@ package sql
 
 import (
 	"context"
+	"database/sql"
 	"testing"
 
 	"github.com/beatlabs/patron/trace"
@@ -63,3 +64,10 @@ func TestSQLStartFinishSpan(t *testing.T) {
 		"key":          "value",
 	}, rawSpan.Tags())
 }
+
+func TestFromDB(t *testing.T) {
+	want := &sql.DB{}
+	db := FromDB(want)
+	got := db.DB()
+	assert.Equal(t, want, got)
+}