diff --git a/collector/pg_table.go b/collector/pg_table.go new file mode 100644 index 000000000..f2b26ad16 --- /dev/null +++ b/collector/pg_table.go @@ -0,0 +1,139 @@ +// 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 tableSizeSubsystem = "table_size" + +func init() { + registerCollector(tableSizeSubsystem, defaultEnabled, NewPGTableSizeCollector) +} + +type PGTableSizeCollector struct { + log log.Logger +} + +func NewPGTableSizeCollector(config collectorConfig) (Collector, error) { + return &PGTableSizeCollector{ + log: config.logger, + }, nil +} + +var ( + pgTableTotalRelationDesc = prometheus.NewDesc( + prometheus.BuildFQName( + namespace, + tableSizeSubsystem, + "total_relation", + ), + "Total Relation Size of the table", + []string{"schemaname", "datname", "relname"}, nil, + ) + pgTableIndexSizeDesc = prometheus.NewDesc( + prometheus.BuildFQName( + namespace, + tableSizeSubsystem, + "index", + ), + "Indexes Size of the Table", + []string{"schemaname", "datname", "relname"}, nil, + ) + pgRelationSizeDesc = prometheus.NewDesc( + prometheus.BuildFQName( + namespace, + tableSizeSubsystem, + "relation", + ), + "Relation Size of the table", + []string{"schemaname", "datname", "relname"}, nil, + ) + pgTableSizeQuery = `SELECT + table_catalog datname, + table_name relname, + table_schema schemaname, + pg_total_relation_size('"'||table_schema||'"."'||table_name||'"') total_relation_size, + pg_relation_size('"'||table_schema||'"."'||table_name||'"') relation_size, + pg_indexes_size('"'||table_schema||'"."'||table_name||'"') indexes_size + FROM information_schema.tables` +) + +// Update implements Collector and exposes database locks. +// It is called by the Prometheus registry when collecting metrics. +func (c PGTableSizeCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error { + db := instance.getDB() + // Query the list of databases + rows, err := db.QueryContext(ctx, pgTableSizeQuery) + if err != nil { + return err + } + defer rows.Close() + for rows.Next() { + var tableSchema, tableName, databaseName sql.NullString + var totalRelationSize, relationSize, indexesSize sql.NullInt64 + + if err := rows.Scan(&databaseName, &tableName, &tableSchema, &totalRelationSize, &relationSize, &indexesSize); err != nil { + return err + } + + if !tableSchema.Valid || !tableName.Valid || !databaseName.Valid { + continue + } + + totalRelationsSizeMetric := 0.0 + relationSizeMetric := 0.0 + indexesSizeMetric := 0.0 + if totalRelationSize.Valid { + totalRelationsSizeMetric = float64(totalRelationSize.Int64) + } + + if relationSize.Valid { + relationSizeMetric = float64(relationSize.Int64) + } + + if indexesSize.Valid { + indexesSizeMetric = float64(indexesSize.Int64) + } + + ch <- prometheus.MustNewConstMetric( + pgTableTotalRelationDesc, + prometheus.CounterValue, totalRelationsSizeMetric, + tableSchema.String, databaseName.String, tableName.String, + ) + + ch <- prometheus.MustNewConstMetric( + pgRelationSizeDesc, + prometheus.CounterValue, relationSizeMetric, + tableSchema.String, databaseName.String, tableName.String, + ) + + ch <- prometheus.MustNewConstMetric( + pgTableIndexSizeDesc, + prometheus.CounterValue, indexesSizeMetric, + tableSchema.String, databaseName.String, tableName.String, + ) + } + + if err := rows.Err(); err != nil { + return err + } + + return nil +} diff --git a/collector/pg_table_test.go b/collector/pg_table_test.go new file mode 100644 index 000000000..617583b30 --- /dev/null +++ b/collector/pg_table_test.go @@ -0,0 +1,54 @@ +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 TestPGTableSizeCollector(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", "relname", "schemaname", "total_relation_size", "relation_size", "indexes_size"}). + AddRow("test", "testrel", "testschema", 69, 42, 27). + AddRow("test2", "testrel2", "testschema2", 14, 10, 2) + + mock.ExpectQuery(sanitizeQuery(pgTableSizeQuery)).WillReturnRows(rows) + + ch := make(chan prometheus.Metric) + go func() { + defer close(ch) + c := PGTableSizeCollector{} + if err := c.Update(context.Background(), inst, ch); err != nil { + t.Errorf("Error calling PGTableSizeCollector.Update: %s", err) + } + }() + + expected := []MetricResult{ + {labels: labelMap{"datname": "test", "relname": "testrel", "schemaname": "testschema"}, value: 69, metricType: dto.MetricType_COUNTER}, + {labels: labelMap{"datname": "test", "relname": "testrel", "schemaname": "testschema"}, value: 42, metricType: dto.MetricType_COUNTER}, + {labels: labelMap{"datname": "test", "relname": "testrel", "schemaname": "testschema"}, value: 27, metricType: dto.MetricType_COUNTER}, + {labels: labelMap{"datname": "test2", "relname": "testrel2", "schemaname": "testschema2"}, value: 14, metricType: dto.MetricType_COUNTER}, + {labels: labelMap{"datname": "test2", "relname": "testrel2", "schemaname": "testschema2"}, value: 10, metricType: dto.MetricType_COUNTER}, + {labels: labelMap{"datname": "test2", "relname": "testrel2", "schemaname": "testschema2"}, value: 2, metricType: dto.MetricType_COUNTER}, + } + 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) + } +}