diff --git a/collector/pg_database.go b/collector/pg_database.go index d2c4b206a..30c4c8af0 100644 --- a/collector/pg_database.go +++ b/collector/pg_database.go @@ -53,12 +53,21 @@ var ( "Disk space used by the database", []string{"datname"}, nil, ) + pgDatabaseConnectionLimitsDesc = prometheus.NewDesc( + prometheus.BuildFQName( + namespace, + databaseSubsystem, + "connection_limit", + ), + "Connection limit set for the database", + []string{"datname"}, nil, + ) - pgDatabaseQuery = "SELECT pg_database.datname FROM pg_database;" + pgDatabaseQuery = "SELECT pg_database.datname, pg_database.datconnlimit FROM pg_database;" pgDatabaseSizeQuery = "SELECT pg_database_size($1)" ) -// Update implements Collector and exposes database size. +// Update implements Collector and exposes database size and connection limits. // It is called by the Prometheus registry when collecting metrics. // The list of databases is retrieved from pg_database and filtered // by the excludeDatabase config parameter. The tradeoff here is that @@ -81,21 +90,32 @@ func (c PGDatabaseCollector) Update(ctx context.Context, instance *instance, ch for rows.Next() { var datname sql.NullString - if err := rows.Scan(&datname); err != nil { + var connLimit sql.NullInt64 + if err := rows.Scan(&datname, &connLimit); err != nil { return err } if !datname.Valid { continue } + database := datname.String // Ignore excluded databases // Filtering is done here instead of in the query to avoid // a complicated NOT IN query with a variable number of parameters - if sliceContains(c.excludedDatabases, datname.String) { + if sliceContains(c.excludedDatabases, database) { continue } - databases = append(databases, datname.String) + databases = append(databases, database) + + connLimitMetric := 0.0 + if connLimit.Valid { + connLimitMetric = float64(connLimit.Int64) + } + ch <- prometheus.MustNewConstMetric( + pgDatabaseConnectionLimitsDesc, + prometheus.GaugeValue, connLimitMetric, database, + ) } // Query the size of the databases @@ -114,6 +134,7 @@ func (c PGDatabaseCollector) Update(ctx context.Context, instance *instance, ch pgDatabaseSizeDesc, prometheus.GaugeValue, sizeMetric, datname, ) + } return rows.Err() } diff --git a/collector/pg_database_test.go b/collector/pg_database_test.go index b5052c5d1..fe94166e9 100644 --- a/collector/pg_database_test.go +++ b/collector/pg_database_test.go @@ -31,8 +31,8 @@ func TestPGDatabaseCollector(t *testing.T) { inst := &instance{db: db} - mock.ExpectQuery(sanitizeQuery(pgDatabaseQuery)).WillReturnRows(sqlmock.NewRows([]string{"datname"}). - AddRow("postgres")) + mock.ExpectQuery(sanitizeQuery(pgDatabaseQuery)).WillReturnRows(sqlmock.NewRows([]string{"datname", "datconnlimit"}). + AddRow("postgres", 15)) mock.ExpectQuery(sanitizeQuery(pgDatabaseSizeQuery)).WithArgs("postgres").WillReturnRows(sqlmock.NewRows([]string{"pg_database_size"}). AddRow(1024)) @@ -47,6 +47,7 @@ func TestPGDatabaseCollector(t *testing.T) { }() expected := []MetricResult{ + {labels: labelMap{"datname": "postgres"}, value: 15, metricType: dto.MetricType_GAUGE}, {labels: labelMap{"datname": "postgres"}, value: 1024, metricType: dto.MetricType_GAUGE}, } convey.Convey("Metrics comparison", t, func() { @@ -71,8 +72,8 @@ func TestPGDatabaseCollectorNullMetric(t *testing.T) { inst := &instance{db: db} - mock.ExpectQuery(sanitizeQuery(pgDatabaseQuery)).WillReturnRows(sqlmock.NewRows([]string{"datname"}). - AddRow("postgres")) + mock.ExpectQuery(sanitizeQuery(pgDatabaseQuery)).WillReturnRows(sqlmock.NewRows([]string{"datname", "datconnlimit"}). + AddRow("postgres", nil)) mock.ExpectQuery(sanitizeQuery(pgDatabaseSizeQuery)).WithArgs("postgres").WillReturnRows(sqlmock.NewRows([]string{"pg_database_size"}). AddRow(nil)) @@ -88,6 +89,7 @@ func TestPGDatabaseCollectorNullMetric(t *testing.T) { expected := []MetricResult{ {labels: labelMap{"datname": "postgres"}, value: 0, metricType: dto.MetricType_GAUGE}, + {labels: labelMap{"datname": "postgres"}, value: 0, metricType: dto.MetricType_GAUGE}, } convey.Convey("Metrics comparison", t, func() { for _, expect := range expected { diff --git a/collector/pg_roles.go b/collector/pg_roles.go new file mode 100644 index 000000000..609c34c33 --- /dev/null +++ b/collector/pg_roles.go @@ -0,0 +1,91 @@ +// 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" + + "github.com/go-kit/log" + "github.com/prometheus/client_golang/prometheus" +) + +const rolesSubsystem = "roles" + +func init() { + registerCollector(rolesSubsystem, defaultEnabled, NewPGRolesCollector) +} + +type PGRolesCollector struct { + log log.Logger +} + +func NewPGRolesCollector(config collectorConfig) (Collector, error) { + return &PGRolesCollector{ + log: config.logger, + }, nil +} + +var ( + pgRolesConnectionLimitsDesc = prometheus.NewDesc( + prometheus.BuildFQName( + namespace, + rolesSubsystem, + "connection_limit", + ), + "Connection limit set for the role", + []string{"rolname"}, nil, + ) + + pgRolesConnectionLimitsQuery = "SELECT pg_roles.rolname, pg_roles.rolconnlimit FROM pg_roles" +) + +// Update implements Collector and exposes roles connection limits. +// It is called by the Prometheus registry when collecting metrics. +func (c PGRolesCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error { + db := instance.getDB() + // Query the list of databases + rows, err := db.QueryContext(ctx, + pgRolesConnectionLimitsQuery, + ) + if err != nil { + return err + } + defer rows.Close() + + for rows.Next() { + var rolname sql.NullString + var connLimit sql.NullInt64 + if err := rows.Scan(&rolname, &connLimit); err != nil { + return err + } + + if !rolname.Valid { + continue + } + rolnameLabel := rolname.String + + if !connLimit.Valid { + continue + } + connLimitMetric := float64(connLimit.Int64) + + ch <- prometheus.MustNewConstMetric( + pgRolesConnectionLimitsDesc, + prometheus.GaugeValue, connLimitMetric, rolnameLabel, + ) + } + + return rows.Err() +} diff --git a/collector/pg_roles_test.go b/collector/pg_roles_test.go new file mode 100644 index 000000000..182a120f9 --- /dev/null +++ b/collector/pg_roles_test.go @@ -0,0 +1,58 @@ +// 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 TestPGRolesCollector(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} + + mock.ExpectQuery(sanitizeQuery(pgRolesConnectionLimitsQuery)).WillReturnRows(sqlmock.NewRows([]string{"rolname", "rolconnlimit"}). + AddRow("postgres", 15)) + + ch := make(chan prometheus.Metric) + go func() { + defer close(ch) + c := PGRolesCollector{} + if err := c.Update(context.Background(), inst, ch); err != nil { + t.Errorf("Error calling PGRolesCollector.Update: %s", err) + } + }() + + expected := []MetricResult{ + {labels: labelMap{"rolname": "postgres"}, value: 15, 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) + } +}