Improve panic error recovery (#176)

* - improve catching and handling panic errors

* - split validateRows into two functions

Author: sasklacz

datasource.go

@@ -128,12 +128,22 @@ func (ds *SQLDatasource) QueryData(ctx context.Context, req *backend.QueryDataRe
 					stack := string(debug.Stack())
 					errorMsg := fmt.Sprintf("SQL datasource query execution panic: %v", r)
 
+					// Log panic without sensitive query data
 					backend.Logger.Error(errorMsg,
 						"panic", r,
 						"refID", query.RefID,
-						"stack", stack)
-
-					response.Set(query.RefID, backend.ErrorResponseWithErrorSource(backend.PluginError(errors.New(errorMsg))))
+						"queryType", query.QueryType,
+						"maxDataPoints", query.MaxDataPoints,
+						"interval", query.Interval)
+
+					// Log stack trace separately at debug level to avoid exposing in production
+					backend.Logger.Debug("Panic stack trace", "stack", stack)
+
+					response.Set(query.RefID, backend.DataResponse{
+						Frames:      nil,
+						Error:       backend.PluginError(errors.New(errorMsg)),
+						ErrorSource: backend.ErrorSourcePlugin,
+					})
 				}
 			}()
 

datasource_test.go

@@ -2,6 +2,8 @@ package sqlds_test
 
 import (
 	"context"
+	"database/sql"
+	"database/sql/driver"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -10,6 +12,7 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
 	"github.com/grafana/sqlds/v4"
+	"github.com/grafana/sqlds/v4/mock"
 	"github.com/grafana/sqlds/v4/test"
 	"github.com/stretchr/testify/assert"
 )
@@ -223,6 +226,68 @@ func Test_query_panic_recovery(t *testing.T) {
 	assert.Nil(t, res.Frames)
 }
 
+func Test_query_panic_in_rows_validation(t *testing.T) {
+	cfg := `{ "timeout": 0, "retries": 0, "retryOn": [] }`
+	opts := test.DriverOpts{}
+
+	// Set up a driver that returns rows which will cause a panic when accessing columns
+	// This will be caught by our validateRows function
+	customQueryFunc := func(args []driver.Value) (driver.Rows, error) {
+		// Return a panicking rows implementation that will cause a panic when columns are accessed
+		return &panickingRows{}, nil
+	}
+
+	// Create a custom driver with a wrapper for the query method
+	driverName := "panic-rows-test"
+
+	// Create and register the handler
+	handler := test.NewDriverHandler(test.Data{}, opts)
+
+	// Create a custom handler that overrides the Query method
+	customHandler := &panickingDBHandler{
+		SqlHandler:      handler,
+		customQueryFunc: customQueryFunc,
+	}
+
+	mock.RegisterDriver(driverName, customHandler)
+
+	// Create datasource with the custom driver
+	testDS := &testDriver{driverName: driverName}
+	ds := sqlds.NewDatasource(testDS)
+
+	// Set up the query request
+	req, settings := setupQueryRequest("panic-rows-validation", cfg)
+	_, err := ds.NewDatasource(context.Background(), settings)
+	assert.Nil(t, err)
+
+	// Execute the query
+	data, err := ds.QueryData(context.Background(), req)
+
+	// Verify that the panic was caught and converted to an error
+	assert.Nil(t, err)
+	assert.NotNil(t, data.Responses)
+
+	res := data.Responses["foo"]
+	assert.NotNil(t, res.Error)
+	assert.Contains(t, res.Error.Error(), "failed to validate rows")
+	assert.NotNil(t, res.Frames) // Error frame is returned, not nil
+}
+
+// panickingRows is a custom rows implementation that panics when columns are accessed
+type panickingRows struct{}
+
+func (r *panickingRows) Columns() []string {
+	panic("panic in Columns method")
+}
+
+func (r *panickingRows) Close() error {
+	return nil
+}
+
+func (r *panickingRows) Next(dest []driver.Value) error {
+	return nil
+}
+
 func queryRequest(t *testing.T, name string, opts test.DriverOpts, cfg string, marcos sqlds.Macros) (*backend.QueryDataRequest, *test.SqlHandler, *sqlds.SQLDatasource) {
 	driver, handler := test.NewDriver(name, test.Data{}, nil, opts, marcos)
 	ds := sqlds.NewDatasource(driver)
@@ -269,3 +334,50 @@ func setupHealthRequest(id string, cfg string) (backend.CheckHealthRequest, back
 	}
 	return req, settings
 }
