feat(postgresql): add disable_pg_stat_database and disable_pg_stat_bgwriter config options (#1430)

* fix(postgresql): add rows.Err() checks

* feat(postgresql): add disable_pg_stat_database and disable_pg_stat_bgwriter config options

Author: kongfei605

conf/input.postgresql/postgresql.toml

@@ -35,10 +35,26 @@
   ## databases are gathered.  Do NOT use with the 'ignored_databases' option.
   # databases = ["app_production", "testing"]
 
+  ## Whether to collect statement-level metrics.
+  ## Requires extension pg_stat_statements enabled, see https://www.postgresql.org/docs/current/pgstatstatements.html
+  # enable_statement_metrics = false
+  ## Max number of statements to collect
+  ## applies only when enable_statement_metrics=true
+  ## 0 means no limit
+  # statement_metrics_limit = 100
+ 
   ## Whether to use prepared statements when connecting to the database.
   ## This should be set to false when connecting through a PgBouncer instance
   ## with pool_mode set to transaction.
   #prepared_statements = true
+
+  ## Whether to skip built-in metrics from pg_stat_database.
+  # disable_pg_stat_database = false
+
+  ## Whether to skip built-in metrics from pg_stat_bgwriter.
+  ## On PostgreSQL 17+, this also skips compatibility metrics collected from pg_stat_checkpointer.
+  # disable_pg_stat_bgwriter = false
+
   # [[instances.metrics]]
   # mesurement = "sessions"
   # label_fields = [ "status", "type" ]

go.mod

@@ -176,6 +176,7 @@ require (
 require (
 	cloud.google.com/go/monitoring v1.16.3
 	github.com/AlekSi/pointer v1.2.0
+	github.com/DATA-DOG/go-sqlmock v1.5.2
 	github.com/IBM/sarama v1.42.1
 	github.com/Knetic/govaluate v3.0.1-0.20171022003610-9aa49832a739+incompatible
 	github.com/NVIDIA/go-dcgm v0.0.0-20240118201113-3385e277e49f

go.sum

@@ -643,6 +643,8 @@ github.com/BurntSushi/toml v1.4.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
 github.com/Code-Hex/go-generics-cache v1.5.1 h1:6vhZGc5M7Y/YD8cIUcY8kcuQLB4cHR7U+0KMqAA0KcU=
 github.com/Code-Hex/go-generics-cache v1.5.1/go.mod h1:qxcC9kRVrct9rHeiYpFWSoW1vxyillCVzX13KZG8dl4=
+github.com/DATA-DOG/go-sqlmock v1.5.2 h1:OcvFkGmslmlZibjAjaHm3L//6LiuBgolP7OputlJIzU=
+github.com/DATA-DOG/go-sqlmock v1.5.2/go.mod h1:88MAG/4G7SMwSE3CeA0ZKzrT5CiOU3OJ+JlNzwDqpNU=
 github.com/DataDog/datadog-go v3.2.0+incompatible/go.mod h1:LButxg5PwREeZtORoXG3tL4fMGNddJ+vMq1mwgfaqoQ=
 github.com/IBM/sarama v1.42.1 h1:wugyWa15TDEHh2kvq2gAy1IHLjEjuYOYgXz/ruC/OSQ=
 github.com/IBM/sarama v1.42.1/go.mod h1:Xxho9HkHd4K/MDUo/T/sOqwtX/17D33++E9Wib6hUdQ=
@@ -1412,6 +1414,7 @@ github.com/keybase/go-keychain v0.0.0-20231219164618-57a3676c3af6/go.mod h1:3VeW
 github.com/kisielk/errcheck v1.1.0/go.mod h1:EZBBE59ingxPouuu3KfxchcWSUPOHkagtvWXihfKN4Q=
 github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
+github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE=
 github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE=
 github.com/klauspost/compress v1.4.1/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
 github.com/klauspost/compress v1.10.10/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=

inputs/postgresql/README.md

@@ -50,6 +50,13 @@ address = ""
 ## with pool_mode set to transaction.
 ## 是否使用prepared statements 连接数据库
 # prepared_statements = true
+
+## Whether to skip built-in metrics from pg_stat_database.
+# disable_pg_stat_database = false
+
+## Whether to skip built-in metrics from pg_stat_bgwriter.
+## On PostgreSQL 17+, this also skips compatibility metrics collected from pg_stat_checkpointer.
+# disable_pg_stat_bgwriter = false
 ```
 ![dashboard](./postgresql.png)
 

inputs/postgresql/postgresql.go

@@ -81,6 +81,8 @@ type Instance struct {
 	PreparedStatements     bool            `toml:"prepared_statements"`
 	EnableStatementMetrics bool            `toml:"enable_statement_metrics"`
 	StatementMetricsLimit  int             `toml:"statement_metrics_limit"`
+	DisablePgStatDatabase  bool            `toml:"disable_pg_stat_database"`
+	DisablePgStatBgwriter  bool            `toml:"disable_pg_stat_bgwriter"`
 	Metrics                []MetricConfig  `toml:"metrics"`
 	TrimServerTagSpace     bool            `toml:"trim_server_tag_space"`
 
@@ -94,7 +96,7 @@ type Instance struct {
 
 var ignoredColumns = map[string]bool{"stats_reset": true}
 
-func (p *Instance) IgnoredColumns() map[string]bool {
+func (ins *Instance) IgnoredColumns() map[string]bool {
 	return ignoredColumns
 }
 
@@ -142,17 +144,12 @@ func (ins *Instance) Init() error {
 }
 
 // closes any necessary channels and connections
-func (p *Instance) Drop() {
+func (ins *Instance) Drop() {
 	// Ignore the returned error as we cannot do anything about it anyway
 	//nolint:errcheck,revive
 }
 
 func (ins *Instance) Gather(slist *types.SampleList) {
-	var (
-		err     error
-		query   string
-		columns []string
-	)
 	addr, err := ins.SanitizedAddress()
 	if err != nil {
 		log.Println("E! can't sanitize address :", err)
@@ -171,40 +168,55 @@ func (ins *Instance) Gather(slist *types.SampleList) {
 	}
 	slist.PushSample(inputName, "up", 1, tags)
 
+	ins.gatherMetrics(slist)
+}
+
+func (ins *Instance) gatherMetrics(slist *types.SampleList) {
+	var (
+		err     error
+		query   string
+		columns []string
+	)
+
 	ins.db.SetMaxOpenConns(ins.MaxOpen)
 	ins.db.SetMaxIdleConns(ins.MaxIdle)
 	ins.db.SetConnMaxLifetime(time.Duration(ins.MaxLifetime))
 
-	if len(ins.Databases) == 0 && len(ins.IgnoredDatabases) == 0 {
-		query = `SELECT * FROM pg_stat_database`
-	} else if len(ins.IgnoredDatabases) != 0 {
-		query = fmt.Sprintf(`SELECT * FROM pg_stat_database WHERE datname NOT IN ('%s')`,
-			strings.Join(ins.IgnoredDatabases, "','"))
-	} else {
-		query = fmt.Sprintf(`SELECT * FROM pg_stat_database WHERE datname IN ('%s')`,
-			strings.Join(ins.Databases, "','"))
-	}
-
-	rows, err := ins.db.Query(query)
-	if err != nil {
-		log.Println("E! failed to execute Query :", err)
-		return
-	}
+	if !ins.DisablePgStatDatabase {
+		if len(ins.Databases) == 0 && len(ins.IgnoredDatabases) == 0 {
+			query = `SELECT * FROM pg_stat_database`
+		} else if len(ins.IgnoredDatabases) != 0 {
+			query = fmt.Sprintf(`SELECT * FROM pg_stat_database WHERE datname NOT IN ('%s')`,
+				strings.Join(ins.IgnoredDatabases, "','"))
+		} else {
+			query = fmt.Sprintf(`SELECT * FROM pg_stat_database WHERE datname IN ('%s')`,
+				strings.Join(ins.Databases, "','"))
+		}
 
-	defer rows.Close()
+		rows, err := ins.db.Query(query)
+		if err != nil {
+			log.Println("E! failed to execute Query :", err)
+			return
+		}
 
-	// grab the column information from the result
-	if columns, err = rows.Columns(); err != nil {
-		log.Println("E! failed to grab column info:", err)
-		return
-	}
+		defer rows.Close()
 
-	for rows.Next() {
-		err = ins.accRow(rows, slist, "", columns, columns, nil)
-		if err != nil {
-			log.Println("E! failed to get row data:", err)
+		// grab the column information from the result
+		if columns, err = rows.Columns(); err != nil {
+			log.Println("E! failed to grab column info:", err)
 			return
 		}
+
+		for rows.Next() {
+			err = ins.accRow(rows, slist, "", columns, columns, nil)
+			if err != nil {
+				log.Println("E! failed to get row data:", err)
+				return
+			}
+		}
+		if err := rows.Err(); err != nil {
+			log.Println("E! failed to iterate pg_stat_database rows:", err)
+		}
 	}
 
 	// Check Postgres Version
@@ -218,89 +230,100 @@ func (ins *Instance) Gather(slist *types.SampleList) {
 		ins.Version = version
 	}
 
-	if ins.Version < 170000 {
-		query = `SELECT * FROM pg_stat_bgwriter`
-		bgWriterRow, err := ins.db.Query(query)
-		if err != nil {
-			log.Println("E! failed to execute Query:", err)
-			return
-		}
-
-		defer bgWriterRow.Close()
-
-		// grab the column information from the result
-		if columns, err = bgWriterRow.Columns(); err != nil {
-			log.Println("E! failed to grab column info:", err)
-			return
-		}
-
-		for bgWriterRow.Next() {
-			err = ins.accRow(bgWriterRow, slist, "", columns, columns, nil)
+	if !ins.DisablePgStatBgwriter {
+		if ins.Version < 170000 {
+			query = `SELECT * FROM pg_stat_bgwriter`
+			bgWriterRow, err := ins.db.Query(query)
 			if err != nil {
-				log.Println("E! failed to get row data:", err)
+				log.Println("E! failed to execute Query:", err)
 				return
 			}
-		}
-	} else {
-		// PG 17+ split pg_stat_bgwriter into pg_stat_bgwriter and pg_stat_checkpointer
 
-		// 1. Query pg_stat_bgwriter (remaining columns)
-		query = `SELECT * FROM pg_stat_bgwriter`
-		bgWriterRow, err := ins.db.Query(query)
-		if err != nil {
-			log.Println("E! failed to execute Query pg_stat_bgwriter:", err)
-			return
-		}
-		defer bgWriterRow.Close()
+			defer bgWriterRow.Close()
 
-		if columns, err = bgWriterRow.Columns(); err != nil {
-			log.Println("E! failed to grab column info for pg_stat_bgwriter:", err)
-			return
-		}
+			// grab the column information from the result
+			if columns, err = bgWriterRow.Columns(); err != nil {
+				log.Println("E! failed to grab column info:", err)
+				return
+			}
 
-		for bgWriterRow.Next() {
-			err = ins.accRow(bgWriterRow, slist, "", columns, columns, nil)
+			for bgWriterRow.Next() {
+				err = ins.accRow(bgWriterRow, slist, "", columns, columns, nil)
+				if err != nil {
+					log.Println("E! failed to get row data:", err)
+					return
+				}
+			}
+			if err := bgWriterRow.Err(); err != nil {
+				log.Println("E! failed to iterate pg_stat_bgwriter rows:", err)
+			}
+		} else {
+			// PG 17+ split pg_stat_bgwriter into pg_stat_bgwriter and pg_stat_checkpointer
+
+			// 1. Query pg_stat_bgwriter (remaining columns)
+			query = `SELECT * FROM pg_stat_bgwriter`
+			bgWriterRow, err := ins.db.Query(query)
 			if err != nil {
-				log.Println("E! failed to get row data from pg_stat_bgwriter:", err)
+				log.Println("E! failed to execute Query pg_stat_bgwriter:", err)
 				return
 			}
-		}
+			defer bgWriterRow.Close()
 
-		// 2. Query pg_stat_checkpointer (moved columns, aliased to old names for compatibility)
-		// num_timed -> checkpoints_timed
-		// num_requested -> checkpoints_req
-		// write_time -> checkpoint_write_time
-		// sync_time -> checkpoint_sync_time
-		// buffers_written -> buffers_checkpoint
-		query = `SELECT 
-			num_timed AS checkpoints_timed,
-			num_requested AS checkpoints_req,
-			write_time AS checkpoint_write_time,
-			sync_time AS checkpoint_sync_time,
-			buffers_written AS buffers_checkpoint,
-			restartpoints_timed,
-			restartpoints_req,
-			restartpoints_done
-			FROM pg_stat_checkpointer`
-
-		checkpointerRow, err := ins.db.Query(query)
-		if err != nil {
-			log.Println("E! failed to get row data:", err)
-			return
-		}
-		defer checkpointerRow.Close()
+			if columns, err = bgWriterRow.Columns(); err != nil {
+				log.Println("E! failed to grab column info for pg_stat_bgwriter:", err)
+				return
+			}
 
-		if columns, err = checkpointerRow.Columns(); err != nil {
-			log.Println("E! failed to grab column info for pg_stat_checkpointer:", err)
-			return
-		}
+			for bgWriterRow.Next() {
+				err = ins.accRow(bgWriterRow, slist, "", columns, columns, nil)
+				if err != nil {
+					log.Println("E! failed to get row data from pg_stat_bgwriter:", err)
+					return
+				}
+			}
+			if err := bgWriterRow.Err(); err != nil {
+				log.Println("E! failed to iterate pg_stat_bgwriter rows:", err)
+			}
 
-		for checkpointerRow.Next() {
-			err = ins.accRow(checkpointerRow, slist, "", columns, columns, nil)
+			// 2. Query pg_stat_checkpointer (moved columns, aliased to old names for compatibility)
+			// num_timed -> checkpoints_timed
+			// num_requested -> checkpoints_req
+			// write_time -> checkpoint_write_time
+			// sync_time -> checkpoint_sync_time
+			// buffers_written -> buffers_checkpoint
+			query = `SELECT 
+				num_timed AS checkpoints_timed,
+				num_requested AS checkpoints_req,
+				write_time AS checkpoint_write_time,
+				sync_time AS checkpoint_sync_time,
+				buffers_written AS buffers_checkpoint,
+				restartpoints_timed,
+				restartpoints_req,
+				restartpoints_done
+				FROM pg_stat_checkpointer`
+
+			checkpointerRow, err := ins.db.Query(query)
 			if err != nil {
-				log.Println("E! failed to get row data from pg_stat_checkpointer:", err)
+				log.Println("E! failed to execute Query pg_stat_checkpointer:", err)
 				return
 			}
+			defer checkpointerRow.Close()
+
+			if columns, err = checkpointerRow.Columns(); err != nil {
+				log.Println("E! failed to grab column info for pg_stat_checkpointer:", err)
+				return
+			}
+
+			for checkpointerRow.Next() {
+				err = ins.accRow(checkpointerRow, slist, "", columns, columns, nil)
+				if err != nil {
+					log.Println("E! failed to get row data from pg_stat_checkpointer:", err)
+					return
+				}
+			}
+			if err := checkpointerRow.Err(); err != nil {
+				log.Println("E! failed to iterate pg_stat_checkpointer rows:", err)
+			}
 		}
 	}
 
@@ -396,6 +419,9 @@ func (ins *Instance) getStatementMetrics(slist *types.SampleList, version int) {
 			return
 		}
 	}
+	if err := statements.Err(); err != nil {
+		log.Println("E! failed to iterate pg_stat_statements rows:", err)
+	}
 }
 
 func (ins *Instance) scrapeMetric(waitMetrics *sync.WaitGroup, slist *types.SampleList, metricConf MetricConfig, tags map[string]string) {
@@ -458,6 +484,9 @@ func (ins *Instance) scrapeMetric(waitMetrics *sync.WaitGroup, slist *types.Samp
 			log.Println("E! no metrics found while parsing")
 		}
 	}
+	if err := rows.Err(); err != nil {
+		log.Println("E! failed to iterate custom query rows:", err)
+	}
 }
 
 func (ins *Instance) parseRow(row map[string]string, metricConf MetricConfig, slist *types.SampleList, tags map[string]string) error {
@@ -523,10 +552,10 @@ func (ins *Instance) accRow(row scanner, slist *types.SampleList, prefix string,
 
 	// deconstruct array of variables and send to Scan
 	err := row.Scan(columnVars...)
-
 	if err != nil {
 		return err
 	}
+
 	if columnMap["datname"] != nil {
 		// extract the database name from the column map
 		if dbNameStr, ok := (*columnMap["datname"]).(string); ok {
@@ -658,4 +687,4 @@ func (ins *Instance) SanitizedAddress() (sanitizedAddress string, err error) {
 	}
 
 	return sanitizedAddress, err
-}
+}