diff --git a/collector/collector.go b/collector/collector.go index 192a31238..bce9ed376 100644 --- a/collector/collector.go +++ b/collector/collector.go @@ -26,12 +26,18 @@ import ( "github.com/prometheus/client_golang/prometheus" ) +type StatStatementsConfig struct { + IncludeQuery bool + QueryLength uint + Limit uint + ExcludeDatabases []string + ExcludeUsers []string +} + var ( - factories = make(map[string]func(collectorConfig) (Collector, error)) - initiatedCollectorsMtx = sync.Mutex{} - initiatedCollectors = make(map[string]Collector) - collectorState = make(map[string]*bool) - forcedCollectors = map[string]bool{} // collectors which have been explicitly enabled or disabled + factories = make(map[string]func(collectorConfig) (Collector, error)) + collectorState = make(map[string]*bool) + forcedCollectors = map[string]bool{} // collectors which have been explicitly enabled or disabled ) const ( @@ -65,6 +71,8 @@ type Collector interface { type collectorConfig struct { logger *slog.Logger excludeDatabases []string + + statStatementsConfig *StatStatementsConfig } func registerCollector(name string, isDefaultEnabled bool, createFunc func(collectorConfig) (Collector, error)) { @@ -83,6 +91,8 @@ type PostgresCollector struct { instance *instance CollectionTimeout time.Duration + + statStatementsConfig *StatStatementsConfig } type Option func(*PostgresCollector) error @@ -100,40 +110,34 @@ func NewPostgresCollector(logger *slog.Logger, excludeDatabases []string, dsn st } } - collectorsToStart := map[string]struct{}{} - if len(enabledCollectors) == 0 { - for name, enabledByDefault := range collectorState { - if *enabledByDefault { - collectorsToStart[name] = struct{}{} - } - } - } else { - for _, name := range enabledCollectors { - _, exist := collectorState[name] - if !exist { - return nil, fmt.Errorf("requested to enable an unknown collector: %s", name) - } - collectorsToStart[name] = struct{}{} + f := make(map[string]bool) + for _, name := range enabledCollectors { + _, exist := collectorState[name] + if !exist { + return nil, fmt.Errorf("requested to enable an unknown collector: %s", name) } + f[name] = true } - collectors := make(map[string]Collector) - initiatedCollectorsMtx.Lock() - defer initiatedCollectorsMtx.Unlock() - for key := range collectorsToStart { - if collector, ok := initiatedCollectors[key]; ok { - collectors[key] = collector - } else { - collector, err := factories[key](collectorConfig{ - logger: logger.With("collector", key), - excludeDatabases: excludeDatabases, - }) - if err != nil { - return nil, err + for key, enabled := range collectorState { + // When enabledCollectors are specified, only start those collectors (regardless of default enabled state). + // When no enabledCollectors are specified, start all collectors that are enabled by default. + if len(f) > 0 { + if !f[key] { + continue } - collectors[key] = collector - initiatedCollectors[key] = collector + } else if !*enabled { + continue } + collector, err := factories[key](collectorConfig{ + logger: logger.With("collector", key), + excludeDatabases: excludeDatabases, + statStatementsConfig: p.statStatementsConfig, + }) + if err != nil { + return nil, err + } + collectors[key] = collector } p.Collectors = collectors @@ -165,6 +169,13 @@ func WithCollectionTimeout(s string) Option { } } +func WithStatStatementsConfig(cfg StatStatementsConfig) Option { + return func(e *PostgresCollector) error { + e.statStatementsConfig = &cfg + return nil + } +} + // Describe implements the prometheus.Collector interface. func (p PostgresCollector) Describe(ch chan<- *prometheus.Desc) { ch <- scrapeDurationDesc diff --git a/collector/pg_stat_statements.go b/collector/pg_stat_statements.go index 9227f5c04..abb3eb467 100644 --- a/collector/pg_stat_statements.go +++ b/collector/pg_stat_statements.go @@ -81,6 +81,19 @@ type PGStatStatementsCollector struct { } func NewPGStatStatementsCollector(config collectorConfig) (Collector, error) { + // When config is provided, use it directly instead of global kingpin flags + if config.statStatementsConfig != nil { + return &PGStatStatementsCollector{ + log: config.logger, + includeQueryStatement: config.statStatementsConfig.IncludeQuery, + statementLength: config.statStatementsConfig.QueryLength, + statementLimit: config.statStatementsConfig.Limit, + excludedDatabases: config.statStatementsConfig.ExcludeDatabases, + excludedUsers: config.statStatementsConfig.ExcludeUsers, + }, nil + } + + // Fall back to kingpin CLI flags for standalone postgres_exporter binary usage. var excludedDatabases []string if *excludedDatabasesFlag != "" { for db := range strings.SplitSeq(*excludedDatabasesFlag, ",") { diff --git a/collector/probe.go b/collector/probe.go index e40d6fee1..5e243997c 100644 --- a/collector/probe.go +++ b/collector/probe.go @@ -31,8 +31,6 @@ type ProbeCollector struct { func NewProbeCollector(logger *slog.Logger, excludeDatabases []string, registry *prometheus.Registry, dsn config.DSN) (*ProbeCollector, error) { collectors := make(map[string]Collector) - initiatedCollectorsMtx.Lock() - defer initiatedCollectorsMtx.Unlock() for key, enabled := range collectorState { // TODO: Handle filters // if !*enabled || (len(f) > 0 && !f[key]) { @@ -41,20 +39,15 @@ func NewProbeCollector(logger *slog.Logger, excludeDatabases []string, registry if !*enabled { continue } - if collector, ok := initiatedCollectors[key]; ok { - collectors[key] = collector - } else { - collector, err := factories[key]( - collectorConfig{ - logger: logger.With("collector", key), - excludeDatabases: excludeDatabases, - }) - if err != nil { - return nil, err - } - collectors[key] = collector - initiatedCollectors[key] = collector + collector, err := factories[key]( + collectorConfig{ + logger: logger.With("collector", key), + excludeDatabases: excludeDatabases, + }) + if err != nil { + return nil, err } + collectors[key] = collector } instance, err := newInstance(dsn.GetConnectionString())