chore: add db.sql stat collector

app/apphandlers/embeddedAppHandler.go

@@ -126,7 +126,7 @@ func (a *embeddedApp) StartRudderCore(ctx context.Context, options *app.Options)
 
 	fileUploaderProvider := fileuploader.NewProvider(ctx, backendconfig.DefaultBackendConfig)
 
-	rsourcesService, err := NewRsourcesService(deploymentType, true)
+	rsourcesService, err := NewRsourcesService(deploymentType, true, stats.Default)
 	if err != nil {
 		return err
 	}
@@ -316,7 +316,7 @@ func (a *embeddedApp) StartRudderCore(ctx context.Context, options *app.Options)
 	if err != nil {
 		return fmt.Errorf("failed to create gw rate limiter: %w", err)
 	}
-	drainConfigManager, err := drain_config.NewDrainConfigManager(config, a.log.Child("drain-config"))
+	drainConfigManager, err := drain_config.NewDrainConfigManager(config, a.log.Child("drain-config"), stats.Default)
 	if err != nil {
 		return fmt.Errorf("drain config manager setup: %v", err)
 	}

app/apphandlers/gatewayAppHandler.go

@@ -112,7 +112,7 @@ func (a *gatewayApp) StartRudderCore(ctx context.Context, options *app.Options)
 	if err != nil {
 		return fmt.Errorf("failed to create rate limiter: %w", err)
 	}
-	rsourcesService, err := NewRsourcesService(deploymentType, false)
+	rsourcesService, err := NewRsourcesService(deploymentType, false, stats.Default)
 	if err != nil {
 		return err
 	}
@@ -121,7 +121,7 @@ func (a *gatewayApp) StartRudderCore(ctx context.Context, options *app.Options)
 		TransformerURL:           config.GetString("DEST_TRANSFORM_URL", "http://localhost:9090"),
 		FeaturesRetryMaxAttempts: 10,
 	})
-	drainConfigManager, err := drain_config.NewDrainConfigManager(config, a.log.Child("drain-config"))
+	drainConfigManager, err := drain_config.NewDrainConfigManager(config, a.log.Child("drain-config"), stats.Default)
 	if err != nil {
 		a.log.Errorw("drain config manager setup failed while starting gateway", "error", err)
 	}

app/apphandlers/processorAppHandler.go

@@ -85,6 +85,7 @@ func (a *processorApp) Setup() error {
 
 func (a *processorApp) StartRudderCore(ctx context.Context, options *app.Options) error {
 	config := config.Default
+	statsFactory := stats.Default
 	if !a.setupDone {
 		return fmt.Errorf("processor service cannot start, database is not setup")
 	}
@@ -133,7 +134,7 @@ func (a *processorApp) StartRudderCore(ctx context.Context, options *app.Options
 
 	fileUploaderProvider := fileuploader.NewProvider(ctx, backendconfig.DefaultBackendConfig)
 
-	rsourcesService, err := NewRsourcesService(deploymentType, true)
+	rsourcesService, err := NewRsourcesService(deploymentType, true, statsFactory)
 	if err != nil {
 		return err
 	}
@@ -208,9 +209,9 @@ func (a *processorApp) StartRudderCore(ctx context.Context, options *app.Options
 			return err
 		}
 		defer client.Close()
-		schemaForwarder = schema_forwarder.NewForwarder(terminalErrFn, schemaDB, &client, backendconfig.DefaultBackendConfig, logger.NewLogger().Child("jobs_forwarder"), config, stats.Default)
+		schemaForwarder = schema_forwarder.NewForwarder(terminalErrFn, schemaDB, &client, backendconfig.DefaultBackendConfig, logger.NewLogger().Child("jobs_forwarder"), config, statsFactory)
 	} else {
-		schemaForwarder = schema_forwarder.NewAbortingForwarder(terminalErrFn, schemaDB, logger.NewLogger().Child("jobs_forwarder"), config, stats.Default)
+		schemaForwarder = schema_forwarder.NewAbortingForwarder(terminalErrFn, schemaDB, logger.NewLogger().Child("jobs_forwarder"), config, statsFactory)
 	}
 
 	modeProvider, err := resolveModeProvider(a.log, deploymentType)
@@ -220,7 +221,7 @@ func (a *processorApp) StartRudderCore(ctx context.Context, options *app.Options
 
 	adaptiveLimit := payload.SetupAdaptiveLimiter(ctx, g)
 
-	enrichers, err := setupPipelineEnrichers(config, a.log, stats.Default)
+	enrichers, err := setupPipelineEnrichers(config, a.log, statsFactory)
 	if err != nil {
 		return fmt.Errorf("setting up pipeline enrichers: %w", err)
 	}
@@ -231,7 +232,7 @@ func (a *processorApp) StartRudderCore(ctx context.Context, options *app.Options
 		}
 	}()
 
-	drainConfigManager, err := drain_config.NewDrainConfigManager(config, a.log.Child("drain-config"))
+	drainConfigManager, err := drain_config.NewDrainConfigManager(config, a.log.Child("drain-config"), statsFactory)
 	if err != nil {
 		return fmt.Errorf("drain config manager setup: %v", err)
 	}
@@ -264,7 +265,7 @@ func (a *processorApp) StartRudderCore(ctx context.Context, options *app.Options
 		trackedUsersReporter,
 		proc.WithAdaptiveLimit(adaptiveLimit),
 	)
-	throttlerFactory, err := throttler.NewFactory(config, stats.Default)
+	throttlerFactory, err := throttler.NewFactory(config, statsFactory)
 	if err != nil {
 		return fmt.Errorf("failed to create throttler factory: %w", err)
 	}
@@ -309,7 +310,7 @@ func (a *processorApp) StartRudderCore(ctx context.Context, options *app.Options
 			archivalDB,
 			fileUploaderProvider,
 			config,
-			stats.Default,
+			statsFactory,
 			archiver.WithAdaptiveLimit(adaptiveLimit),
 		),
 	}
@@ -330,8 +331,8 @@ func (a *processorApp) StartRudderCore(ctx context.Context, options *app.Options
 	})
 
 	g.Go(func() error {
-		replicationLagStat := stats.Default.NewStat("rsources_log_replication_lag", stats.GaugeType)
-		replicationSlotStat := stats.Default.NewStat("rsources_log_replication_slot", stats.GaugeType)
+		replicationLagStat := statsFactory.NewStat("rsources_log_replication_lag", stats.GaugeType)
+		replicationSlotStat := statsFactory.NewStat("rsources_log_replication_slot", stats.GaugeType)
 		rsourcesService.Monitor(ctx, replicationLagStat, replicationSlotStat)
 		return nil
 	})

