Query runner: Add parameter handling and Postgres settings support

Author: lfittl

output/pganalyze_collector/server_message.pb.go

@@ -39,10 +39,18 @@ message ServerMessage {
     bool pause = 1;
   }
 
+  message QueryRunPostgresSetting {
+    string name = 1;
+    string value = 2;
+  }
+
   message QueryRun {
     int64 id = 1;
     QueryRunType type = 2;
     string database_name = 3;
     string query_text = 4;
+    repeated NullString query_parameters = 5;
+    repeated string query_parameter_types = 6;
+    repeated QueryRunPostgresSetting postgres_settings = 7;
   }
 }

protobuf/server_message.proto

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/lib/pq"
 	"github.com/pganalyze/collector/input/postgres"
 	"github.com/pganalyze/collector/output"
 	"github.com/pganalyze/collector/output/pganalyze_collector"
@@ -65,11 +66,18 @@ func runQueryOnDatabase(ctx context.Context, server *state.Server, collectionOpt
 	server.QueryRuns[id].BackendPid = pid
 	server.QueryRunsMutex.Unlock()
 
+	for _, setting := range query.PostgresSettings {
+		_, err = db.ExecContext(ctx, postgres.QueryMarkerSQL+"SET %s = %s", pq.QuoteIdentifier(setting.Name), pq.QuoteLiteral(setting.Value))
+		if err != nil {
+			return "", err
+		}
+	}
+
 	// We don't include QueryMarkerSQL so query runs are reported separately in pganalyze
 	marker := fmt.Sprintf("/* pganalyze:no-alert,pganalyze-query-run:%d */ ", query.Id)
 
 	if query.Type == pganalyze_collector.QueryRunType_EXPLAIN {
-		return postgres.RunExplainAnalyzeForQueryRun(ctx, db, query.QueryText, nil, nil, marker)
+		return postgres.RunExplainAnalyzeForQueryRun(ctx, db, query.QueryText, query.QueryParameters, query.QueryParameterTypes, marker)
 	} else {
 		logger.PrintVerbose("Unhandled query run type %d for %d", query.Type, query.Id)
 		return "", errors.New("Unhandled query run type")

runner/query_run.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/gorilla/websocket"
+	"github.com/guregu/null"
 	"github.com/pganalyze/collector/output/pganalyze_collector"
 	"github.com/pganalyze/collector/state"
 	"github.com/pganalyze/collector/util"
@@ -140,12 +141,23 @@ func connect(ctx context.Context, server *state.Server, globalCollectionOpts sta
 			} else if message.GetQueryRun() != nil {
 				q := message.GetQueryRun()
 				logger.PrintVerbose("Query run %d received: %s", q.Id, q.QueryText)
+				parameters := []null.String{}
+				for _, p := range q.QueryParameters {
+					parameters = append(parameters, null.NewString(p.Value, p.Valid))
+				}
+				settings := []state.QueryRunPostgresSetting{}
+				for _, s := range q.PostgresSettings {
+					settings = append(settings, state.QueryRunPostgresSetting{Name: s.Name, Value: s.Value})
+				}
 				server.QueryRunsMutex.Lock()
 				server.QueryRuns[q.Id] = &state.QueryRun{
-					Id:           q.Id,
-					Type:         q.Type,
-					DatabaseName: q.DatabaseName,
-					QueryText:    q.QueryText,
+					Id:                  q.Id,
+					Type:                q.Type,
+					DatabaseName:        q.DatabaseName,
+					QueryText:           q.QueryText,
+					QueryParameters:     parameters,
+					QueryParameterTypes: q.QueryParameterTypes,
+					PostgresSettings:    settings,
 				}
 				server.QueryRunsMutex.Unlock()
 			}

runner/websocket.go

@@ -8,6 +8,7 @@ import (
 
 	raven "github.com/getsentry/raven-go"
 	"github.com/gorilla/websocket"
+	"github.com/guregu/null"
 	"github.com/pganalyze/collector/config"
 	"github.com/pganalyze/collector/output/pganalyze_collector"
 )
@@ -253,16 +254,24 @@ type CollectionStatus struct {
 	LogSnapshotDisabledReason string
 }
 
+type QueryRunPostgresSetting struct {
+	Name  string
+	Value string
+}
+
 type QueryRun struct {
-	Id           int64
-	Type         pganalyze_collector.QueryRunType
-	DatabaseName string
-	QueryText    string
-	Result       string
-	Error        string
-	StartedAt    time.Time
-	FinishedAt   time.Time
-	BackendPid   int
+	Id                  int64
+	Type                pganalyze_collector.QueryRunType
+	DatabaseName        string
+	QueryText           string
+	QueryParameters     []null.String
+	QueryParameterTypes []string
+	PostgresSettings    []QueryRunPostgresSetting
+	Result              string
+	Error               string
+	StartedAt           time.Time
+	FinishedAt          time.Time
+	BackendPid          int
 }
 
 type Server struct {