Performance improvement for GET /projects/params endpoint. (#1093)

Performance improvement for GET /projects/params endpoint.

pkg/api/aim2/dao/repositories/metric.go

@@ -29,8 +29,8 @@ type SearchResultMap = map[string]SearchResult
 // MetricRepositoryProvider provides an interface to work with models.Metric entity.
 type MetricRepositoryProvider interface {
 	BaseRepositoryProvider
-	// GetLatestMetricsByExperiments returns latest metrics by provided experiments.
-	GetLatestMetricsByExperiments(
+	// GetMetricKeysAndContextsByExperiments returns metric keys and contexts by provided experiments.
+	GetMetricKeysAndContextsByExperiments(
 		ctx context.Context, namespaceID uint, experiments []int,
 	) ([]models.LatestMetric, error)
 	// SearchMetrics returns a sql.Rows cursor for streaming the metrics matching the request.
@@ -57,18 +57,20 @@ func NewMetricRepository(db *gorm.DB) *MetricRepository {
 	}
 }
 
-// GetLatestMetricsByExperiments returns latest metrics by provided experiments.
-func (r MetricRepository) GetLatestMetricsByExperiments(
+// GetMetricKeysAndContextsByExperiments returns metric keys and contexts by provided experiments.
+func (r MetricRepository) GetMetricKeysAndContextsByExperiments(
 	ctx context.Context, namespaceID uint, experiments []int,
 ) ([]models.LatestMetric, error) {
-	query := r.db.WithContext(ctx).Distinct().Model(
+	query := r.db.WithContext(ctx).Distinct().Select(
+		"key", "context_id",
+	).Model(
 		&models.LatestMetric{},
 	).Joins(
 		"JOIN runs USING(run_uuid)",
 	).Joins(
 		"INNER JOIN experiments ON experiments.experiment_id = runs.experiment_id AND experiments.namespace_id = ?",
 		namespaceID,
-	).Joins(
+	).Preload(
 		"Context",
 	).Where(
 		"runs.lifecycle_stage = ?", models.LifecycleStageActive,

pkg/api/aim2/services/project/service.go

@@ -104,7 +104,7 @@ func (s Service) GetProjectParams(
 	}
 
 	if slices.Contains(req.Sequences, "metric") {
-		metrics, err := s.metricRepository.GetLatestMetricsByExperiments(ctx, namespaceID, req.Experiments)
+		metrics, err := s.metricRepository.GetMetricKeysAndContextsByExperiments(ctx, namespaceID, req.Experiments)
 		if err != nil {
 			return nil, api.NewInternalError("error getting metrics: %s", err)
 		}