capture duration and label with error source (#116)

* capture duration and label with error source

Author: scottlepp

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"
@@ -239,7 +240,8 @@ 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 := QueryDB(ctx, dbConn.db, ds.c.Converters(), fillMode, q, args...)
+	dbQuery := NewQuery(dbConn.db, dbConn.settings, ds.c.Converters(), fillMode)
+	res, err := dbQuery.Run(ctx, q, args...)
 	if err == nil {
 		return res, nil
 	}
@@ -263,7 +265,9 @@ func (ds *SQLDatasource) handleQuery(ctx context.Context, req backend.DataQuery,
 				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, args...)
+
+				dbQuery := NewQuery(db, dbConn.settings, ds.c.Converters(), fillMode)
+				res, err = dbQuery.Run(ctx, q, args...)
 				if err == nil {
 					return res, err
 				}
@@ -284,7 +288,8 @@ func (ds *SQLDatasource) handleQuery(ctx context.Context, req backend.DataQuery,
 				continue
 			}
 
-			res, err = QueryDB(ctx, db, ds.c.Converters(), fillMode, q, args...)
+			dbQuery := NewQuery(db, dbConn.settings, ds.c.Converters(), fillMode)
+			res, err = dbQuery.Run(ctx, q, args...)
 			if err == nil {
 				return res, err
 			}

query.go

@@ -7,12 +7,15 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/grafana/dataplane/sdata/timeseries"
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/data"
 	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
 // FormatQueryOption defines how the user has chosen to represent the data
@@ -36,6 +39,12 @@ const (
 // Deprecated: use sqlutil.Query directly instead
 type Query = sqlutil.Query
 
+var duration = promauto.NewHistogramVec(prometheus.HistogramOpts{
+	Namespace: "plugins",
+	Name:      "plugin_external_requests_duration_seconds",
+	Help:      "Duration of requests to external services",
+}, []string{"datasource_name", "datasource_type", "error_source"})
+
 // GetQuery wraps sqlutil's GetQuery to add headers if needed
 func GetQuery(query backend.DataQuery, headers http.Header, setHeaders bool) (*Query, error) {
 	model, err := sqlutil.GetQuery(query)
@@ -51,29 +60,63 @@ func GetQuery(query backend.DataQuery, headers http.Header, setHeaders bool) (*Q
 	return model, nil
 }
 
-// QueryDB sends the query to the connection and converts the rows to a dataframe.
-func QueryDB(ctx context.Context, db Connection, converters []sqlutil.Converter, fillMode *data.FillMissing, query *Query, args ...interface{}) (data.Frames, error) {
-	// Query the rows from the database
-	rows, err := db.QueryContext(ctx, query.RawSQL, args...)
+type DBQuery struct {
+	DSName     string
+	DB         Connection
+	Settings   backend.DataSourceInstanceSettings
+	converters []sqlutil.Converter
+	fillMode   *data.FillMissing
+}
+
+func NewQuery(db Connection, settings backend.DataSourceInstanceSettings, converters []sqlutil.Converter, fillMode *data.FillMissing) *DBQuery {
+	dsName, ok := sanitizeLabelName(settings.Name)
+	if !ok {
+		backend.Logger.Warn("Failed to sanitize datasource name for prometheus label", dsName)
+	}
+	return &DBQuery{
+		DB:         db,
+		DSName:     dsName,
+		converters: converters,
+		fillMode:   fillMode,
+	}
+}
+
+// Run sends the query to the connection and converts the rows to a dataframe.
+func (q *DBQuery) Run(ctx context.Context, query *Query, args ...interface{}) (data.Frames, error) {
+	start := time.Now()
+	var errWithSource error
+
+	defer func() {
+		if q.DSName != "" {
+			errorSource := "none"
+			if errWithSource != nil {
+				errorSource = string(ErrorSource(errWithSource))
+			}
+
+			duration.WithLabelValues(q.DSName, q.Settings.Type, errorSource).Observe(time.Since(start).Seconds())
+		}
+	}()
+
+	rows, err := q.DB.QueryContext(ctx, query.RawSQL, args...)
 	if err != nil {
 		errType := ErrorQuery
 		if errors.Is(err, context.Canceled) {
 			errType = context.Canceled
 		}
-		err := DownstreamError(fmt.Errorf("%w: %s", errType, err.Error()))
-		return sqlutil.ErrorFrameFromQuery(query), err
+		errWithSource := DownstreamError(fmt.Errorf("%w: %s", errType, err.Error()))
+		return sqlutil.ErrorFrameFromQuery(query), errWithSource
 	}
 
 	// Check for an error response
 	if err := rows.Err(); err != nil {
 		if errors.Is(err, sql.ErrNoRows) {
 			// Should we even response with an error here?
 			// The panel will simply show "no data"
-			err := DownstreamError(fmt.Errorf("%s: %w", "No results from query", err))
-			return sqlutil.ErrorFrameFromQuery(query), err
+			errWithSource := DownstreamError(fmt.Errorf("%s: %w", "No results from query", err))
+			return sqlutil.ErrorFrameFromQuery(query), errWithSource
 		}
-		err := DownstreamError(fmt.Errorf("%s: %w", "Error response from database", err))
-		return sqlutil.ErrorFrameFromQuery(query), err
+		errWithSource := DownstreamError(fmt.Errorf("%s: %w", "Error response from database", err))
+		return sqlutil.ErrorFrameFromQuery(query), errWithSource
 	}
 
 	defer func() {
@@ -83,10 +126,10 @@ func QueryDB(ctx context.Context, db Connection, converters []sqlutil.Converter,
 	}()
 
 	// Convert the response to frames
-	res, err := getFrames(rows, -1, converters, fillMode, query)
+	res, err := getFrames(rows, -1, q.converters, q.fillMode, query)
 	if err != nil {
-		err := PluginError(fmt.Errorf("%w: %s", err, "Could not process SQL results"))
-		return sqlutil.ErrorFrameFromQuery(query), err
+		errWithSource := PluginError(fmt.Errorf("%w: %s", err, "Could not process SQL results"))
+		return sqlutil.ErrorFrameFromQuery(query), errWithSource
 	}
 
 	return res, nil
@@ -203,3 +246,29 @@ func applyHeaders(query *Query, headers http.Header) *Query {
 
 	return query
 }
+
+// sanitizeLabelName removes all invalid chars from the label name.
+// If the label name is empty or contains only invalid chars, it will return false indicating it was not sanitized.
+// copied from https://github.com/grafana/grafana/blob/main/pkg/infra/metrics/metricutil/utils.go#L14
+func sanitizeLabelName(name string) (string, bool) {
+	if len(name) == 0 {
+		backend.Logger.Warn(fmt.Sprintf("label name cannot be empty: %s", name))
+		return "", false
+	}
+
+	out := strings.Builder{}
+	for i, b := range name {
+		if (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') || b == '_' || (b >= '0' && b <= '9' && i > 0) {
+			out.WriteRune(b)
+		} else if b == ' ' {
+			out.WriteRune('_')
+		}
+	}
+
+	if out.Len() == 0 {
+		backend.Logger.Warn(fmt.Sprintf("label name only contains invalid chars: %q", name))
+		return "", false
+	}
+
+	return out.String(), true
+}

query_integration_test.go

@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
 
 	_ "github.com/go-sql-driver/mysql"
@@ -75,7 +76,12 @@ func TestQuery_MySQL(t *testing.T) {
 			RawSQL: "SELECT SLEEP(5)",
 		}
 
-		_, err := QueryDB(ctx, db, []sqlutil.Converter{}, nil, q)
+		settings := backend.DataSourceInstanceSettings{
+			Name: "foo",
+		}
+
+		sqlQuery := NewQuery(db, settings, []sqlutil.Converter{}, nil)
+		_, err := sqlQuery.Run(ctx, q)
 		if err == nil {
 			t.Fatal("expected an error but received none")
 		}

query_test.go

@@ -8,8 +8,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/data"
 	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -78,7 +80,12 @@ func TestQuery_Timeout(t *testing.T) {
 
 		defer conn.Close()
 
-		_, err := QueryDB(ctx, conn, []sqlutil.Converter{}, nil, &Query{})
+		settings := backend.DataSourceInstanceSettings{
+			Name: "foo",
+		}
+
+		sqlQuery := NewQuery(conn, settings, []sqlutil.Converter{}, nil)
+		_, err := sqlQuery.Run(ctx, &Query{})
 
 		if !errors.Is(err, context.Canceled) {
 			t.Fatal("expected error to be context.Canceled, received", err)
@@ -100,7 +107,12 @@ func TestQuery_Timeout(t *testing.T) {
 
 		defer conn.Close()
 
-		_, err := QueryDB(ctx, conn, []sqlutil.Converter{}, nil, &Query{})
+		settings := backend.DataSourceInstanceSettings{
+			Name: "foo",
+		}
+
+		sqlQuery := NewQuery(conn, settings, []sqlutil.Converter{}, nil)
+		_, err := sqlQuery.Run(ctx, &Query{})
 
 		if !errors.Is(err, ErrorQuery) {
 			t.Fatal("expected function to complete, received error: ", err)
@@ -152,3 +164,27 @@ func TestFixFrameForLongToMulti(t *testing.T) {
 		require.Equal(t, err, fmt.Errorf("can not convert to wide series, input is missing a time field"))
 	})
 }
+
+func TestLabelNameSanitization(t *testing.T) {
+	testcases := []struct {
+		input    string
+		expected string
+		err      bool
+	}{
+		{input: "job", expected: "job"},
+		{input: "job._loal['", expected: "job_loal"},
+		{input: "", expected: "", err: true},
+		{input: ";;;", expected: "", err: true},
+		{input: "Data source", expected: "Data_source"},
+	}
+
+	for _, tc := range testcases {
+		got, ok := sanitizeLabelName(tc.input)
+		if tc.err {
+			assert.Equal(t, false, ok)
+		} else {
+			assert.Equal(t, true, ok)
+			assert.Equal(t, tc.expected, got)
+		}
+	}
+}