Enable metric.context filter in Aim metrics search

pkg/api/aim/query/clause.go

@@ -1,6 +1,10 @@
 package query
 
 import (
+	"database/sql/driver"
+	"reflect"
+	"strings"
+
 	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm/clause"
@@ -52,3 +56,146 @@ func (regexp Regexp) writeColumn(builder clause.Builder) {
 		builder.WriteQuoted(regexp.Column)
 	}
 }
+
+// Json clause for string match at a json path.
+type Json struct {
+	clause.Column
+	JsonPath  string
+	Dialector string
+}
+
+// Build builds positive statement.
+func (json Json) Build(builder clause.Builder) {
+	json.writeColumn(builder)
+	jsonPath := json.JsonPath
+	switch json.Dialector {
+	case postgres.Dialector{}.Name():
+		//nolint:errcheck,gosec
+		builder.WriteString("#>>")
+		jsonPath = "{" + strings.ReplaceAll(jsonPath, ",", ".") + "}"
+	default:
+		//nolint:errcheck,gosec
+		builder.WriteString("->>")
+	}
+	builder.AddVar(builder, jsonPath)
+}
+
+// NegationBuild builds negative statement.
+func (json Json) NegationBuild(builder clause.Builder) {
+	json.writeColumn(builder)
+	jsonPath := json.JsonPath
+	switch json.Dialector {
+	case postgres.Dialector{}.Name():
+		//nolint:errcheck,gosec
+		builder.WriteString("#>>")
+		jsonPath = "{" + strings.ReplaceAll(jsonPath, ",", ".") + "}"
+	default:
+		//nolint:errcheck,gosec
+		builder.WriteString("->>")
+	}
+	builder.AddVar(builder, jsonPath)
+}
+
+func (json Json) writeColumn(builder clause.Builder) {
+	switch json.Dialector {
+	case sqlite.Dialector{}.Name():
+		//nolint:errcheck,gosec
+		builder.WriteString("IFNULL(")
+		builder.WriteQuoted(json.Column)
+		//nolint:errcheck,gosec
+		builder.WriteString(", JSON('{}'))")
+	default:
+		builder.WriteQuoted(json.Column)
+	}
+}
+
+type JsonEq struct {
+	Left  Json
+	Value any
+}
+
+func (eq JsonEq) Build(builder clause.Builder) {
+	eq.Left.Build(builder)
+	switch eq.Value.(type) {
+	case []string, []int, []int32, []int64, []uint, []uint32, []uint64, []interface{}:
+		rv := reflect.ValueOf(eq.Value)
+		if rv.Len() == 0 {
+			//nolint:errcheck,gosec
+			builder.WriteString(" IN (NULL)")
+		} else {
+			//nolint:errcheck,gosec
+			builder.WriteString(" IN (")
+			for i := 0; i < rv.Len(); i++ {
+				if i > 0 {
+					//nolint:errcheck,gosec
+					builder.WriteByte(',')
+				}
+				builder.AddVar(builder, rv.Index(i).Interface())
+			}
+			//nolint:errcheck,gosec
+			builder.WriteByte(')')
+		}
+	default:
+		if eqNil(eq.Value) {
+			//nolint:errcheck,gosec
+			builder.WriteString(" IS NULL")
+		} else {
+			//nolint:errcheck,gosec
+			builder.WriteString(" = ")
+			builder.AddVar(builder, eq.Value)
+		}
+	}
+}
+
+func (eq JsonEq) NegationBuild(builder clause.Builder) {
+	JsonNeq(eq).Build(builder)
+}
+
+// JsonNeq not equal to for where
+type JsonNeq JsonEq
+
+func (neq JsonNeq) Build(builder clause.Builder) {
+	neq.Left.Build(builder)
+	switch neq.Value.(type) {
+	case []string, []int, []int32, []int64, []uint, []uint32, []uint64, []interface{}:
+		//nolint:errcheck,gosec
+		builder.WriteString(" NOT IN (")
+		rv := reflect.ValueOf(neq.Value)
+		for i := 0; i < rv.Len(); i++ {
+			if i > 0 {
+				//nolint:errcheck,gosec
+				builder.WriteByte(',')
+			}
+			builder.AddVar(builder, rv.Index(i).Interface())
+		}
+		//nolint:errcheck,gosec
+		builder.WriteByte(')')
+	default:
+		if eqNil(neq.Value) {
+			//nolint:errcheck,gosec
+			builder.WriteString(" IS NOT NULL")
+		} else {
+			//nolint:errcheck,gosec
+			builder.WriteString(" <> ")
+			builder.AddVar(builder, neq.Value)
+		}
+	}
+}
+
+func (neq JsonNeq) NegationBuild(builder clause.Builder) {
+	JsonEq(neq).Build(builder)
+}
+
+func eqNil(value interface{}) bool {
+	if valuer, ok := value.(driver.Valuer); ok && !eqNilReflect(valuer) {
+		//nolint:errcheck
+		value, _ = valuer.Value()
+	}
+
+	return value == nil || eqNilReflect(value)
+}
+
+func eqNilReflect(value interface{}) bool {
+	reflectValue := reflect.ValueOf(value)
+	return reflectValue.Kind() == reflect.Ptr && reflectValue.IsNil()
+}

pkg/api/aim/query/query.go

@@ -312,6 +312,7 @@ func (pq *parsedQuery) parseCall(node *ast.Call) (any, error) {
 	}
 }
 
+// nolint:gocyclo
 func (pq *parsedQuery) parseCompare(node *ast.Compare) (any, error) {
 	exprs := make([]clause.Expression, len(node.Ops))
 
@@ -345,6 +346,11 @@ func (pq *parsedQuery) parseCompare(node *ast.Compare) (any, error) {
 			default:
 				return nil, fmt.Errorf("unsupported comparison %q", ast.Dump(node))
 			}
+		case Json:
+			exprs[i], err = newSqlJsonPathComparison(op, left, right)
+			if err != nil {
+				return nil, err
+			}
 		default:
 			switch right := right.(type) {
 			case clause.Column:
@@ -602,6 +608,31 @@ func (pq *parsedQuery) parseName(node *ast.Name) (any, error) {
 						}, nil
 					case "first_step":
 						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 metrics.context_id = contexts.id",
+									}
+									pq.joins["metric_contexts"] = j
+								}
+
+								// Add a WHERE clause for the context key
+								return Json{
+									Column: clause.Column{
+										Table: "contexts",
+										Name:  "json",
+									},
+									JsonPath:  contextKey,
+									Dialector: pq.qp.Dialector,
+								}, nil
+							},
+						), nil
 					default:
 						return nil, fmt.Errorf("unsupported metrics attribute %q", attr)
 					}
@@ -840,6 +871,23 @@ func newSqlComparison(op ast.CmpOp, left clause.Column, right any) (clause.Expre
 	}
 }
 
+func newSqlJsonPathComparison(op ast.CmpOp, left Json, right any) (clause.Expression, error) {
+	switch op {
+	case ast.Eq:
+		return JsonEq{
+			Left:  left,
+			Value: right,
+		}, nil
+	case ast.NotEq:
+		return JsonNeq{
+			Left:  left,
+			Value: right,
+		}, nil
+	default:
+		return nil, fmt.Errorf("unsupported comparison operation %q", op)
+	}
+}
+
 func reverseComparison(op ast.CmpOp, left any, right clause.Column) (ast.CmpOp, clause.Column, any, error) {
 	switch op {
 	case ast.Lt: