Context query subscript

pkg/api/aim/query/query.go

@@ -16,6 +16,11 @@ import (
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/models"
 )
 
+const (
+	TableContexts      = "contexts"
+	TableLatestMetrics = "latest_metrics"
+)
+
 type DefaultExpression struct {
 	Contains   string
 	Expression string
@@ -154,6 +159,7 @@ func (qp *QueryParser) Parse(q string) (ParsedQuery, error) {
 	return pq, nil
 }
 
+// Filter will add the appropriate Joins and Where clauses to the tx.
 func (pq *parsedQuery) Filter(tx *gorm.DB) *gorm.DB {
 	for _, j := range pq.joins {
 		tx.Joins(j.query, j.args...)
@@ -204,6 +210,10 @@ func (pq *parsedQuery) _parseNode(node ast.Expr) (any, error) {
 		return pq.parseAttribute(n)
 	case *ast.Compare:
 		return pq.parseCompare(n)
+	case *ast.Dict:
+		return pq.parseDictionary(n)
+	case *ast.Tuple:
+		return pq.parseTuple(n)
 	default:
 		return nil, fmt.Errorf("unsupported expression %q", ast.Dump(n))
 	}
@@ -404,6 +414,38 @@ func (pq *parsedQuery) parseCompare(node *ast.Compare) (any, error) {
 	}, nil
 }
 
+// parseDictionary will return `clause.And` having joined JsonEq conditions
+// derived from the dictionary
+func (pq *parsedQuery) parseDictionary(node *ast.Dict) (any, error) {
+	clauses := make([]clause.Expression, len(node.Keys))
+	for i, key := range node.Keys {
+		clauses[i] = JsonEq{
+			Left: Json{
+				Column: clause.Column{
+					Table: TableContexts,
+					Name:  "json",
+				},
+				JsonPath:  string(key.(*ast.Str).S),
+				Dialector: pq.qp.Dialector,
+			},
+			Value: string(node.Values[i].(*ast.Str).S),
+		}
+	}
+	return clause.And(clauses...), nil
+}
+
+func (pq *parsedQuery) parseTuple(node *ast.Tuple) (any, error) {
+	var err error
+	list := make([]any, len(node.Elts))
+	for i, e := range node.Elts {
+		list[i], err = pq.parseNode(e)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return list, nil
+}
+
 func (pq *parsedQuery) parseList(node *ast.List) (any, error) {
 	var err error
 	list := make([]any, len(node.Elts))
@@ -488,41 +530,7 @@ func (pq *parsedQuery) parseName(node *ast.Name) (any, error) {
 								if err != nil {
 									return nil, err
 								}
-								switch v := v.(type) {
-								case string:
-									j, ok := pq.joins[fmt.Sprintf("metrics:%s", v)]
-									if !ok {
-										alias := fmt.Sprintf("metrics_%d", len(pq.joins))
-										j = join{
-											alias: alias,
-											query: fmt.Sprintf(
-												"LEFT JOIN latest_metrics %s ON %s.run_uuid = %s.run_uuid AND %s.key = ?",
-												alias, table, alias, alias,
-											),
-											args: []any{v},
-										}
-										pq.joins[fmt.Sprintf("metrics:%s", v)] = j
-									}
-									return attributeGetter(func(attr string) (any, error) {
-										var name string
-										switch attr {
-										case "last":
-											name = "value"
-										case "last_step":
-											name = "last_iter"
-										case "first_step":
-											return 0, nil
-										default:
-											return nil, fmt.Errorf("unsupported metrics attribute %q", attr)
-										}
-										return clause.Column{
-											Table: j.alias,
-											Name:  name,
-										}, nil
-									}), nil
-								default:
-									return nil, fmt.Errorf("unsupported index value type %t", v)
-								}
+								return pq.metricSubscriptSlicer(v)
 							default:
 								return nil, fmt.Errorf("unsupported slicer %q", ast.Dump(s))
 							}
@@ -555,7 +563,7 @@ func (pq *parsedQuery) parseName(node *ast.Name) (any, error) {
 										Name:  "value",
 									}, nil
 								default:
-									return nil, fmt.Errorf("unsupported index value type %t", v)
+									return nil, fmt.Errorf("unsupported index value type %T", v)
 								}
 							default:
 								return nil, fmt.Errorf("unsupported slicer %q", ast.Dump(s))
@@ -610,22 +618,12 @@ func (pq *parsedQuery) parseName(node *ast.Name) (any, error) {
 						return 0, nil
 					case "context":
 						return attributeGetter(
-							func(contextKey string) (any, error) {
-								// create the join for contexts
-								_, ok := pq.joins["metric_contexts"]
-								if !ok {
-									alias := "contexts"
-									j := join{
-										alias: alias,
-										query: "LEFT JOIN contexts ON latest_metrics.context_id = contexts.id",
-									}
-									pq.joins["metric_contexts"] = j
-								}
 
+							func(contextKey string) (any, error) {
 								// Add a WHERE clause for the context key
 								return Json{
 									Column: clause.Column{
-										Table: "contexts",
+										Table: TableContexts,
 										Name:  "json",
 									},
 									JsonPath:  contextKey,
@@ -728,6 +726,65 @@ func (pq *parsedQuery) parseName(node *ast.Name) (any, error) {
 	}
 }
 
+func (pq *parsedQuery) metricSubscriptSlicer(v any) (any, error) {
+	switch v := v.(type) {
+	case string:
+		// case of metric key
+		keyEqual := pq.metricSubscriptStringExpression(v)
+		pq.conditions = append(pq.conditions, keyEqual)
+		return metricAttributeGetter(TableLatestMetrics)
+	case clause.Expression:
+		// case of metric context dictionary
+		pq.conditions = append(pq.conditions, v)
+		return metricAttributeGetter(TableLatestMetrics)
+	case []any:
+		// case of subscript tuple (string and context dictionary)
+		if len(v) != 2 {
+			return nil, fmt.Errorf("unsupported tuple length %d (should be 2)", len(v))
+		}
+		metricKey, ok := v[0].(string)
+		if !ok {
+			return nil, fmt.Errorf("unsupported tuple value type %T (should be string at 0)", v)
+		}
+		metricContextExpression, ok := v[1].(clause.Expression)
+		if !ok {
+			return nil, fmt.Errorf("unsupported index value type %T (should be clause.Expression at 1)", v)
+		}
+		metricKeyExpression := pq.metricSubscriptStringExpression(metricKey)
+		pq.conditions = append(pq.conditions, clause.And(metricKeyExpression, metricContextExpression))
+		return metricAttributeGetter(TableLatestMetrics)
+	default:
+		return nil, fmt.Errorf("unsupported index value type %T", v)
+	}
+}
+
+func (pq *parsedQuery) metricSubscriptStringExpression(v string) clause.Expression {
+	return clause.Eq{
+		Column: fmt.Sprintf("%s.%s", TableLatestMetrics, "key"),
+		Value:  v,
+	}
+}
+
+func metricAttributeGetter(table string) (any, error) {
+	return attributeGetter(func(attr string) (any, error) {
+		var name string
+		switch attr {
+		case "last":
+			name = "value"
+		case "last_step":
+			name = "last_iter"
+		case "first_step":
+			return 0, nil
+		default:
+			return nil, fmt.Errorf("unsupported metrics attribute %q", attr)
+		}
+		return clause.Column{
+			Table: table,
+			Name:  name,
+		}, nil
+	}), nil
+}
+
 func (pq *parsedQuery) parseNameConstant(node *ast.NameConstant) (any, error) {
 	switch node.Value.Type() {
 	case py.NoneTypeType:

pkg/api/aim/query/query_test.go

@@ -110,24 +110,21 @@ func (s *QueryTestSuite) TestPostgresDialector_Ok() {
 			name:  "TestNegativeInteger",
 			query: `run.metrics['my_metric'].last < -1`,
 			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
-				`LEFT JOIN latest_metrics metrics_0 ON runs.run_uuid = metrics_0.run_uuid AND metrics_0.key = $1 ` +
-				`WHERE ("metrics_0"."value" < $2 AND "runs"."lifecycle_stage" <> $3)`,
+				`WHERE ("latest_metrics"."key" = $1 AND ("latest_metrics"."value" < $2 AND "runs"."lifecycle_stage" <> $3))`,
 			expectedVars: []interface{}{"my_metric", -1, models.LifecycleStageDeleted},
 		},
 		{
 			name:  "TestNegativeFloat",
 			query: `run.metrics['my_metric'].last < -1.0`,
 			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
-				`LEFT JOIN latest_metrics metrics_0 ON runs.run_uuid = metrics_0.run_uuid AND metrics_0.key = $1 ` +
-				`WHERE ("metrics_0"."value" < $2 AND "runs"."lifecycle_stage" <> $3)`,
+				`WHERE ("latest_metrics"."key" = $1 AND ("latest_metrics"."value" < $2 AND "runs"."lifecycle_stage" <> $3))`,
 			expectedVars: []interface{}{"my_metric", -1.0, models.LifecycleStageDeleted},
 		},
 		{
 			name:          "TestMetricContext",
 			query:         `metric.context.key1 == 'value1'`,
 			selectMetrics: true,
 			expectedSQL: `SELECT ID FROM "metrics" ` +
-				`LEFT JOIN contexts ON latest_metrics.context_id = contexts.id ` +
 				`WHERE ("contexts"."json"#>>$1 = $2 AND "runs"."lifecycle_stage" <> $3)`,
 			expectedVars: []interface{}{"{key1}", "value1", models.LifecycleStageDeleted},
 		},
@@ -136,10 +133,25 @@ func (s *QueryTestSuite) TestPostgresDialector_Ok() {
 			query:         `metric.context.key1 != 'value1'`,
 			selectMetrics: true,
 			expectedSQL: `SELECT ID FROM "metrics" ` +
-				`LEFT JOIN contexts ON latest_metrics.context_id = contexts.id ` +
 				`WHERE ("contexts"."json"#>>$1 <> $2 AND "runs"."lifecycle_stage" <> $3)`,
 			expectedVars: []interface{}{"{key1}", "value1", models.LifecycleStageDeleted},
 		},
+		{
+			name:  "TestMetricContextSlice",
+			query: `run.metrics[{"key1": "value1"}].last < -1`,
+			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
+				`WHERE ("contexts"."json"#>>$1 = $2 AND ("latest_metrics"."value" < $3 ` +
+				`AND "runs"."lifecycle_stage" <> $4))`,
+			expectedVars: []interface{}{"{key1}", "value1", -1, models.LifecycleStageDeleted},
+		},
+		{
+			name:  "TestMetricContextSliceTuple",
+			query: `run.metrics["my_metric", {"key1": "value1"}].last < -1`,
+			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
+				`WHERE (("latest_metrics"."key" = $1 AND "contexts"."json"#>>$2 = $3) ` +
+				`AND ("latest_metrics"."value" < $4 AND "runs"."lifecycle_stage" <> $5))`,
+			expectedVars: []interface{}{"my_metric", "{key1}", "value1", -1, models.LifecycleStageDeleted},
+		},
 	}
 
 	for _, tt := range tests {
@@ -251,24 +263,23 @@ func (s *QueryTestSuite) TestSqliteDialector_Ok() {
 			name:  "TestNegativeInteger",
 			query: `run.metrics['my_metric'].last < -1`,
 			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
-				`LEFT JOIN latest_metrics metrics_0 ON runs.run_uuid = metrics_0.run_uuid AND metrics_0.key = $1 ` +
-				`WHERE ("metrics_0"."value" < $2 AND "runs"."lifecycle_stage" <> $3)`,
+				`WHERE ("latest_metrics"."key" = $1 AND ("latest_metrics"."value" < $2 ` +
+				`AND "runs"."lifecycle_stage" <> $3))`,
 			expectedVars: []interface{}{"my_metric", -1, models.LifecycleStageDeleted},
 		},
 		{
 			name:  "TestNegativeFloat",
 			query: `run.metrics['my_metric'].last < -1.0`,
 			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
-				`LEFT JOIN latest_metrics metrics_0 ON runs.run_uuid = metrics_0.run_uuid AND metrics_0.key = $1 ` +
-				`WHERE ("metrics_0"."value" < $2 AND "runs"."lifecycle_stage" <> $3)`,
+				`WHERE ("latest_metrics"."key" = $1 AND ("latest_metrics"."value" < $2 ` +
+				`AND "runs"."lifecycle_stage" <> $3))`,
 			expectedVars: []interface{}{"my_metric", -1.0, models.LifecycleStageDeleted},
 		},
 		{
 			name:          "TestMetricContext",
 			query:         `metric.context.key1 == 'value1'`,
 			selectMetrics: true,
 			expectedSQL: `SELECT ID FROM "metrics" ` +
-				`LEFT JOIN contexts ON latest_metrics.context_id = contexts.id ` +
 				`WHERE (IFNULL("contexts"."json", JSON('{}'))->>$1 = $2 AND "runs"."lifecycle_stage" <> $3)`,
 			expectedVars: []interface{}{"key1", "value1", models.LifecycleStageDeleted},
 		},
@@ -277,10 +288,33 @@ func (s *QueryTestSuite) TestSqliteDialector_Ok() {
 			query:         `metric.context.key1 != 'value1'`,
 			selectMetrics: true,
 			expectedSQL: `SELECT ID FROM "metrics" ` +
-				`LEFT JOIN contexts ON latest_metrics.context_id = contexts.id ` +
 				`WHERE (IFNULL("contexts"."json", JSON('{}'))->>$1 <> $2 AND "runs"."lifecycle_stage" <> $3)`,
 			expectedVars: []interface{}{"key1", "value1", models.LifecycleStageDeleted},
 		},
+		{
+			name:  "TestMetricKeySlice",
+			query: `run.metrics["key1"].last < -1`,
+			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
+				`WHERE ("latest_metrics"."key" = $1 AND ` +
+				`("latest_metrics"."value" < $2 AND "runs"."lifecycle_stage" <> $3))`,
+			expectedVars: []interface{}{"key1", -1, models.LifecycleStageDeleted},
+		},
+		{
+			name:  "TestMetricContextSlice",
+			query: `run.metrics[{"key1": "value1"}].last < -1`,
+			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
+				`WHERE (IFNULL("contexts"."json", JSON('{}'))->>$1 = $2 AND ` +
+				`("latest_metrics"."value" < $3 AND "runs"."lifecycle_stage" <> $4))`,
+			expectedVars: []interface{}{"key1", "value1", -1, models.LifecycleStageDeleted},
+		},
+		{
+			name:  "TestMetricContextSliceTuple",
+			query: `run.metrics["my_metric", {"key1": "value1"}].last < -1`,
+			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
+				`WHERE (("latest_metrics"."key" = $1 AND IFNULL("contexts"."json", JSON('{}'))->>$2 = $3) ` +
+				`AND ("latest_metrics"."value" < $4 AND "runs"."lifecycle_stage" <> $5))`,
+			expectedVars: []interface{}{"my_metric", "key1", "value1", -1, models.LifecycleStageDeleted},
+		},
 	}
 
 	for _, tt := range tests {
@@ -328,6 +362,11 @@ func (s *QueryTestSuite) Test_Error() {
 			query:         `metric.context.parent.nested == 'value1'`,
 			expectedError: SyntaxError{},
 		},
+		{
+			name:          "TestMetricContextSliceTupleOrder",
+			query:         `run.metrics[{"key1": "value1"}, "my_metric"].last < -1`,
+			expectedError: SyntaxError{},
+		},
 	}
 	for _, tt := range tests {
 		s.Run(tt.name, func() {

pkg/api/aim/runs.go

@@ -466,6 +466,9 @@ func SearchRuns(c *fiber.Ctx) error {
 	}
 
 	var runs []database.Run
+	// add joins needed by pq.Filter
+	tx.Joins("JOIN latest_metrics USING(run_uuid)").
+		Joins("JOIN contexts ON latest_metrics.context_id = contexts.id")
 	pq.Filter(tx).Find(&runs)
 	if tx.Error != nil {
 		return fmt.Errorf("error searching runs: %w", tx.Error)
@@ -660,7 +663,9 @@ func SearchMetrics(c *fiber.Ctx) error {
 				"INNER JOIN experiments ON experiments.experiment_id = runs.experiment_id AND experiments.namespace_id = ?",
 				ns.ID,
 			).
-			Joins("LEFT JOIN latest_metrics USING(run_uuid)"))).
+			Joins("JOIN latest_metrics USING(run_uuid)").
+			Joins("JOIN contexts ON latest_metrics.context_id = contexts.id"),
+		)).
 		Order("runs.row_num DESC").
 		Find(&runs); tx.Error != nil {
 		return fmt.Errorf("error searching run metrics: %w", tx.Error)
@@ -718,7 +723,7 @@ func SearchMetrics(c *fiber.Ctx) error {
 					"runs.row_num",
 					"latest_metrics.key",
 					"latest_metrics.context_id",
-					"latest_metrics_context.json AS context_json",
+					"contexts.json AS context_json",
 					fmt.Sprintf("(latest_metrics.last_iter + 1)/ %f AS interval", float32(q.Steps)),
 				).
 				Table("runs").
@@ -727,8 +732,7 @@ func SearchMetrics(c *fiber.Ctx) error {
 					ns.ID,
 				).
 				Joins("LEFT JOIN latest_metrics USING(run_uuid)").
-				Joins(`LEFT JOIN contexts latest_metrics_context `+
-					`ON latest_metrics.context_id = latest_metrics_context.id`)),
+				Joins("LEFT JOIN contexts ON latest_metrics.context_id = contexts.id")),
 		).
 		Where("MOD(metrics.iter + 1 + runmetrics.interval / 2, runmetrics.interval) < 1").
 		Order("runmetrics.row_num DESC").