From 69990e175cd02d7aaedb1ddac93f248720450011 Mon Sep 17 00:00:00 2001
From: Stratos Neiros <nstratos@gmail.com>
Date: Fri, 13 Mar 2020 19:27:05 +0200
Subject: [PATCH] trace/sql: Wrap and return standard library sql.DB (#178)

Signed-off-by: Stratos Neiros <s.neiros@thebeat.co>
---
 trace/sql/sql.go      | 12 ++++++++++++
 trace/sql/sql_test.go |  8 ++++++++
 2 files changed, 20 insertions(+)

diff --git a/trace/sql/sql.go b/trace/sql/sql.go
index 281ac71ef..140f99ce1 100644
--- a/trace/sql/sql.go
+++ b/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", "")
diff --git a/trace/sql/sql_test.go b/trace/sql/sql_test.go
index dd1724483..78625b90a 100644
--- a/trace/sql/sql_test.go
+++ b/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)
+}