Merge pull request #22 from upfluence/am/register-driver-wrapper

sqlutil: Add public API to register driver wrapper

Author: jlevesy

sqlutil/open.go

@@ -19,11 +19,18 @@ var (
 
 	ErrNoDBProvided = errors.New("sql/sqlutil: No DB provided")
 
-	driverWrapperMu = &sync.Mutex{}
-	driverWrapper   = map[string]driverWrapperFunc{"postgres": postgres.NewDB}
+	driverWrappersMu = &sync.Mutex{}
+	driverWrappers   = map[string]DriverWrapperFunc{"postgres": postgres.NewDB}
 )
 
-type driverWrapperFunc func(sql.DB, sqlparser.SQLParser) sql.DB
+func RegisterDriverWrapper(d string, fn DriverWrapperFunc) {
+	driverWrappersMu.Lock()
+	defer driverWrappersMu.Unlock()
+
+	driverWrappers[d] = fn
+}
+
+type DriverWrapperFunc func(sql.DB, sqlparser.SQLParser) sql.DB
 
 type DBOption func(*dbInput)
 
@@ -84,7 +91,7 @@ func (i *dbInput) buildDB(p sqlparser.SQLParser) (sql.DB, error) {
 
 	db := simple.FromStdDB(plainDB, i.driver)
 
-	if wfn, ok := driverWrapper[i.driver]; ok {
+	if wfn, ok := driverWrappers[i.driver]; ok {
 		db = wfn(db, p)
 	}
 

sqlutil/open_cgo.go

@@ -9,9 +9,7 @@ import (
 )
 
 func init() {
-	driverWrapperMu.Lock()
-	driverWrapper["sqlite3"] = newSQLite3DB
-	driverWrapperMu.Unlock()
+	RegisterDriverWrapper("sqlite3", newSQLite3DB)
 }
 
 func newSQLite3DB(db sql.DB, _ sqlparser.SQLParser) sql.DB {

sqlutil/open_test.go

@@ -1,8 +1,11 @@
 package sqlutil
 
 import (
+	"database/sql"
 	"testing"
 
+	"github.com/lib/pq"
+
 	"github.com/upfluence/sql/backend/postgres"
 )
 
@@ -17,3 +20,18 @@ func TestOpenPostgresDB(t *testing.T) {
 		t.Errorf("invalid wrapping of the DB")
 	}
 }
+
+func TestRegisterDriverWrapper(t *testing.T) {
+	sql.Register("bizbuz", &pq.Driver{})
+	RegisterDriverWrapper("bizbuz", postgres.NewDB)
+
+	db, err := Open(WithMaster("bizbuz", "foobar"))
+
+	if err != nil {
+		t.Errorf("Open() = (_, %+v) wanted nil", err)
+	}
+
+	if !postgres.IsPostgresDB(db) {
+		t.Errorf("invalid wrapping of the DB")
+	}
+}

x/migration/driver.go

@@ -1,27 +1,41 @@
 package migration
 
+import "sync"
+
 var (
 	defaultDriver = &driver{name: "default"}
 
-	driverMap = map[string]Driver{
-		"postgres": &driver{
-			name:       "postgres",
-			extensions: []string{"postgres", "psql"},
-		},
+	PostgresDriver Driver = &driver{
+		name:       "postgres",
+		extensions: []string{"postgres", "psql"},
+	}
+
+	driversMu = &sync.Mutex{}
+	drivers   = map[string]Driver{
+		"postgres": PostgresDriver,
 		"sqlite3": &driver{
 			name:       "sqlite3",
 			extensions: []string{"sqlite3", "sqlite"},
 		},
 	}
 )
 
+func RegisterDriver(n string, d Driver) {
+	driversMu.Lock()
+	defer driversMu.Unlock()
+
+	drivers[n] = d
+}
+
 type Driver interface {
 	Name() string
 	Extensions() []string
 }
 
 func fetchDriver(dname string) Driver {
-	d, ok := driverMap[dname]
+	driversMu.Lock()
+	d, ok := drivers[dname]
+	driversMu.Unlock()
 
 	if ok {
 		return d

x/migration/driver_test.go

@@ -0,0 +1,19 @@
+package migration
+
+import "testing"
+
+func TestRegisterDriver(t *testing.T) {
+	d := fetchDriver("foo")
+
+	if n := d.Name(); n != "default" {
+		t.Errorf("driver.Name() = %q [ want: default ]", n)
+	}
+
+	RegisterDriver("foo", PostgresDriver)
+
+	d = fetchDriver("foo")
+
+	if n := d.Name(); n != "postgres" {
+		t.Errorf("driver.Name() = %q [ want: postgres ]", n)
+	}
+}