Add option to use simple or extended query wire protocols.

mdevan · mdevan · commit 1a97fa469e31 · 2025-02-21T11:52:21.000+05:30
diff --git a/cmd/pgmetrics/main.go b/cmd/pgmetrics/main.go
@@ -61,6 +61,8 @@ Collection options:
                                    queries (default: 500)
       --statements-limit=LIMIT collect only utmost LIMIT number of row from
                                    pg_stat_statements (default: 100)
+      --query-proto=PROTO      which query wire protocol to use; "simple" or
+                                   "extended" (default: "simple")
       --only-listed            collect info only from the databases listed as
                                    command-line args (use with Heroku)
       --all-dbs                collect info from all user databases
@@ -151,7 +153,8 @@ type options struct {
 	tooLongSec uint
 	nopager    bool
 	// connection
-	passNone bool
+	passNone   bool
+	queryProto string
 }
 
 func (o *options) defaults() {
@@ -169,6 +172,7 @@ func (o *options) defaults() {
 	o.nopager = false
 	// connection
 	o.passNone = false
+	o.queryProto = "simple"
 }
 
 func (o *options) usage(code int) {
@@ -234,6 +238,7 @@ func (o *options) parse() (args []string) {
 	s.StringVarLong(&o.CollectConfig.User, "username", 'U', "")
 	s.BoolVarLong(&o.passNone, "no-password", 'w', "")
 	s.StringVarLong(&o.CollectConfig.Role, "role", 0, "")
+	s.StringVarLong(&o.queryProto, "query-proto", 0, "")
 
 	// parse
 	s.Parse(os.Args)
@@ -296,6 +301,13 @@ func (o *options) parse() (args []string) {
 			os.Exit(2)
 		}
 	}
+	if o.queryProto != "simple" && o.queryProto != "extended" {
+		fmt.Fprintln(os.Stderr, `option --query-proto must be "simple" or "extended"`)
+		printTry()
+		os.Exit(2)
+	} else {
+		o.CollectConfig.UseExtendedQP = o.queryProto == "extended"
+	}
 
 	// help action
 	if o.helpShort || o.help == "short" || o.help == "variables" {
diff --git a/collector/collect.go b/collector/collect.go
@@ -99,6 +99,7 @@ type CollectConfig struct {
 	AllDBs          bool
 	AzureResourceID string
 	Pgpool          bool // collect only pgpool information
+	UseExtendedQP   bool // use extended query protocol instead of simple
 
 	// connection
 	Host     string
@@ -115,21 +116,13 @@ func DefaultCollectConfig() CollectConfig {
 		// ------------------ general
 		TimeoutSec:          5,
 		LockTimeoutMillisec: 50,
-		//NoSizes: false,
 
 		// ------------------ collection
-		//Schema: "",
-		//ExclSchema: "",
-		//Table: "",
-		//ExclTable: "",
-		//Omit: nil,
-		//OnlyListedDBs: false,
 		SQLLength:  500,
 		StmtsLimit: 100,
 		LogSpan:    5,
 
 		// ------------------ connection
-		//Password: "",
 	}
 
 	// connection: host
@@ -231,8 +224,13 @@ func Collect(o CollectConfig, dbnames []string) *pgmetrics.Model {
 	// set application name
 	connstr += makeKV("application_name", "pgmetrics")
 
-	// use simple protocol for maximum compatibility (pgx-specific keyword)
-	connstr += makeKV("default_query_exec_mode", "simple_protocol")
+	// Using simple protocol for maximum compatibility is the default for
+	// pgmetrics. This is selected by adding default_query_exec_mode=simple_protocol
+	// to the connection string. To use extended query protocol, nothing needs
+	// to be added.
+	if !o.UseExtendedQP {
+		connstr += makeKV("default_query_exec_mode", "simple_protocol")
+	}
 
 	// if "all DBs" was specified, collect the names of databases first
 	if o.AllDBs {
