Migrate pg_locks to collector package (prometheus-community#817)

Migrate the `pg_locks_count` query from `main` to the `collector` package. Signed-off-by: SuperQ <superq@gmail.com>

Author: SuperQ

README.md

@@ -73,7 +73,10 @@ This will build the docker image as `prometheuscommunity/postgres_exporter:${bra
 
 
 * `[no-]collector.database`
- Enable the database collector (default: enabled).
+ Enable the `database` collector (default: enabled).
+
+* `[no-]collector.locks`
+ Enable the `locks` collector (default: enabled).
 
 * `[no-]collector.postmaster`
  Enable the `postmaster` collector (default: enabled).

cmd/postgres_exporter/postgres_exporter.go

@@ -176,15 +176,6 @@ var builtinMetricMaps = map[string]intermediateMetricMap{
 true,
 0,
 },
-"pg_locks": {
-map[string]ColumnMapping{
-"datname": {LABEL, "Name of this database", nil, nil},
-"mode": {LABEL, "Type of Lock", nil, nil},
-"count": {GAUGE, "Number of locks", nil, nil},
-},
-true,
-0,
-},
 "pg_stat_replication": {
 map[string]ColumnMapping{
 "procpid": {DISCARD, "Process ID of a WAL sender process", nil, semver.MustParseRange("<9.2.0")},

cmd/postgres_exporter/queries.go

@@ -46,31 +46,6 @@ type OverrideQuery struct {
 // Overriding queries for namespaces above.
 // TODO: validate this is a closed set in tests, and there are no overlaps
 var queryOverrides = map[string][]OverrideQuery{
-"pg_locks": {
-{
-semver.MustParseRange(">0.0.0"),
-`SELECT pg_database.datname,tmp.mode,COALESCE(count,0) as count
-FROM
-(
- VALUES ('accesssharelock'),
- ('rowsharelock'),
- ('rowexclusivelock'),
- ('shareupdateexclusivelock'),
- ('sharelock'),
- ('sharerowexclusivelock'),
- ('exclusivelock'),
- ('accessexclusivelock'),
- ('sireadlock')
-) AS tmp(mode) CROSS JOIN pg_database
-LEFT JOIN
- (SELECT database, lower(mode) AS mode,count(*) AS count
- FROM pg_locks WHERE database IS NOT NULL
- GROUP BY database, lower(mode)
-) AS tmp2
-ON tmp.mode=tmp2.mode and pg_database.oid = tmp2.database ORDER BY 1`,
-},
-},
-
 "pg_stat_replication": {
 {
 semver.MustParseRange(">=10.0.0"),

collector/pg_locks.go

@@ -0,0 +1,129 @@
+// Copyright 2023 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package collector
+
+import (
+"context"
+"database/sql"
+
+"github.com/go-kit/log"
+"github.com/prometheus/client_golang/prometheus"
+)
+
+const locksSubsystem = "locks"
+
+func init() {
+registerCollector(locksSubsystem, defaultEnabled, NewPGLocksCollector)
+}
+
+type PGLocksCollector struct {
+log log.Logger
+}
+
+func NewPGLocksCollector(config collectorConfig) (Collector, error) {
+return &PGLocksCollector{
+log: config.logger,
+}, nil
+}
+
+var (
+pgLocksDesc = prometheus.NewDesc(
+prometheus.BuildFQName(
+namespace,
+locksSubsystem,
+"count",
+),
+"Number of locks",
+[]string{"datname", "mode"}, nil,
+)
+
+pgLocksQuery = `
+SELECT 
+ pg_database.datname as datname,
+ tmp.mode as mode,
+ COALESCE(count, 0) as count 
+FROM 
+ (
+ VALUES 
+ ('accesssharelock'), 
+ ('rowsharelock'), 
+ ('rowexclusivelock'), 
+ ('shareupdateexclusivelock'), 
+ ('sharelock'), 
+ ('sharerowexclusivelock'), 
+ ('exclusivelock'), 
+ ('accessexclusivelock'), 
+ ('sireadlock')
+ ) AS tmp(mode)
+ CROSS JOIN pg_database 
+ LEFT JOIN (
+ SELECT 
+ database, 
+ lower(mode) AS mode, 
+ count(*) AS count 
+ FROM 
+ pg_locks 
+ WHERE 
+ database IS NOT NULL 
+ GROUP BY 
+ database, 
+ lower(mode)
+ ) AS tmp2 ON tmp.mode = tmp2.mode 
+ and pg_database.oid = tmp2.database 
+ORDER BY 
+ 1
+`
+)
+
+// Update implements Collector and exposes database locks.
+// It is called by the Prometheus registry when collecting metrics.
+func (c PGLocksCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
+db := instance.getDB()
+// Query the list of databases
+rows, err := db.QueryContext(ctx,
+pgLocksQuery,
+)
+if err != nil {
+return err
+}
+defer rows.Close()
+
+var datname, mode sql.NullString
+var count sql.NullInt64
+
+for rows.Next() {
+if err := rows.Scan(&datname, &mode, &count); err != nil {
+return err
+}
+
+if !datname.Valid || !mode.Valid {
+continue
+}
+
+countMetric := 0.0
+if count.Valid {
+countMetric = float64(count.Int64)
+}
+
+ch <- prometheus.MustNewConstMetric(
+pgLocksDesc,
+prometheus.GaugeValue, countMetric,
+datname.String, mode.String,
+)
+}
+if err := rows.Err(); err != nil {
+return err
+}
+return nil
+}

collector/pg_locks_test.go

@@ -0,0 +1,60 @@
+// Copyright 2023 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package collector
+
+import (
+"context"
+"testing"
+
+"github.com/DATA-DOG/go-sqlmock"
+"github.com/prometheus/client_golang/prometheus"
+dto "github.com/prometheus/client_model/go"
+"github.com/smartystreets/goconvey/convey"
+)
+
+func TestPGLocksCollector(t *testing.T) {
+db, mock, err := sqlmock.New()
+if err != nil {
+t.Fatalf("Error opening a stub db connection: %s", err)
+}
+defer db.Close()
+
+inst := &instance{db: db}
+
+rows := sqlmock.NewRows([]string{"datname", "mode", "count"}).
+AddRow("test", "exclusivelock", 42)
+
+mock.ExpectQuery(sanitizeQuery(pgLocksQuery)).WillReturnRows(rows)
+
+ch := make(chan prometheus.Metric)
+go func() {
+defer close(ch)
+c := PGLocksCollector{}
+if err := c.Update(context.Background(), inst, ch); err != nil {
+t.Errorf("Error calling PGLocksCollector.Update: %s", err)
+}
+}()
+
+expected := []MetricResult{
+{labels: labelMap{"datname": "test", "mode": "exclusivelock"}, value: 42, metricType: dto.MetricType_GAUGE},
+}
+convey.Convey("Metrics comparison", t, func() {
+for _, expect := range expected {
+m := readMetric(<-ch)
+convey.So(expect, convey.ShouldResemble, m)
+}
+})
+if err := mock.ExpectationsWereMet(); err != nil {
+t.Errorf("there were unfulfilled exceptions: %s", err)
+}
+}