Remove metric querying support

pkg/api/aim2/api/request/run.go

@@ -49,13 +49,20 @@ type SearchRunsRequest struct {
 	ExcludeTraces bool   `query:"exclude_traces"`
 }
 
+// MetricTuple represents a metric with key and context.
+type MetricTuple struct {
+	Key     string `query:"metric"`
+	Context string `query:"context"`
+}
+
 // SearchMetricsRequest is a request struct for `GET /runs/search/metric` endpoint.
 type SearchMetricsRequest struct {
 	BaseSearchRequest
-	Query      string `query:"q"`
-	Steps      int    `query:"p"`
-	XAxis      string `query:"x_axis"`
-	SkipSystem bool   `query:"skip_system"`
+	Metrics    []MetricTuple `query:"m"`
+	Query      string        `query:"q"`
+	Steps      int           `query:"p"`
+	XAxis      string        `query:"x_axis"`
+	SkipSystem bool          `query:"skip_system"`
 }
 
 // SearchAlignedMetricsRequest is a request struct for `GET /runs/search/metric/align` endpoint.

pkg/api/aim2/services/run/helpers.go → pkg/api/aim2/common/helpers.go

@@ -1,4 +1,4 @@
-package run
+package common
 
 import (
 	"encoding/json"

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

@@ -4,12 +4,14 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"strings"
 
 	"github.com/gofiber/fiber/v2"
 	"github.com/rotisserie/eris"
 	"gorm.io/gorm"
 
 	"github.com/G-Research/fasttrackml/pkg/api/aim2/api/request"
+	"github.com/G-Research/fasttrackml/pkg/api/aim2/common"
 	"github.com/G-Research/fasttrackml/pkg/api/aim2/dao/models"
 	"github.com/G-Research/fasttrackml/pkg/api/aim2/query"
 	"github.com/G-Research/fasttrackml/pkg/common/db/types"
@@ -103,7 +105,7 @@ func (r MetricRepository) SearchMetrics(
 		return nil, 0, nil, err
 	}
 
-	if !pq.IsMetricSelected() {
+	if req.Metrics == nil || len(req.Metrics) == 0 {
 		return nil, 0, nil, eris.New("No metrics are selected")
 	}
 
@@ -169,30 +171,45 @@ func (r MetricRepository) SearchMetrics(
 		result[r.ID] = SearchResult{int64(r.RowNum), run}
 	}
 
+	ids, err := r.findContextIDs(ctx, &req)
+	if err != nil {
+		return nil, 0, nil, eris.Wrap(err, "error finding context ids")
+	}
+
+	var metricKeyContextConditionSlice []string
+	for i, tuple := range req.Metrics {
+		condition := fmt.Sprintf("(latest_metrics.key = '%s' AND contexts.id = %d)", tuple.Key, ids[i])
+		metricKeyContextConditionSlice = append(metricKeyContextConditionSlice, condition)
+	}
+	metricKeyContextCondition := strings.Join(metricKeyContextConditionSlice, " OR ")
+
+	subQuery := r.db.WithContext(ctx).
+		Select(
+			"runs.run_uuid",
+			"runs.row_num",
+			"latest_metrics.key",
+			"latest_metrics.context_id",
+			"contexts.json AS context_json",
+			fmt.Sprintf("(latest_metrics.last_iter + 1)/ %f AS interval", float32(req.Steps)),
+		).
+		Table("runs").
+		Joins(
+			"INNER JOIN experiments ON experiments.experiment_id = runs.experiment_id AND experiments.namespace_id = ?",
+			namespaceID,
+		).
+		Joins("LEFT JOIN latest_metrics USING(run_uuid)").
+		Joins("LEFT JOIN contexts ON latest_metrics.context_id = contexts.id").
+		Where(metricKeyContextCondition)
+
 	tx := r.db.WithContext(ctx).
 		Select(`
-			metrics.*,
-			runmetrics.context_json`,
+        metrics.*,
+        runmetrics.context_json`,
 		).
 		Table("metrics").
 		Joins(
 			"INNER JOIN (?) runmetrics USING(run_uuid, key, context_id)",
-			pq.Filter(r.db.WithContext(ctx).
-				Select(
-					"runs.run_uuid",
-					"runs.row_num",
-					"latest_metrics.key",
-					"latest_metrics.context_id",
-					"contexts.json AS context_json",
-					fmt.Sprintf("(latest_metrics.last_iter + 1)/ %f AS interval", float32(req.Steps)),
-				).
-				Table("runs").
-				Joins(
-					"INNER JOIN experiments ON experiments.experiment_id = runs.experiment_id AND experiments.namespace_id = ?",
-					namespaceID,
-				).
-				Joins("LEFT JOIN latest_metrics USING(run_uuid)").
-				Joins("LEFT JOIN contexts ON latest_metrics.context_id = contexts.id")),
+			pq.Filter(subQuery),
 		).
 		Where("MOD(metrics.iter + 1 + runmetrics.interval / 2, runmetrics.interval) < 1").
 		Order("runmetrics.row_num DESC").
@@ -234,3 +251,29 @@ func (r MetricRepository) GetContextListByContextObjects(
 	}
 	return contexts, nil
 }
+
+func (r MetricRepository) findContextIDs(ctx context.Context, req *request.SearchMetricsRequest) ([]uint, error) {
+	contextList := []types.JSONB{}
+	contextsMap := map[string]types.JSONB{}
+	for _, r := range req.Metrics {
+		data := types.JSONB(r.Context)
+		contextList = append(contextList, data)
+		contextsMap[string(data)] = data
+	}
+
+	contexts, err := r.GetContextListByContextObjects(ctx, contextsMap)
+	if err != nil {
+		return nil, fmt.Errorf("error getting context list: %w", err)
+	}
+
+	ids := make([]uint, len(contextList))
+	for _, context := range contexts {
+		for i := 0; i < len(contextList); i++ {
+			if common.CompareJson(contextList[i], context.Json) {
+				ids[i] = context.ID
+			}
+		}
+	}
+
+	return ids, nil
+}

pkg/api/aim2/query/query.go

@@ -35,7 +35,6 @@ type QueryParser struct {
 
 type ParsedQuery interface {
 	Filter(*gorm.DB) *gorm.DB
-	IsMetricSelected() bool
 }
 
 type parsedQuery struct {
@@ -187,10 +186,6 @@ func (pq *parsedQuery) Filter(tx *gorm.DB) *gorm.DB {
 	return tx
 }
 
-func (pq *parsedQuery) IsMetricSelected() bool {
-	return pq.metricSelected
-}
-
 func (pq *parsedQuery) parseNode(node ast.Expr) (any, error) {
 	ret, err := pq._parseNode(node)
 	if err != nil && !errors.Is(err, SyntaxError{}) {
@@ -644,51 +639,6 @@ func (pq *parsedQuery) parseName(node *ast.Name) (any, error) {
 					}
 				},
 			), nil
-		case "metric":
-			table, ok := pq.qp.Tables["metrics"]
-			if !ok {
-				return nil, errors.New("unsupported name identifier 'metric'")
-			}
-			return attributeGetter(
-				func(attr string) (any, error) {
-					switch attr {
-					case "name":
-						pq.metricSelected = true
-						return clause.Column{
-							Table: table,
-							Name:  "key",
-						}, nil
-					case "last":
-						return clause.Column{
-							Table: table,
-							Name:  "value",
-						}, nil
-					case "last_step":
-						return clause.Column{
-							Table: table,
-							Name:  "last_iter",
-						}, nil
-					case "first_step":
-						return 0, nil
-					case "context":
-						return attributeGetter(
-							func(contextKey string) (any, error) {
-								// Add a WHERE clause for the context key
-								return Json{
-									Column: clause.Column{
-										Table: TableContexts,
-										Name:  "json",
-									},
-									JsonPath:  contextKey,
-									Dialector: pq.qp.Dialector,
-								}, nil
-							},
-						), nil
-					default:
-						return nil, fmt.Errorf("unsupported metrics attribute %q", attr)
-					}
-				},
-			), nil
 		case "re":
 			return attributeGetter(
 				func(attr string) (any, error) {