refactor(metrics): rename MetricsPayload to MetricsCollection and enhance metrics collection structure

Author: leo-marble

models/metrics_collection.go

@@ -24,7 +24,7 @@ type MetricData struct {
 	Frequency      MetricCollectionFrequency
 }
 
-type MetricsPayload struct {
+type MetricsCollection struct {
 	CollectionID uuid.UUID // Unique ID for this collection run, could be use as idempotency key
 	Timestamp    time.Time
 	Metrics      []MetricData

usecases/metrics_collection/collector.go

@@ -2,11 +2,20 @@ package metrics_collection
 
 import (
 	"context"
+	"slices"
 	"time"
 
 	"github.com/checkmarble/marble-backend/models"
+	"github.com/checkmarble/marble-backend/repositories"
+	"github.com/checkmarble/marble-backend/usecases/executor_factory"
+	"github.com/checkmarble/marble-backend/utils"
+	"github.com/google/uuid"
 )
 
+type OrganizationRepository interface {
+	AllOrganizations(ctx context.Context, exec repositories.Executor) ([]models.Organization, error)
+}
+
 // GlobalCollector is a collector that is not specific to an organization.
 // It is used to collect metrics that are not specific to an organization.
 // For example, the app version, the number of users
@@ -23,23 +32,97 @@ type Collectors struct {
 	version          string
 	globalCollectors []GlobalCollector
 	collectors       []Collector
+
+	organizationRepository OrganizationRepository
+	executorFactory        executor_factory.ExecutorFactory
 }
 
-func (c Collectors) GetGlobalCollectors() []GlobalCollector {
-	return c.globalCollectors
+func (c Collectors) CollectMetrics(ctx context.Context, from time.Time, to time.Time) (models.MetricsCollection, error) {
+	metrics := make([]models.MetricData, 0)
+
+	// Collect global metrics
+	globalMetrics, err := c.collectGlobalMetrics(ctx, from, to)
+	if err != nil {
+		return models.MetricsCollection{}, err
+	}
+	metrics = slices.Concat(metrics, globalMetrics)
+
+	// Collect organization-specific metrics
+	orgMetrics, err := c.collectOrganizationMetrics(ctx, from, to)
+	if err != nil {
+		return models.MetricsCollection{}, err
+	}
+	metrics = slices.Concat(metrics, orgMetrics)
+
+	payload := models.MetricsCollection{
+		CollectionID: uuid.New(),
+		Timestamp:    time.Now(),
+		Metrics:      metrics,
+		Version:      c.version,
+	}
+
+	return payload, nil
 }
 
-func (c Collectors) GetCollectors() []Collector {
-	return c.collectors
+// Collects global metrics from all collectors
+// If a collector fails, it will log a warning and continue to the next collector (don't fail the whole function)
+func (c Collectors) collectGlobalMetrics(ctx context.Context, from time.Time, to time.Time) ([]models.MetricData, error) {
+	metrics := make([]models.MetricData, 0)
+	logger := utils.LoggerFromContext(ctx)
+
+	for _, collector := range c.globalCollectors {
+		value, err := collector.Collect(ctx, from, to)
+		if err != nil {
+			logger.WarnContext(ctx, "Failed to collect global metrics", "error", err)
+			continue
+		}
+		metrics = slices.Concat(metrics, value)
+	}
+
+	return metrics, nil
+}
+
+// Collects organization metrics from all collectors, fetching all organizations from the database first
+// If a collector fails, it will log a warning and continue to the next collector (don't fail the whole function)
+func (c Collectors) collectOrganizationMetrics(ctx context.Context, from time.Time, to time.Time) ([]models.MetricData, error) {
+	metrics := make([]models.MetricData, 0)
+	logger := utils.LoggerFromContext(ctx)
+
+	orgs, err := c.getListOfOrganizations(ctx)
+	if err != nil {
+		return []models.MetricData{}, err
+	}
+
+	for _, org := range orgs {
+		for _, collector := range c.collectors {
+			value, err := collector.Collect(ctx, org.Id, from, to)
+			if err != nil {
+				logger.WarnContext(ctx, "Failed to collect organization metrics", "error", err)
+				continue
+			}
+			metrics = slices.Concat(metrics, value)
+		}
+	}
+
+	return metrics, nil
 }
 
-func (c Collectors) GetVersion() string {
-	return c.version
+// Fetches all organizations from the database
+// NOTE: Add caching to avoid fetching the same organizations every time (but how can we invalidate the cache?)
+func (c Collectors) getListOfOrganizations(ctx context.Context) ([]models.Organization, error) {
+	orgs, err := c.organizationRepository.AllOrganizations(ctx, c.executorFactory.NewExecutor())
+	if err != nil {
+		return []models.Organization{}, err
+	}
+	return orgs, nil
 }
 
 // Use version to track the version of the collectors, could be used to track changes
 // and tell the server which collectors is used by the client
-func NewCollectorsV1() Collectors {
+func NewCollectorsV1(
+	executorFactory executor_factory.ExecutorFactory,
+	organizationRepository OrganizationRepository,
+) Collectors {
 	return Collectors{
 		version: "v1",
 		collectors: []Collector{
@@ -48,5 +131,7 @@ func NewCollectorsV1() Collectors {
 		globalCollectors: []GlobalCollector{
 			NewStubGlobalCollector(),
 		},
+		executorFactory:        executorFactory,
+		organizationRepository: organizationRepository,
 	}
 }

usecases/scheduled_execution/metrics_collection_job.go

@@ -2,15 +2,10 @@ package scheduled_execution
 
 import (
 	"context"
-	"slices"
 	"time"
 
 	"github.com/checkmarble/marble-backend/models"
-	"github.com/checkmarble/marble-backend/repositories"
-	"github.com/checkmarble/marble-backend/usecases/executor_factory"
-	"github.com/checkmarble/marble-backend/usecases/metrics_collection"
 	"github.com/checkmarble/marble-backend/utils"
-	"github.com/google/uuid"
 	"github.com/riverqueue/river"
 )
 
@@ -33,23 +28,21 @@ func NewMetricsCollectionPeriodicJob() *river.PeriodicJob {
 	)
 }
 
+type MetricsCollectionUsecase interface {
+	CollectMetrics(ctx context.Context, from time.Time, to time.Time) (models.MetricsCollection, error)
+}
+
 type MetricCollectionWorker struct {
 	river.WorkerDefaults[models.MetricsCollectionArgs]
 
-	executorFactory        executor_factory.ExecutorFactory
-	organizationRepository repositories.OrganizationRepository
-	collectors             metrics_collection.Collectors
+	collectors MetricsCollectionUsecase
 }
 
 func NewMetricCollectionWorker(
-	executorFactory executor_factory.ExecutorFactory,
-	organizationRepository repositories.OrganizationRepository,
-	collectors metrics_collection.Collectors,
+	collectors MetricsCollectionUsecase,
 ) MetricCollectionWorker {
 	return MetricCollectionWorker{
-		executorFactory:        executorFactory,
-		organizationRepository: organizationRepository,
-		collectors:             collectors,
+		collectors: collectors,
 	}
 }
 
@@ -66,100 +59,20 @@ func (w MetricCollectionWorker) Work(ctx context.Context, job *river.Job[models.
 	from := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)
 
 	// Create the metric collection usecase
-	metrics, err := w.collectMetrics(ctx, from, now)
+	metricsCollection, err := w.collectors.CollectMetrics(ctx, from, now)
 	if err != nil {
 		logger.ErrorContext(ctx, "Failed to collect metrics", "error", err)
 		return err
 	}
 
-	logger.DebugContext(ctx, "Collected metrics", "metrics", metrics)
+	logger.DebugContext(ctx, "Collected metrics", "metrics", metricsCollection)
 
 	// TODO: Update the watermarks with now value
 	logger.DebugContext(ctx, "Updating watermarks", "new_timestamp", now)
 
 	// TODO: Store or send the metrics somewhere
 	// For now, just log the number of metrics collected
-	logger.DebugContext(ctx, "Metric collection completed", "metrics_count", len(metrics.Metrics))
+	logger.DebugContext(ctx, "Metric collection completed", "metrics_count", len(metricsCollection.Metrics))
 
 	return nil
 }
-
-func (w *MetricCollectionWorker) collectMetrics(ctx context.Context, from time.Time, to time.Time) (models.MetricsPayload, error) {
-	metrics := make([]models.MetricData, 0)
-
-	// Collect global metrics
-	globalMetrics, err := w.collectGlobalMetrics(ctx, from, to)
-	if err != nil {
-		return models.MetricsPayload{}, err
-	}
-	metrics = slices.Concat(metrics, globalMetrics)
-
-	// Collect organization-specific metrics
-	orgMetrics, err := w.collectOrganizationMetrics(ctx, from, to)
-	if err != nil {
-		return models.MetricsPayload{}, err
-	}
-	metrics = slices.Concat(metrics, orgMetrics)
-
-	payload := models.MetricsPayload{
-		CollectionID: uuid.New(),
-		Timestamp:    time.Now(),
-		Metrics:      metrics,
-		Version:      w.collectors.GetVersion(),
-	}
-
-	return payload, nil
-}
-
-// Collects global metrics from all collectors
-// If a collector fails, it will log a warning and continue to the next collector (don't fail the whole function)
-func (w *MetricCollectionWorker) collectGlobalMetrics(ctx context.Context, from time.Time, to time.Time) ([]models.MetricData, error) {
-	metrics := make([]models.MetricData, 0)
-	logger := utils.LoggerFromContext(ctx)
-
-	for _, collector := range w.collectors.GetGlobalCollectors() {
-		value, err := collector.Collect(ctx, from, to)
-		if err != nil {
-			logger.WarnContext(ctx, "Failed to collect global metrics", "error", err)
-			continue
-		}
-		metrics = slices.Concat(metrics, value)
-	}
-
-	return metrics, nil
-}
-
-// Collects organization metrics from all collectors, fetching all organizations from the database first
-// If a collector fails, it will log a warning and continue to the next collector (don't fail the whole function)
-func (w *MetricCollectionWorker) collectOrganizationMetrics(ctx context.Context, from time.Time, to time.Time) ([]models.MetricData, error) {
-	metrics := make([]models.MetricData, 0)
-	logger := utils.LoggerFromContext(ctx)
-
-	orgs, err := w.getListOfOrganizations(ctx)
-	if err != nil {
-		return []models.MetricData{}, err
-	}
-
-	for _, org := range orgs {
-		for _, collector := range w.collectors.GetCollectors() {
-			value, err := collector.Collect(ctx, org.Id, from, to)
-			if err != nil {
-				logger.WarnContext(ctx, "Failed to collect organization metrics", "error", err)
-				continue
-			}
-			metrics = slices.Concat(metrics, value)
-		}
-	}
-
-	return metrics, nil
-}
-
-// Fetches all organizations from the database
-// NOTE: Add caching to avoid fetching the same organizations every time (but how can we invalidate the cache?)
-func (w *MetricCollectionWorker) getListOfOrganizations(ctx context.Context) ([]models.Organization, error) {
-	orgs, err := w.organizationRepository.AllOrganizations(ctx, w.executorFactory.NewExecutor())
-	if err != nil {
-		return []models.Organization{}, err
-	}
-	return orgs, nil
-}

usecases/usecases.go

@@ -346,8 +346,9 @@ func (usecases *Usecases) NewTaskQueueWorker(riverClient *river.Client[pgx.Tx])
 
 func (usecases *Usecases) NewMetricsCollectionWorker() scheduled_execution.MetricCollectionWorker {
 	return scheduled_execution.NewMetricCollectionWorker(
-		usecases.NewExecutorFactory(),
-		&usecases.Repositories.MarbleDbRepository,
-		metrics_collection.NewCollectorsV1(),
+		metrics_collection.NewCollectorsV1(
+			usecases.NewExecutorFactory(),
+			&usecases.Repositories.MarbleDbRepository,
+		),
 	)
 }