diff --git a/collector/pg_connections_by_client.go b/collector/pg_connections_by_client.go new file mode 100644 index 000000000..4b3fc1c7a --- /dev/null +++ b/collector/pg_connections_by_client.go @@ -0,0 +1,100 @@ +// Copyright 2024 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" + "log/slog" + + "github.com/prometheus/client_golang/prometheus" +) + +const connectionsByClientSubsystem = "connections_by_client" + +func init() { + registerCollector(connectionsByClientSubsystem, defaultEnabled, NewPGConnectionsByClientCollector) +} + +type PGConnectionByClientCollector struct { + log *slog.Logger +} + +func NewPGConnectionsByClientCollector(config collectorConfig) (Collector, error) { + return &PGConnectionByClientCollector{ + log: config.logger, + }, nil +} + +var ( + pgClientCountDesc = prometheus.NewDesc( + prometheus.BuildFQName( + namespace, + connectionsByClientSubsystem, + "count", + ), + "Number of clients", + []string{"name"}, nil, + ) + + pgConnectionsByClientQuery = ` + SELECT + count(*) as count, + client_hostname + FROM pg_stat_activity + WHERE client_hostname is not null + GROUP BY client_hostname; + ` +) + +func (c PGConnectionByClientCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error { + db := instance.getDB() + + rows, err := db.QueryContext(ctx, + pgConnectionsByClientQuery, + ) + if err != nil { + return err + } + defer rows.Close() + var count sql.NullInt64 + var name sql.NullString + + for rows.Next() { + if err := rows.Scan(&count, &name); err != nil { + return err + } + + if !name.Valid { + continue + } + + countMetric := 0.0 + if count.Valid { + countMetric = float64(count.Int64) + } + + ch <- prometheus.MustNewConstMetric( + pgClientCountDesc, + prometheus.GaugeValue, + countMetric, + name.String, + ) + } + if err := rows.Err(); err != nil { + return err + } + + return nil +} diff --git a/collector/pg_connections_by_client_test.go b/collector/pg_connections_by_client_test.go new file mode 100644 index 000000000..bbedd75f9 --- /dev/null +++ b/collector/pg_connections_by_client_test.go @@ -0,0 +1,61 @@ +// Copyright 2024 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 TestPGConnectionsByClientCollector(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{"count", "client_name"}). + AddRow(42, "test") + + mock.ExpectQuery(sanitizeQuery(pgConnectionsByClientQuery)).WillReturnRows(rows) + + ch := make(chan prometheus.Metric) + go func() { + defer close(ch) + c := PGConnectionByClientCollector{} + if err := c.Update(context.Background(), inst, ch); err != nil { + t.Errorf("Error calling PGConnectionsByClientCollector.Update: %s", err) + } + }() + + expected := []MetricResult{ + {labels: labelMap{"name": "test"}, 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) + } +}