Enable dataproxy.row_limit configuration option from Grafana (#162)

Author: kevinwcyu

datasource.go

@@ -19,6 +19,7 @@ import (
 )
 
 const defaultKeySuffix = "default"
+const defaultRowLimit = int64(-1)
 
 var (
 	ErrorMissingMultipleConnectionsConfig = backend.PluginError(errors.New("received connection arguments but the feature is not enabled"))
@@ -50,6 +51,10 @@ type SQLDatasource struct {
 	CustomRoutes              map[string]func(http.ResponseWriter, *http.Request)
 	metrics                   Metrics
 	EnableMultipleConnections bool
+	// EnableRowLimit: enables using the dataproxy.row_limit setting to limit the number of rows returned by the query
+	// https://grafana.com/docs/grafana/latest/setup-grafana/configure-grafana/#row_limit
+	EnableRowLimit bool
+	rowLimit       int64
 	// PreCheckHealth (optional). Performs custom health check before the Connect method
 	PreCheckHealth func(ctx context.Context, req *backend.CheckHealthRequest) *backend.CheckHealthResult
 	// PostCheckHealth (optional).Performs custom health check after the Connect method
@@ -73,6 +78,17 @@ func (ds *SQLDatasource) NewDatasource(ctx context.Context, settings backend.Dat
 	ds.CallResourceHandler = httpadapter.New(mux)
 	ds.metrics = NewMetrics(settings.Name, settings.Type, EndpointQuery)
 
+	config := backend.GrafanaConfigFromContext(ctx)
+	ds.rowLimit = defaultRowLimit
+	if ds.EnableRowLimit && config != nil {
+		sqlConfig, err := config.SQL()
+		if err != nil {
+			backend.Logger.Error(fmt.Sprintf("failed setting row limit from sql config: %s", err))
+		} else {
+			ds.rowLimit = sqlConfig.RowLimit
+		}
+	}
+
 	return ds, nil
 }
 
@@ -191,7 +207,7 @@ 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"
-	dbQuery := NewQuery(dbConn.db, dbConn.settings, ds.driver().Converters(), fillMode)
+	dbQuery := NewQuery(dbConn.db, dbConn.settings, ds.driver().Converters(), fillMode, ds.rowLimit)
 	res, err := dbQuery.Run(ctx, q, args...)
 	if err == nil {
 		return res, nil
@@ -217,7 +233,7 @@ func (ds *SQLDatasource) handleQuery(ctx context.Context, req backend.DataQuery,
 					time.Sleep(time.Duration(ds.DriverSettings().Pause * int(time.Second)))
 				}
 
-				dbQuery := NewQuery(db, dbConn.settings, ds.driver().Converters(), fillMode)
+				dbQuery := NewQuery(db, dbConn.settings, ds.driver().Converters(), fillMode, ds.rowLimit)
 				res, err = dbQuery.Run(ctx, q, args...)
 				if err == nil {
 					return res, err
@@ -239,7 +255,7 @@ func (ds *SQLDatasource) handleQuery(ctx context.Context, req backend.DataQuery,
 				continue
 			}
 
-			dbQuery := NewQuery(db, dbConn.settings, ds.driver().Converters(), fillMode)
+			dbQuery := NewQuery(db, dbConn.settings, ds.driver().Converters(), fillMode, ds.rowLimit)
 			res, err = dbQuery.Run(ctx, q, args...)
 			if err == nil {
 				return res, err
@@ -289,3 +305,7 @@ func (ds *SQLDatasource) errors(response *Response) error {
 	}
 	return err
 }
+
+func (ds *SQLDatasource) GetRowLimit() int64 {
+	return ds.rowLimit
+}

query.go

@@ -57,15 +57,17 @@ type DBQuery struct {
 	metrics    Metrics
 	DSName     string
 	converters []sqlutil.Converter
+	rowLimit   int64
 }
 
-func NewQuery(db Connection, settings backend.DataSourceInstanceSettings, converters []sqlutil.Converter, fillMode *data.FillMissing) *DBQuery {
+func NewQuery(db Connection, settings backend.DataSourceInstanceSettings, converters []sqlutil.Converter, fillMode *data.FillMissing, rowLimit int64) *DBQuery {
 	return &DBQuery{
 		DB:         db,
 		DSName:     settings.Name,
 		converters: converters,
 		fillMode:   fillMode,
 		metrics:    NewMetrics(settings.Name, settings.Type, EndpointQuery),
+		rowLimit:   rowLimit,
 	}
 }
 
@@ -105,14 +107,14 @@ func (q *DBQuery) Run(ctx context.Context, query *Query, args ...interface{}) (d
 
 	start = time.Now()
 	// Convert the response to frames
-	res, err := getFrames(rows, -1, q.converters, q.fillMode, query)
+	res, err := getFrames(rows, q.rowLimit, q.converters, q.fillMode, query)
 	if err != nil {
 		// We default to plugin error source
 		errSource := backend.ErrorSourcePlugin
 		if backend.IsDownstreamHTTPError(err) || isProcessingDownstreamError(err) {
 			errSource = backend.ErrorSourceDownstream
 		}
-		errWithSource := backend.NewErrorWithSource(fmt.Errorf("%w: %s", err, "Could not process SQL results"), errSource)	
+		errWithSource := backend.NewErrorWithSource(fmt.Errorf("%w: %s", err, "Could not process SQL results"), errSource)
 		q.metrics.CollectDuration(Source(errSource), StatusError, time.Since(start).Seconds())
 		return sqlutil.ErrorFrameFromQuery(query), errWithSource
 	}

query_integration_test.go

@@ -79,7 +79,7 @@ func TestQuery_MySQL(t *testing.T) {
 			Name: "foo",
 		}
 
-		sqlQuery := NewQuery(db, settings, []sqlutil.Converter{}, nil)
+		sqlQuery := NewQuery(db, settings, []sqlutil.Converter{}, nil, defaultRowLimit)
 		_, err := sqlQuery.Run(ctx, q)
 		if err == nil {
 			t.Fatal("expected an error but received none")

query_test.go

@@ -84,7 +84,7 @@ func TestQuery_Timeout(t *testing.T) {
 			Name: "foo",
 		}
 
-		sqlQuery := NewQuery(conn, settings, []sqlutil.Converter{}, nil)
+		sqlQuery := NewQuery(conn, settings, []sqlutil.Converter{}, nil, defaultRowLimit)
 		_, err := sqlQuery.Run(ctx, &Query{})
 
 		if !errors.Is(err, context.Canceled) {
@@ -111,7 +111,7 @@ func TestQuery_Timeout(t *testing.T) {
 			Name: "foo",
 		}
 
-		sqlQuery := NewQuery(conn, settings, []sqlutil.Converter{}, nil)
+		sqlQuery := NewQuery(conn, settings, []sqlutil.Converter{}, nil, defaultRowLimit)
 		_, err := sqlQuery.Run(ctx, &Query{})
 
 		if !errors.Is(err, ErrorQuery) {