add optional dispose method

datasource.go

@@ -6,12 +6,13 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
 	"net/http"
 	"strings"
 	"sync"
 	"time"
 
+	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
+
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/resource/httpadapter"
@@ -111,6 +112,24 @@ func NewDatasource(c Driver) *SQLDatasource {
 // Dispose cleans up datasource instance resources.
 // Note: Called when testing and saving a datasource
 func (ds *SQLDatasource) Dispose() {
+	if ds.driverSettings.Dispose {
+		var connections []dbConnection
+		ds.dbConnections.Range(func(key, value any) bool {
+			connection, ok := value.(dbConnection)
+			if ok {
+				connections = append(connections, connection)
+			}
+			return true
+		})
+		backend.Logger.Debug(fmt.Sprintf("Disposing %d connections", len(connections)))
+		for _, conn := range connections {
+			if err := conn.db.Close(); err != nil {
+				backend.Logger.Warn(fmt.Sprintf("closing connection failed: %s", err.Error()))
+			}
+			k := defaultKey(getDatasourceUID(conn.settings))
+			ds.dbConnections.Delete(k)
+		}
+	}
 }
 
 // QueryData creates the Responses list and executes each query
@@ -150,6 +169,12 @@ func (ds *SQLDatasource) QueryData(ctx context.Context, req *backend.QueryDataRe
 	wg.Wait()
 
 	errs := ds.errors(response)
+
+	// Dispose after alerts and other non-user queries
+	if req.PluginContext.User == nil && ds.driverSettings.Dispose {
+		ds.Dispose()
+	}
+
 	if ds.driverSettings.Errors {
 		return response.Response(), errs
 	}

datasource_test.go

@@ -99,6 +99,17 @@ func Test_no_errors(t *testing.T) {
 	assert.Equal(t, expected, result.Message)
 }
 
+func TestDispose(t *testing.T) {
+	req, _, ds := healthRequest(t, "pass", test.DriverOpts{}, `{ "dispose": true} `)
+	_, err := ds.CheckHealth(context.Background(), &req)
+	assert.Nil(t, err)
+
+	ds.Dispose()
+
+	_, err = ds.CheckHealth(context.Background(), &req)
+	assert.NotNil(t, err)
+}
+
 func queryRequest(t *testing.T, name string, opts test.DriverOpts, cfg string) (*backend.QueryDataRequest, *test.SqlHandler, *sqlds.SQLDatasource) {
 	driver, handler := test.NewDriver(name, test.Data{}, nil, opts)
 	ds := sqlds.NewDatasource(driver)

driver.go

@@ -20,6 +20,7 @@ type DriverSettings struct {
 	RetryOn        []string
 	ForwardHeaders bool
 	Errors         bool
+	Dispose        bool
 }
 
 // Driver is a simple interface that defines how to connect to a backend SQL datasource