app/apphandlers/setup.go

@@ -53,7 +53,7 @@ func rudderCoreNodeSetup() error {
 }
 
 // NewRsourcesService produces a rsources.JobService through environment configuration (env variables & config file)
-func NewRsourcesService(deploymentType deployment.Type, shouldSetupSharedDB bool) (rsources.JobService, error) {
+func NewRsourcesService(deploymentType deployment.Type, shouldSetupSharedDB bool, stats stats.Stats) (rsources.JobService, error) {
 	var rsourcesConfig rsources.JobServiceConfig
 	rsourcesConfig.MaxPoolSize = config.GetInt("Rsources.MaxPoolSize", 3)
 	rsourcesConfig.MinPoolSize = config.GetInt("Rsources.MinPoolSize", 1)
@@ -80,7 +80,7 @@ func NewRsourcesService(deploymentType deployment.Type, shouldSetupSharedDB bool
 
 	rsourcesConfig.ShouldSetupSharedDB = shouldSetupSharedDB
 
-	return rsources.NewJobService(rsourcesConfig)
+	return rsources.NewJobService(rsourcesConfig, stats)
 }
 
 func resolveModeProvider(log logger.Logger, deploymentType deployment.Type) (cluster.ChangeEventProvider, error) {

enterprise/reporting/error_index/error_index_reporting.go

@@ -15,6 +15,7 @@ import (
 	"github.com/rudderlabs/rudder-go-kit/filemanager"
 	"github.com/rudderlabs/rudder-go-kit/logger"
 	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-go-kit/stats/collectors"
 	kitsync "github.com/rudderlabs/rudder-go-kit/sync"
 
 	"github.com/rudderlabs/rudder-server/jobsdb"
@@ -188,6 +189,10 @@ func (eir *ErrorIndexReporter) DatabaseSyncer(c types.SyncerConfig) types.Report
 	if err != nil {
 		panic(fmt.Errorf("failed to open error index db: %w", err))
 	}
+	err = eir.statsFactory.RegisterCollector(collectors.NewDatabaseSQLStats("jobsdb-err_idx", dbHandle))
+	if err != nil {
+		panic(fmt.Errorf("failed to register collector: %w", err))
+	}
 	errIndexDB := jobsdb.NewForReadWrite(
 		"err_idx",
 		jobsdb.WithDBHandle(dbHandle),

enterprise/reporting/error_reporting.go

@@ -25,6 +25,7 @@ import (
 	"github.com/rudderlabs/rudder-go-kit/config"
 	"github.com/rudderlabs/rudder-go-kit/logger"
 	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-go-kit/stats/collectors"
 
 	obskit "github.com/rudderlabs/rudder-observability-kit/go/labels"
 	migrator "github.com/rudderlabs/rudder-server/services/sql-migrator"
@@ -290,6 +291,10 @@ func (edr *ErrorDetailReporter) migrate(c types.SyncerConfig) (*sql.DB, error) {
 		return nil, err
 	}
 	dbHandle.SetMaxOpenConns(edr.maxOpenConnections)
+	err = edr.stats.RegisterCollector(collectors.NewDatabaseSQLStats("error_detail_reporting", dbHandle))
+	if err != nil {
+		return nil, fmt.Errorf("could not register database stats collector: %w", err)
+	}
 
 	m := &migrator.Migrator{
 		Handle:          dbHandle,

enterprise/reporting/flusher/factory.go

@@ -12,6 +12,7 @@ import (
 	"github.com/rudderlabs/rudder-go-kit/config"
 	"github.com/rudderlabs/rudder-go-kit/logger"
 	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-go-kit/stats/collectors"
 	"github.com/rudderlabs/rudder-server/enterprise/reporting/flusher/aggregator"
 	"github.com/rudderlabs/rudder-server/utils/misc"
 )
@@ -28,9 +29,13 @@ func CreateRunner(ctx context.Context, table string, log logger.Logger, stats st
 	maxOpenConns := conf.GetIntVar(4, 1, "Reporting.flusher.maxOpenConnections")
 	db, err := sql.Open("postgres", connStr)
 	if err != nil {
-		return nil, fmt.Errorf("error opening postgres %w", err)
+		return nil, fmt.Errorf("opening postgres: %w", err)
 	}
 	db.SetMaxOpenConns(maxOpenConns)
+	err = stats.RegisterCollector(collectors.NewDatabaseSQLStats(fmt.Sprintf("reporting_flusher_%s", table), db))
+	if err != nil {
+		return nil, fmt.Errorf("registering collector: %w", err)
+	}
 
 	if table == "tracked_users_reports" {
 

enterprise/trackedusers/users_reporter.go

@@ -10,6 +10,7 @@ import (
 	"github.com/rudderlabs/rudder-server/utils/timeutil"
 
 	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-go-kit/stats/collectors"
 
 	"github.com/rudderlabs/rudder-server/jobsdb"
 	migrator "github.com/rudderlabs/rudder-server/services/sql-migrator"
@@ -85,6 +86,10 @@ func (u *UniqueUsersReporter) MigrateDatabase(dbConn string, conf *config.Config
 		return err
 	}
 	dbHandle.SetMaxOpenConns(1)
+	err = u.stats.RegisterCollector(collectors.NewDatabaseSQLStats("tracked_users_reports", dbHandle))
+	if err != nil {
+		return fmt.Errorf("registering collector: %w", err)
+	}
 
 	m := &migrator.Migrator{
 		Handle:                     dbHandle,

internal/drain-config/drainConfig.go

@@ -11,6 +11,8 @@ import (
 
 	"github.com/rudderlabs/rudder-go-kit/config"
 	"github.com/rudderlabs/rudder-go-kit/logger"
+	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-go-kit/stats/collectors"
 	migrator "github.com/rudderlabs/rudder-server/services/sql-migrator"
 	"github.com/rudderlabs/rudder-server/utils/misc"
 )
@@ -42,8 +44,8 @@ type drainConfigManager struct {
 // NewDrainConfigManager returns a drainConfigManager
 // If migration fails while setting up drain config, drainConfigManager object will be returned along with error
 // Consumers must handle errors and non-nil drainConfigManager object according to their use case.
-func NewDrainConfigManager(conf *config.Config, log logger.Logger) (*drainConfigManager, error) {
-	db, err := setupDBConn(conf)
+func NewDrainConfigManager(conf *config.Config, log logger.Logger, stats stats.Stats) (*drainConfigManager, error) {
+	db, err := setupDBConn(conf, stats)
 	if err != nil {
 		log.Errorw("db setup", "error", err)
 		return nil, fmt.Errorf("db setup: %v", err)
@@ -182,7 +184,7 @@ func migrate(db *sql.DB) error {
 }
 
 // setupDBConn sets up the database connection
-func setupDBConn(conf *config.Config) (*sql.DB, error) {
+func setupDBConn(conf *config.Config, stats stats.Stats) (*sql.DB, error) {
 	psqlInfo := misc.GetConnectionString(conf, "drain-config")
 	if conf.IsSet("SharedDB.dsn") {
 		psqlInfo = conf.GetString("SharedDB.dsn", "")
@@ -193,5 +195,9 @@ func setupDBConn(conf *config.Config) (*sql.DB, error) {
 	}
 	db.SetMaxIdleConns(conf.GetInt("drainConfig.maxIdleConns", 1))
 	db.SetMaxOpenConns(conf.GetInt("drainConfig.maxOpenConns", 2))
+	err = stats.RegisterCollector(collectors.NewDatabaseSQLStats("drain_config", db))
+	if err != nil {
+		return nil, fmt.Errorf("registering collector: %v", err)
+	}
 	return db, nil
 }

internal/drain-config/drainConfig_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/rudderlabs/rudder-go-kit/bytesize"
 	"github.com/rudderlabs/rudder-go-kit/config"
 	"github.com/rudderlabs/rudder-go-kit/logger"
+	"github.com/rudderlabs/rudder-go-kit/stats/memstats"
 	"github.com/rudderlabs/rudder-go-kit/testhelper/docker/resource/postgres"
 	drain_config "github.com/rudderlabs/rudder-server/internal/drain-config"
 )
@@ -25,8 +26,10 @@ func TestDrainConfigRoutine(t *testing.T) {
 	conf.Set("drainConfig.cleanupFrequency", "50ms")
 	ctx := context.Background()
 	log := logger.NOP
+	sts, err := memstats.New()
+	require.NoError(t, err, "should create stats")
 
-	drainConfigManager, err := drain_config.NewDrainConfigManager(conf, log)
+	drainConfigManager, err := drain_config.NewDrainConfigManager(conf, log, sts)
 	require.NoError(t, err, "should create drain config manager")
 	go func() {
 		require.NoError(t, drainConfigManager.DrainConfigRoutine(ctx), "drain config routine should exit")
@@ -55,8 +58,10 @@ func TestDrainConfigHttpHandler(t *testing.T) {
 	conf.Set("drainConfig.pollFrequency", "100ms")
 	ctx, cancel := context.WithCancel(context.Background())
 	log := logger.NOP
+	sts, err := memstats.New()
+	require.NoError(t, err, "should create stats")
 
-	drainConfigManager, err := drain_config.NewDrainConfigManager(conf, log)
+	drainConfigManager, err := drain_config.NewDrainConfigManager(conf, log, sts)
 	require.NoError(t, err, "should create drain config manager")
 	go func() {
 		require.NoError(t, drainConfigManager.DrainConfigRoutine(ctx), "drain config routine should exit")

jobsdb/jobsdb.go

@@ -53,6 +53,8 @@ import (
 
 	"github.com/rudderlabs/rudder-go-kit/config"
 	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-go-kit/stats/collectors"
+
 	"github.com/rudderlabs/rudder-server/services/rmetrics"
 	"github.com/rudderlabs/rudder-server/utils/misc"
 
@@ -633,6 +635,20 @@ const (
 	ReadWrite OwnerType = ""
 )
 
+func (ot OwnerType) Identifier() string {
+	switch ot {
+	case Read:
+		return "r"
+	case Write:
+		return "w"
+	case ReadWrite:
+		return "rw"
+	default:
+		return ""
+
+	}
+}
+
 func init() {
 	for _, js := range jobStates {
 		if !js.isValid {
@@ -747,6 +763,15 @@ func (jd *Handle) init() {
 		psqlInfo := misc.GetConnectionString(jd.config, "jobsdb_"+jd.tablePrefix)
 		jd.dbHandle, err = sql.Open("postgres", psqlInfo)
 		jd.assertError(err)
+
+		jd.assertError(
+			stats.Default.RegisterCollector(
+				collectors.NewDatabaseSQLStats(
+					"jobsdb_"+jd.tablePrefix+"_"+jd.ownerType.Identifier(),
+					jd.dbHandle,
+				),
+			),
+		)
 	}
 
 	var maxConns int

processor/processor_test.go

@@ -4076,7 +4076,7 @@ var _ = Describe("Static Function Tests", func() {
 					},
 				},
 			}
-			response := ConvertToFilteredTransformerResponse(events, true, func(event transformer.TransformerEvent) (bool, string) { return false, "" })
+			response := ConvertToFilteredTransformerResponse(events, true, func(_ transformer.TransformerEvent) (bool, string) { return false, "" })
 			Expect(response).To(Equal(expectedResponse))
 		})
 

services/notifier/notifier.go

@@ -12,6 +12,7 @@ import (
 	"github.com/lib/pq"
 
 	"github.com/rudderlabs/rudder-go-kit/sqlutil"
+	"github.com/rudderlabs/rudder-go-kit/stats/collectors"
 
 	"golang.org/x/sync/errgroup"
 
@@ -266,6 +267,10 @@ func (n *Notifier) setupDatabase(
 		return fmt.Errorf("could not open: %w", err)
 	}
 	database.SetMaxOpenConns(n.config.maxOpenConnections)
+	err = n.statsFactory.RegisterCollector(collectors.NewDatabaseSQLStats("notifier-"+n.workspaceIdentifier, database))
+	if err != nil {
+		return fmt.Errorf("registering collector: %w", err)
+	}
 
 	if err := database.PingContext(ctx); err != nil {
 		return fmt.Errorf("could not ping: %w", err)