query retries (#77)

query retries

Author: scottlepp

datasource.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net/http"
 	"sync"
+	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
@@ -224,13 +225,22 @@ func (ds *SQLDatasource) handleQuery(ctx context.Context, req backend.DataQuery,
 	// If there's a query error that didn't exceed the
 	// context deadline retry the query
 	if errors.Is(err, ErrorQuery) && !errors.Is(err, context.DeadlineExceeded) {
-		db, err := ds.c.Connect(dbConn.settings, q.ConnectionArgs)
-		if err != nil {
-			return nil, err
-		}
-		ds.storeDBConnection(cacheKey, dbConnection{db, dbConn.settings})
+		for i := 0; i < ds.driverSettings.Retries; i++ {
+			backend.Logger.Warn(fmt.Sprintf("query failed. retrying %d times", i))
+			db, err := ds.c.Connect(dbConn.settings, q.ConnectionArgs)
+			if err != nil {
+				return nil, err
+			}
+			ds.storeDBConnection(cacheKey, dbConnection{db, dbConn.settings})
 
-		return QueryDB(ctx, db, ds.c.Converters(), fillMode, q)
+			if ds.driverSettings.Pause > 0 {
+				time.Sleep(time.Duration(ds.driverSettings.Pause * int(time.Second)))
+			}
+			res, err = QueryDB(ctx, db, ds.c.Converters(), fillMode, q)
+			if err == nil {
+				return res, err
+			}
+		}
 	}
 
 	// allow retries on timeouts
@@ -243,7 +253,10 @@ func (ds *SQLDatasource) handleQuery(ctx context.Context, req backend.DataQuery,
 			}
 			ds.storeDBConnection(cacheKey, dbConnection{db, dbConn.settings})
 
-			return QueryDB(ctx, db, ds.c.Converters(), fillMode, q)
+			res, err = QueryDB(ctx, db, ds.c.Converters(), fillMode, q)
+			if err == nil {
+				return res, err
+			}
 		}
 	}
 

datasource_test.go

@@ -3,13 +3,16 @@ package sqlds
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"testing"
 	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
 	"github.com/grafana/sqlds/v2/mock"
 	"github.com/stretchr/testify/assert"
 )
@@ -24,6 +27,14 @@ func (d fakeDriver) Connect(backend.DataSourceInstanceSettings, json.RawMessage)
 	return d.db, nil
 }
 
+func (d fakeDriver) Macros() Macros {
+	return Macros{}
+}
+
+func (d fakeDriver) Converters() []sqlutil.Converter {
+	return []sqlutil.Converter{}
+}
+
 // func (d fakeDriver) Settings(backend.DataSourceInstanceSettings) DriverSettings
 
 func Test_getDBConnectionFromQuery(t *testing.T) {
@@ -121,7 +132,7 @@ func Test_Dispose(t *testing.T) {
 	})
 }
 
-func Test_retries(t *testing.T) {
+func Test_timeout_retries(t *testing.T) {
 	dsUID := "timeout"
 	settings := backend.DataSourceInstanceSettings{UID: dsUID}
 
@@ -157,14 +168,67 @@ func Test_retries(t *testing.T) {
 	assert.Equal(t, expected, result.Message)
 }
 
+func Test_error_retries(t *testing.T) {
+	testCounter = 0
+	dsUID := "error"
+	settings := backend.DataSourceInstanceSettings{UID: dsUID}
+
+	handler := testSqlHandler{
+		error: errors.New("foo"),
+	}
+	mockDriver := "sqlmock-error"
+	mock.RegisterDriver(mockDriver, handler)
+	db, err := sql.Open(mockDriver, "")
+	if err != nil {
+		t.Errorf("failed to connect to mock driver: %v", err)
+	}
+	timeoutDriver := fakeDriver{
+		db: db,
+	}
+	retries := 5
+	max := time.Duration(10) * time.Second
+	driverSettings := DriverSettings{Retries: retries, Timeout: max, Pause: 1}
+	ds := &SQLDatasource{c: timeoutDriver, driverSettings: driverSettings}
+
+	key := defaultKey(dsUID)
+	// Add the mandatory default db
+	ds.storeDBConnection(key, dbConnection{db, settings})
+	ctx := context.Background()
+
+	qry := `{ "rawSql": "foo" }`
+
+	req := &backend.QueryDataRequest{
+		PluginContext: backend.PluginContext{
+			DataSourceInstanceSettings: &settings,
+		},
+		Queries: []backend.DataQuery{
+			{
+				RefID: "foo",
+				JSON:  []byte(qry),
+			},
+		},
+	}
+
+	data, err := ds.QueryData(ctx, req)
+	assert.Nil(t, err)
+	assert.Equal(t, retries+1, testCounter)
+	assert.NotNil(t, data.Responses)
+
+}
+
 var testCounter = 0
 var testTimeout = 1
+var testRows = 0
 
 type testSqlHandler struct {
 	mock.DBHandler
+	error
 }
 
 func (s testSqlHandler) Ping(ctx context.Context) error {
+	if s.error != nil {
+		return s.error
+	}
 	testCounter++                              // track the retries for the test assertion
 	time.Sleep(time.Duration(testTimeout + 1)) // simulate a connection delay
 	select {
@@ -173,3 +237,30 @@ func (s testSqlHandler) Ping(ctx context.Context) error {
 		return ctx.Err()
 	}
 }
+
+func (s testSqlHandler) Query(args []driver.Value) (driver.Rows, error) {
+	fmt.Println("query")
+	if s.error != nil {
+		testCounter++
+		return s, s.error
+	}
+	return s, nil
+}
+
+func (s testSqlHandler) Columns() []string {
+	return []string{"foo", "bar"}
+}
+
+func (s testSqlHandler) Next(dest []driver.Value) error {
+	testRows++
+	if testRows > 5 {
+		return io.EOF
+	}
+	dest[0] = "foo"
+	dest[1] = "bar"
+	return nil
+}
+
+func (s testSqlHandler) Close() error {
+	return nil
+}

driver.go

@@ -15,6 +15,7 @@ type DriverSettings struct {
 	Timeout  time.Duration
 	FillMode *data.FillMissing
 	Retries  int
+	Pause    int
 }
 
 // Driver is a simple interface that defines how to connect to a backend SQL datasource