From 1af976f7d5effee3c9c608a46ed14f4d89792fef Mon Sep 17 00:00:00 2001
From: Scott Lepper <scott.lepper@gmail.com>
Date: Mon, 12 Dec 2022 17:50:09 -0500
Subject: [PATCH 1/2] allow setting messages to retry on

---
 datasource.go | 52 ++++++++++++++++++++++++++++++++++++++-------------
 driver.go     |  1 +
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/datasource.go b/datasource.go
index d9d14d7..fc6f906 100644
--- a/datasource.go
+++ b/datasource.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"strings"
 	"sync"
 	"time"
 
@@ -225,19 +226,26 @@ 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) {
-		for i := 0; i < ds.driverSettings.Retries; i++ {
-			backend.Logger.Warn(fmt.Sprintf("query failed. retrying %d times", i))
-			db, err := ds.dbReconnect(dbConn, q, cacheKey)
-			if err != nil {
-				return nil, err
-			}
-
-			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
+		// only retry on messages that contain specific errors
+		if shouldRetry(ds.driverSettings.RetryOn, err.Error()) {
+			for i := 0; i < ds.driverSettings.Retries; i++ {
+				backend.Logger.Warn(fmt.Sprintf("query failed: %s. Retrying %d times", err.Error(), i))
+				db, err := ds.dbReconnect(dbConn, q, cacheKey)
+				if err != nil {
+					return nil, err
+				}
+
+				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
+				}
+				if !shouldRetry(ds.driverSettings.RetryOn, err.Error()) {
+					return res, err
+				}
+				backend.Logger.Warn(fmt.Sprintf("Retry failed: %s", err.Error()))
 			}
 		}
 	}
@@ -302,6 +310,15 @@ func (ds *SQLDatasource) checkWithRetries(conn dbConnection) (*backend.CheckHeal
 		if err == nil {
 			return result, err
 		}
+
+		if !shouldRetry(ds.driverSettings.RetryOn, err.Error()) {
+			break
+		}
+
+		if ds.driverSettings.Pause > 0 {
+			time.Sleep(time.Duration(ds.driverSettings.Pause * int(time.Second)))
+		}
+		backend.Logger.Warn(fmt.Sprintf("connect failed: %s. Retrying %d times", err.Error(), i))
 	}
 
 	// TODO: failed health checks don't return an error
@@ -332,3 +349,12 @@ func (ds *SQLDatasource) ping(conn dbConnection) error {
 
 	return conn.db.PingContext(ctx)
 }
+
+func shouldRetry(retryOn []string, err string) bool {
+	for _, r := range retryOn {
+		if strings.Contains(err, r) {
+			return true
+		}
+	}
+	return false
+}
diff --git a/driver.go b/driver.go
index 3fe1b77..92f4dac 100644
--- a/driver.go
+++ b/driver.go
@@ -16,6 +16,7 @@ type DriverSettings struct {
 	FillMode *data.FillMissing
 	Retries  int
 	Pause    int
+	RetryOn  []string
 }
 
 // Driver is a simple interface that defines how to connect to a backend SQL datasource

From 70ede6a9de37f48bceee80e116b7aa4121300092 Mon Sep 17 00:00:00 2001
From: Scott Lepper <scott.lepper@gmail.com>
Date: Mon, 12 Dec 2022 17:58:06 -0500
Subject: [PATCH 2/2] fix tests

---
 datasource_test.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/datasource_test.go b/datasource_test.go
index 9615b9d..c5f775d 100644
--- a/datasource_test.go
+++ b/datasource_test.go
@@ -148,7 +148,7 @@ func Test_timeout_retries(t *testing.T) {
 	}
 	retries := 5
 	max := time.Duration(testTimeout) * time.Second
-	driverSettings := DriverSettings{Retries: retries, Timeout: max}
+	driverSettings := DriverSettings{Retries: retries, Timeout: max, RetryOn: []string{"deadline"}}
 	ds := &SQLDatasource{c: timeoutDriver, driverSettings: driverSettings}
 
 	key := defaultKey(dsUID)
@@ -190,7 +190,7 @@ func Test_error_retries(t *testing.T) {
 	}
 	retries := 5
 	max := time.Duration(10) * time.Second
-	driverSettings := DriverSettings{Retries: retries, Timeout: max, Pause: 1}
+	driverSettings := DriverSettings{Retries: retries, Timeout: max, Pause: 1, RetryOn: []string{"foo"}}
 	ds := &SQLDatasource{c: timeoutDriver, driverSettings: driverSettings}
 
 	key := defaultKey(dsUID)