Add health check for async queries (#124)

Add health check for async queries

Author: sarahzinger

go.mod

@@ -75,6 +75,7 @@ require (
 	github.com/prometheus/common v0.46.0 // indirect
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
+	github.com/stretchr/objx v0.5.0 // indirect
 	github.com/unknwon/bra v0.0.0-20200517080246-1e3013ecaff8 // indirect
 	github.com/unknwon/com v1.0.1 // indirect
 	github.com/unknwon/log v0.0.0-20150304194804-e617c87089d3 // indirect

go.sum

@@ -227,6 +227,7 @@ github.com/smartystreets/goconvey v0.0.0-20181108003508-044398e4856c/go.mod h1:X
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=

pkg/awsds/asyncDatasource.go

@@ -130,6 +130,29 @@ func (ds *AsyncAWSDatasource) QueryData(ctx context.Context, req *backend.QueryD
 	return response.Response(), nil
 }
 
+func (ds *AsyncAWSDatasource) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
+	datasourceUID := req.PluginContext.DataSourceInstanceSettings.UID
+	key := defaultKey(datasourceUID)
+	dbConn, ok := ds.getDBConnection(key)
+	if !ok {
+		return &backend.CheckHealthResult{
+			Status:  backend.HealthStatusError,
+			Message: "No database connection found for datasource uid: " + datasourceUID,
+		}, nil
+	}
+	err := dbConn.db.Ping(ctx)
+	if err != nil {
+		return &backend.CheckHealthResult{
+			Status:  backend.HealthStatusError,
+			Message: err.Error(),
+		}, nil
+	}
+	return &backend.CheckHealthResult{
+		Status:  backend.HealthStatusOk,
+		Message: "Data source is working",
+	}, nil
+}
+
 func (ds *AsyncAWSDatasource) getAsyncDBFromQuery(q *AsyncQuery, datasourceUID string) (AsyncDB, error) {
 	if !ds.EnableMultipleConnections && len(q.ConnectionArgs) > 0 {
 		return nil, sqlds.ErrorMissingMultipleConnectionsConfig

pkg/awsds/asyncDatasource_test.go

@@ -4,13 +4,16 @@ import (
 	"context"
 	"database/sql/driver"
 	"encoding/json"
+	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/sqlds/v3"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 )
 
 type fakeAsyncDB struct{}
@@ -136,3 +139,91 @@ func Test_Async_QueryData_uses_synchronous_flow_when_header_has_alert_and_expres
 		})
 	}
 }
+
+type MockDB struct {
+	mock.Mock
+}
+
+func (m *MockDB) Ping(context context.Context) error {
+	args := m.Called(context)
+	return args.Error(0)
+}
+
+func (m *MockDB) Begin() (driver.Tx, error) {
+	args := m.Called()
+	return args.Get(0).(driver.Tx), args.Error(1)
+}
+
+func (m *MockDB) CancelQuery(ctx context.Context, queryID string) error {
+	args := m.Called(ctx, queryID)
+	return args.Error(0)
+}
+func (m *MockDB) Close() error {
+	args := m.Called()
+	return args.Error(0)
+}
+func (m *MockDB) GetQueryID(ctx context.Context, query string, args ...interface{}) (bool, string, error) {
+	arg := m.Called(ctx, query, args)
+	return arg.Bool(0), arg.String(1), arg.Error(2)
+}
+func (m *MockDB) GetRows(ctx context.Context, queryID string) (driver.Rows, error) {
+	args := m.Called(ctx, queryID)
+	return args.Get(0).(driver.Rows), args.Error(1)
+}
+func (m *MockDB) Prepare(query string) (driver.Stmt, error) {
+	args := m.Called(query)
+	return args.Get(0).(driver.Stmt), args.Error(1)
+}
+func (m *MockDB) QueryStatus(ctx context.Context, queryID string) (QueryStatus, error) {
+	args := m.Called(ctx, queryID)
+	return args.Get(0).(QueryStatus), args.Error(1)
+}
+func (m *MockDB) StartQuery(ctx context.Context, query string, args ...interface{}) (string, error) {
+	arg := m.Called(ctx, query, args)
+	return arg.String(0), arg.Error(1)
+}
+
+func Test_AsyncDatasource_CheckHealth(t *testing.T) {
+	tests := []struct {
+		desc             string
+		mockPingResponse error
+		expected         *backend.CheckHealthResult
+	}{
+		{
+			desc:             "it returns an error when ping fails",
+			mockPingResponse: fmt.Errorf("your auth wasn't right"),
+			expected: &backend.CheckHealthResult{
+				Status:  backend.HealthStatusError,
+				Message: "your auth wasn't right",
+			},
+		},
+		{
+			desc:             "it returns an ok when the query succeeds",
+			mockPingResponse: nil,
+			expected: &backend.CheckHealthResult{
+				Status:  backend.HealthStatusOk,
+				Message: "Data source is working",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			db := new(MockDB)
+			db.On("Ping", context.Background()).Return(tt.mockPingResponse)
+			dbC := dbConnection{
+				db,
+				backend.DataSourceInstanceSettings{UID: "uid1"},
+			}
+			ds := &AsyncAWSDatasource{dbConnections: sync.Map{}}
+			ds.dbConnections.Store(defaultKey("uid1"), dbC)
+
+			result, err := ds.CheckHealth(context.Background(), &backend.CheckHealthRequest{
+				PluginContext: backend.PluginContext{
+					DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{UID: "uid1"},
+				},
+			})
+			assert.NoError(t, err)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}