+
+// testDriver implements sqlds.Driver interface for testing
+type testDriver struct {
+	driverName string
+}
+
+func (d *testDriver) Connect(ctx context.Context, cfg backend.DataSourceInstanceSettings, msg json.RawMessage) (*sql.DB, error) {
+	return sql.Open(d.driverName, "")
+}
+
+func (d *testDriver) Settings(ctx context.Context, config backend.DataSourceInstanceSettings) sqlds.DriverSettings {
+	settings, _ := test.LoadSettings(ctx, config)
+	return settings
+}
+
+func (d *testDriver) Macros() sqlds.Macros {
+	return nil
+}
+
+func (d *testDriver) Converters() []sqlutil.Converter {
+	return nil
+}
+
+// panickingDBHandler implements mock.DBHandler and causes panics when querying
+type panickingDBHandler struct {
+	test.SqlHandler
+	customQueryFunc func(args []driver.Value) (driver.Rows, error)
+}
+
+func (h *panickingDBHandler) Query(args []driver.Value) (driver.Rows, error) {
+	if h.customQueryFunc != nil {
+		return h.customQueryFunc(args)
+	}
+	return h.SqlHandler.Query(args)
+}
+
+func (h *panickingDBHandler) Ping(ctx context.Context) error {
+	return nil
+}
+
+func (h *panickingDBHandler) Columns() []string {
+	return []string{"test_column"}
+}
+
+func (h *panickingDBHandler) Next(dest []driver.Value) error {
+	return errors.New("no more rows")
+}

query.go

@@ -9,6 +9,8 @@ import (
 	"net/http"
 	"time"
 
+	"runtime/debug"
+
 	"github.com/grafana/dataplane/sdata/timeseries"
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/data"
@@ -123,7 +125,14 @@ func (q *DBQuery) Run(ctx context.Context, query *Query, args ...interface{}) (d
 	return res, nil
 }
 
+// getFrames converts rows to dataframes
 func getFrames(rows *sql.Rows, limit int64, converters []sqlutil.Converter, fillMode *data.FillMissing, query *Query) (data.Frames, error) {
+	// Validate rows before processing to prevent panics
+	if err := validateRows(rows); err != nil {
+		backend.Logger.Error("Invalid SQL rows", "error", err.Error())
+		return nil, err
+	}
+
 	frame, err := sqlutil.FrameFromRows(rows, limit, converters...)
 	if err != nil {
 		return nil, err
@@ -185,6 +194,34 @@ func getFrames(rows *sql.Rows, limit int64, converters []sqlutil.Converter, fill
 	return data.Frames{frame}, nil
 }
 
+// accessColumns checks whether we can access rows.Columns, checking
+// for error or panic. In the case of panic, logs the stack trace at debug level
+// for security
+func accessColumns(rows *sql.Rows) (columnErr error) {
+	defer func() {
+		if r := recover(); r != nil {
+			columnErr = fmt.Errorf("panic accessing columns: %v", r)
+			stack := string(debug.Stack())
+			backend.Logger.Debug("accessColumns panic stack trace", "stack", stack)
+		}
+	}()
+	_, columnErr = rows.Columns()
+	return columnErr
+}
+
+// validateRows performs safety checks on SQL rows to prevent panics
+func validateRows(rows *sql.Rows) error {
+	if rows == nil {
+		return fmt.Errorf("rows is nil")
+	}
+
+	err := accessColumns(rows)
+	if err != nil {
+		return fmt.Errorf("failed to validate rows: %w", err)
+	}
+	return nil
+}
+
 // fixFrameForLongToMulti edits the passed in frame so that it's first time field isn't nullable and has the correct meta
 func fixFrameForLongToMulti(frame *data.Frame) error {
 	if frame == nil {