fix: refactor collectors initialization to support config flags

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

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, ",") {

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())