Add query args to be able to modify the current DB (#30)

Author: andresmgot

datasource.go

@@ -14,27 +14,32 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/data"
 )
 
+const defaultKey = "_default"
+
 type sqldatasource struct {
 	Completable
 
-	db *sql.DB
-	c  Driver
-
+	dbConnections  sync.Map
+	c              Driver
 	driverSettings DriverSettings
 	settings       backend.DataSourceInstanceSettings
 
 	backend.CallResourceHandler
 	CustomRoutes map[string]func(http.ResponseWriter, *http.Request)
+	// Enabling multiple connections may cause that concurrent connection limits
+	// are hit. The datasource enabling this should make sure connections are cached
+	// if necessary.
+	EnableMultipleConnections bool
 }
 
 // NewDatasource creates a new `sqldatasource`.
 // It uses the provided settings argument to call the ds.Driver to connect to the SQL server
 func (ds *sqldatasource) NewDatasource(settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
-	db, err := ds.c.Connect(settings)
+	db, err := ds.c.Connect(settings, nil)
 	if err != nil {
 		return nil, err
 	}
-	ds.db = db
+	ds.dbConnections.Store(defaultKey, db)
 	ds.settings = settings
 	mux := http.NewServeMux()
 	err = ds.registerRoutes(mux)
@@ -57,7 +62,14 @@ func NewDatasource(c Driver) *sqldatasource {
 
 // Dispose cleans up datasource instance resources.
 func (ds *sqldatasource) Dispose() {
-	ds.db.Close()
+	ds.dbConnections.Range(func(key, db interface{}) bool {
+		err := db.(*sql.DB).Close()
+		if err != nil {
+			backend.Logger.Error(err.Error())
+		}
+		ds.dbConnections.Delete(key)
+		return true
+	})
 }
 
 // QueryData creates the Responses list and executes each query
@@ -88,6 +100,34 @@ func (ds *sqldatasource) QueryData(ctx context.Context, req *backend.QueryDataRe
 
 }
 
+func (ds *sqldatasource) getDB(q *Query) (*sql.DB, string, error) {
+	// The database connection may vary depending on query arguments
+	// The raw arguments are used as key to store the db connection in memory so they can be reused
+	key := defaultKey
+	db, ok := ds.dbConnections.Load(key)
+	if !ok {
+		return nil, "", fmt.Errorf("unable to get default db connection")
+	}
+	if !ds.EnableMultipleConnections || len(q.ConnectionArgs) == 0 {
+		return db.(*sql.DB), key, nil
+	}
+
+	key = string(q.ConnectionArgs)
+	if cachedDB, ok := ds.dbConnections.Load(key); ok {
+		return cachedDB.(*sql.DB), key, nil
+	}
+
+	var err error
+	db, err = ds.c.Connect(ds.settings, q.ConnectionArgs)
+	if err != nil {
+		return nil, "", err
+	}
+	// Assign this connection in the cache
+	ds.dbConnections.Store(key, db)
+
+	return db.(*sql.DB), key, nil
+}
+
 // handleQuery will call query, and attempt to reconnect if the query failed
 func (ds *sqldatasource) handleQuery(ctx context.Context, req backend.DataQuery) (data.Frames, error) {
 	// Convert the backend.DataQuery into a Query object
@@ -108,6 +148,12 @@ func (ds *sqldatasource) handleQuery(ctx context.Context, req backend.DataQuery)
 		fillMode = q.FillMissing
 	}
 
+	// Retrieve the database connection
+	db, cacheKey, err := ds.getDB(q)
+	if err != nil {
+		return nil, err
+	}
+
 	if ds.driverSettings.Timeout != 0 {
 		tctx, cancel := context.WithTimeout(ctx, ds.driverSettings.Timeout)
 		defer cancel()
@@ -119,7 +165,7 @@ func (ds *sqldatasource) handleQuery(ctx context.Context, req backend.DataQuery)
 	//  * Some datasources (snowflake) expire connections or have an authentication token that expires if not used in 1 or 4 hours.
 	//    Because the datasource driver does not include an option for permanent connections, we retry the connection
 	//    if the query fails. NOTE: this does not include some errors like "ErrNoRows"
-	res, err := query(ctx, ds.db, ds.c.Converters(), fillMode, q)
+	res, err := query(ctx, db, ds.c.Converters(), fillMode, q)
 	if err == nil {
 		return res, nil
 	}
@@ -129,20 +175,25 @@ func (ds *sqldatasource) handleQuery(ctx context.Context, req backend.DataQuery)
 	}
 
 	if errors.Is(err, ErrorQuery) {
-		ds.db, err = ds.c.Connect(ds.settings)
+		db, err = ds.c.Connect(ds.settings, q.ConnectionArgs)
 		if err != nil {
 			return nil, err
 		}
+		ds.dbConnections.Store(cacheKey, db)
 
-		return query(ctx, ds.db, ds.c.Converters(), fillMode, q)
+		return query(ctx, db, ds.c.Converters(), fillMode, q)
 	}
 
 	return nil, err
 }
 
 // CheckHealth pings the connected SQL database
 func (ds *sqldatasource) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
-	if err := ds.db.Ping(); err != nil {
+	db, ok := ds.dbConnections.Load(defaultKey)
+	if !ok {
+		return nil, fmt.Errorf("unable to get default db connection")
+	}
+	if err := db.(*sql.DB).Ping(); err != nil {
 		return &backend.CheckHealthResult{
 			Status:  backend.HealthStatusError,
 			Message: err.Error(),

datasource_test.go

@@ -0,0 +1,108 @@
+package sqlds
+
+import (
+	"database/sql"
+	"encoding/json"
+	"testing"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+)
+
+type fakeDriver struct {
+	db *sql.DB
+
+	Driver
+}
+
+func (d *fakeDriver) Connect(backend.DataSourceInstanceSettings, json.RawMessage) (db *sql.DB, err error) {
+	return d.db, nil
+}
+
+func Test_getDB(t *testing.T) {
+	db := &sql.DB{}
+	db2 := &sql.DB{}
+	d := &fakeDriver{db: db2}
+	tests := []struct {
+		desc       string
+		args       string
+		existingDB *sql.DB
+		expectedDB *sql.DB
+	}{
+		{
+			"it should return the default db with no args",
+			defaultKey,
+			db,
+			db,
+		},
+		{
+			"it should return the cached connection for the given args",
+			"foo",
+			db,
+			db,
+		},
+		{
+			"it should create a new connection with the given args",
+			"foo",
+			nil,
+			db2,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			ds := &sqldatasource{c: d, EnableMultipleConnections: true}
+			if tt.existingDB != nil {
+				ds.dbConnections.Store(tt.args, tt.existingDB)
+			}
+			if tt.args != defaultKey {
+				// Add the mandatory default db
+				ds.dbConnections.Store(defaultKey, db)
+			}
+			res, key, err := ds.getDB(&Query{ConnectionArgs: json.RawMessage(tt.args)})
+			if err != nil {
+				t.Fatalf("unexpected error %v", err)
+			}
+			if key != tt.args {
+				t.Fatalf("unexpected cache key %s", key)
+			}
+			if res != tt.expectedDB {
+				t.Fatalf("unexpected result %v", res)
+			}
+		})
+	}
+}
+
+func Test_Dispose(t *testing.T) {
+	t.Run("it should close all db connections", func(t *testing.T) {
+		db1, mock1, err := sqlmock.New()
+		if err != nil {
+			t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
+		}
+		db2, mock2, err := sqlmock.New()
+		if err != nil {
+			t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
+		}
+
+		ds := &sqldatasource{}
+		ds.dbConnections.Store(defaultKey, db1)
+		ds.dbConnections.Store("foo", db2)
+
+		mock1.ExpectClose()
+		mock2.ExpectClose()
+		ds.Dispose()
+
+		err = mock1.ExpectationsWereMet()
+		if err != nil {
+			t.Error(err)
+		}
+		err = mock2.ExpectationsWereMet()
+		if err != nil {
+			t.Error(err)
+		}
+
+		ds.dbConnections.Range(func(key, value interface{}) bool {
+			t.Errorf("db connections were not deleted")
+			return false
+		})
+	})
+}

driver.go

@@ -2,6 +2,7 @@ package sqlds
 
 import (
 	"database/sql"
+	"encoding/json"
 	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -18,7 +19,7 @@ type DriverSettings struct {
 // Plugin creators will need to implement this in order to create a managed datasource
 type Driver interface {
 	// Connect connects to the database. It does not need to call `db.Ping()`
-	Connect(backend.DataSourceInstanceSettings) (*sql.DB, error)
+	Connect(backend.DataSourceInstanceSettings, json.RawMessage) (*sql.DB, error)
 	// Settings are read whenever the plugin is initialized, or after the data source settings are updated
 	Settings(backend.DataSourceInstanceSettings) DriverSettings
 	Macros() Macros

go.mod

@@ -3,6 +3,7 @@ module github.com/grafana/sqlds
 go 1.15
 
 require (
+	github.com/DATA-DOG/go-sqlmock v1.5.0
 	github.com/grafana/grafana-plugin-sdk-go v0.94.0
 	github.com/stretchr/testify v1.7.0
 )

go.sum

@@ -1,6 +1,8 @@
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
+github.com/DATA-DOG/go-sqlmock v1.5.0 h1:Shsta01QNfFxHCfpW6YH2STWB0MudeXXEWMr20OEh60=
+github.com/DATA-DOG/go-sqlmock v1.5.0/go.mod h1:f/Ixk793poVmq4qj/V1dPUg2JEAKC73Q5eFN3EC/SaM=
 github.com/Knetic/govaluate v3.0.1-0.20171022003610-9aa49832a739+incompatible/go.mod h1:r7JcOSlj0wfOMncg0iLm8Leh48TZaKVeNIfJntJ2wa0=
 github.com/Shopify/sarama v1.19.0/go.mod h1:FVkBWblsNy7DGZRfXLU0O9RCGt5g3g3yEuWXgklEdEo=
 github.com/Shopify/toxiproxy v2.1.4+incompatible/go.mod h1:OXgGpZ6Cli1/URJOF1DMxUHB2q5Ap20/P/eIdh4G0pI=

macros_test.go

@@ -1,29 +1,17 @@
 package sqlds
 
 import (
-	"database/sql"
 	"fmt"
 	"testing"
 	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
-	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-type MockDB struct{}
-
-func (h *MockDB) Connect(backend.DataSourceInstanceSettings) (db *sql.DB, err error) {
-	return
-}
-
-func (h *MockDB) Settings(backend.DataSourceInstanceSettings) (settings DriverSettings) {
-	return
-}
-
-func (h *MockDB) Converters() (sc []sqlutil.Converter) {
-	return
+type MockDB struct {
+	Driver
 }
 
 func (h *MockDB) Macros() (macros Macros) {

query.go

@@ -26,8 +26,9 @@ const (
 // For the sake of backwards compatibility, when making changes to this type, ensure that changes are
 // only additive.
 type Query struct {
-	RawSQL string            `json:"rawSql"`
-	Format FormatQueryOption `json:"format"`
+	RawSQL         string            `json:"rawSql"`
+	Format         FormatQueryOption `json:"format"`
+	ConnectionArgs json.RawMessage   `json:"connectionArgs"`
 
 	RefID         string            `json:"-"`
 	Interval      time.Duration     `json:"-"`
@@ -45,15 +46,16 @@ type Query struct {
 // This is mostly useful in the Interpolate function, where the RawSQL value is modified in a loop
 func (q *Query) WithSQL(query string) *Query {
 	return &Query{
-		RawSQL:        query,
-		RefID:         q.RefID,
-		Interval:      q.Interval,
-		TimeRange:     q.TimeRange,
-		MaxDataPoints: q.MaxDataPoints,
-		FillMissing:   q.FillMissing,
-		Schema:        q.Schema,
-		Table:         q.Table,
-		Column:        q.Column,
+		RawSQL:         query,
+		ConnectionArgs: q.ConnectionArgs,
+		RefID:          q.RefID,
+		Interval:       q.Interval,
+		TimeRange:      q.TimeRange,
+		MaxDataPoints:  q.MaxDataPoints,
+		FillMissing:    q.FillMissing,
+		Schema:         q.Schema,
+		Table:          q.Table,
+		Column:         q.Column,
 	}
 }
 
@@ -67,16 +69,17 @@ func GetQuery(query backend.DataQuery) (*Query, error) {
 
 	// Copy directly from the well typed query
 	return &Query{
-		RawSQL:        model.RawSQL,
-		Format:        model.Format,
-		RefID:         query.RefID,
-		Interval:      query.Interval,
-		TimeRange:     query.TimeRange,
-		MaxDataPoints: query.MaxDataPoints,
-		FillMissing:   model.FillMissing,
-		Schema:        model.Schema,
-		Table:         model.Table,
-		Column:        model.Column,
+		RawSQL:         model.RawSQL,
+		Format:         model.Format,
+		ConnectionArgs: model.ConnectionArgs,
+		RefID:          query.RefID,
+		Interval:       query.Interval,
+		TimeRange:      query.TimeRange,
+		MaxDataPoints:  query.MaxDataPoints,
+		FillMissing:    model.FillMissing,
+		Schema:         model.Schema,
+		Table:          model.Table,
+		Column:         model.Column,
 	}, nil
 }