feat: stats for notifier repo module (#6184)

# Description

- This PR adds comprehensive stats emission to all notifier repository
operations, providing better observability and monitoring capabilities
for database performance across the notifier module


## Linear Ticket

- Resolves WAR-922

## Security

- [x] The code changed/added as part of this pull request won't create
any security issues with how the software is being used.

Author: achettyiitr

services/notifier/notifier.go

@@ -227,7 +227,7 @@ func (n *Notifier) Setup(
 	if err := n.setupDatabase(ctx, dsn); err != nil {
 		return fmt.Errorf("could not setup db: %w", err)
 	}
-	n.repo = newRepo(n.db)
+	n.repo = newRepo(n.db, WithStats(n.statsFactory))
 
 	groupCtx, groupCancel := context.WithCancel(ctx)
 	n.background.group, n.background.groupCtx = errgroup.WithContext(groupCtx)

services/notifier/repo.go

@@ -10,6 +10,8 @@ import (
 
 	"github.com/lib/pq"
 
+	"github.com/rudderlabs/rudder-go-kit/stats"
+
 	"github.com/rudderlabs/rudder-go-kit/bytesize"
 	"github.com/rudderlabs/rudder-go-kit/config"
 
@@ -46,15 +48,23 @@ func WithNow(now func() time.Time) Opt {
 	}
 }
 
+func WithStats(s stats.Stats) Opt {
+	return func(r *repo) {
+		r.statsFactory = s
+	}
+}
+
 type repo struct {
-	db  *sqlmw.DB
-	now func() time.Time
+	db           *sqlmw.DB
+	now          func() time.Time
+	statsFactory stats.Stats
 }
 
 func newRepo(db *sqlmw.DB, opts ...Opt) *repo {
 	r := &repo{
-		db:  db,
-		now: timeutil.Now,
+		db:           db,
+		now:          timeutil.Now,
+		statsFactory: stats.NOP,
 	}
 	for _, opt := range opts {
 		opt(r)
@@ -67,6 +77,10 @@ func (n *repo) resetForWorkspace(
 	ctx context.Context,
 	workspaceIdentifier string,
 ) error {
+	defer n.timerStat("reset_for_workspace", stats.Tags{
+		"workspaceIdentifier": workspaceIdentifier,
+	})()
+
 	_, err := n.db.ExecContext(ctx, `
 		DELETE FROM `+notifierTableName+`
 		WHERE workspace = $1;
@@ -86,6 +100,11 @@ func (n *repo) insert(
 	workspaceIdentifier string,
 	batchID string,
 ) error {
+	defer n.timerStat("insert", stats.Tags{
+		"workspaceIdentifier": workspaceIdentifier,
+		"jobType":             string(publishRequest.JobType),
+	})()
+
 	txn, err := n.db.BeginTx(ctx, &sql.TxOptions{})
 	if err != nil {
 		return fmt.Errorf("inserting: begin transaction: %w", err)
@@ -166,6 +185,8 @@ func (n *repo) pendingByBatchID(
 	ctx context.Context,
 	batchID string,
 ) (int64, error) {
+	defer n.timerStat("pending_by_batch_id", nil)()
+
 	var count int64
 
 	err := n.db.QueryRowContext(ctx, `
@@ -195,6 +216,8 @@ func (n *repo) getByBatchID(
 	ctx context.Context,
 	batchID string,
 ) ([]Job, error) {
+	defer n.timerStat("get_by_batch_id", nil)()
+
 	query := `
 		SELECT
 			id,
@@ -300,6 +323,8 @@ func (n *repo) deleteByBatchID(
 	ctx context.Context,
 	batchID string,
 ) error {
+	defer n.timerStat("delete_by_batch_id", nil)()
+
 	_, err := n.db.ExecContext(ctx, `
 		DELETE FROM `+notifierTableName+` WHERE batch_id = $1;
 	`,
@@ -328,6 +353,10 @@ func (n *repo) claim(
 	ctx context.Context,
 	workerID string,
 ) (*Job, error) {
+	defer n.timerStat("claim", stats.Tags{
+		"workerID": workerID,
+	})()
+
 	row := n.db.QueryRowContext(ctx, `
 		UPDATE
   		  `+notifierTableName+`
@@ -376,6 +405,12 @@ func (n *repo) onClaimFailed(
 	claimError error,
 	maxAttempt int,
 ) error {
+	defer n.timerStat("on_claim_failed", stats.Tags{
+		"workspaceIdentifier": job.WorkspaceIdentifier,
+		"jobStatus":           string(job.Status),
+		"jobType":             string(job.Type),
+	})()
+
 	query := fmt.Sprint(`
 		UPDATE
 		  ` + notifierTableName + `
@@ -415,6 +450,12 @@ func (n *repo) onClaimSuccess(
 	job *Job,
 	payload json.RawMessage,
 ) error {
+	defer n.timerStat("on_claim_success", stats.Tags{
+		"workspaceIdentifier": job.WorkspaceIdentifier,
+		"jobStatus":           string(job.Status),
+		"jobType":             string(job.Type),
+	})()
+
 	_, err := n.db.ExecContext(ctx, `
 		UPDATE
 		  `+notifierTableName+`
@@ -442,6 +483,8 @@ func (n *repo) orphanJobIDs(
 	ctx context.Context,
 	intervalInSeconds int,
 ) ([]int64, error) {
+	defer n.timerStat("orphan_job_ids", nil)()
+
 	rows, err := n.db.QueryContext(ctx, `
 		UPDATE
           `+notifierTableName+`
@@ -489,6 +532,8 @@ func (n *repo) orphanJobIDs(
 }
 
 func (n *repo) refreshClaim(ctx context.Context, jobId int64) error {
+	defer n.timerStat("refresh_claim", nil)()
+
 	_, err := n.db.ExecContext(ctx, `
 		UPDATE
 		  `+notifierTableName+`
@@ -507,3 +552,13 @@ func (n *repo) refreshClaim(ctx context.Context, jobId int64) error {
 	}
 	return nil
 }
+
+// timerStat returns a function that records the duration of a database action.
+func (n *repo) timerStat(action string, extraTags stats.Tags) func() {
+	statName := "notifier_repo_query_duration_seconds"
+	tags := stats.Tags{"action": action, "repoType": notifierTableName}
+	for k, v := range extraTags {
+		tags[k] = v
+	}
+	return n.statsFactory.NewTaggedStat(statName, stats.TimerType, tags).RecordDuration()
+}

services/notifier/repo_test.go

@@ -16,6 +16,8 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/rudderlabs/rudder-go-kit/config"
+	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-go-kit/stats/memstats"
 	"github.com/rudderlabs/rudder-go-kit/testhelper/docker/resource/postgres"
 
 	migrator "github.com/rudderlabs/rudder-server/services/sql-migrator"
@@ -57,6 +59,9 @@ func TestRepo(t *testing.T) {
 			pool, err := dockertest.NewPool("")
 			require.NoError(t, err)
 
+			statsStore, err := memstats.New()
+			require.NoError(t, err)
+
 			pgResource, err := postgres.Setup(pool, t)
 			require.NoError(t, err)
 
@@ -80,7 +85,7 @@ func TestRepo(t *testing.T) {
 			ctx := context.Background()
 			now := time.Now().Truncate(time.Second).UTC()
 			db := sqlmw.New(pgResource.DB)
-			r := newRepo(db, WithNow(func() time.Time {
+			r := newRepo(db, WithStats(statsStore), WithNow(func() time.Time {
 				return now
 			}))
 
@@ -106,6 +111,12 @@ func TestRepo(t *testing.T) {
 
 					err := r.insert(ctx, &publishRequest, workspaceIdentifier, batchID)
 					require.NoError(t, err)
+					require.Greater(t, statsStore.Get("notifier_repo_query_duration_seconds", stats.Tags{
+						"action":              "insert",
+						"repoType":            "pg_notifier_queue",
+						"workspaceIdentifier": workspaceIdentifier,
+						"jobType":             string(publishRequest.JobType),
+					}).LastDuration(), time.Duration(0))
 
 					jobs, err := r.getByBatchID(ctx, batchID)
 					require.NoError(t, err)