Add startswith and endswith to query parser for context (#848)

* Add startswith and endswith to query parser for context
G-Research · Feb 5, 2024 · 552905b · 552905b
1 parent 19f6420
commit 552905b
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 26 deletions.
diff --git a/pkg/api/aim/query/query.go b/pkg/api/aim/query/query.go
@@ -248,46 +248,58 @@ func (pq *parsedQuery) parseAttribute(node *ast.Attribute) (any, error) {
 		case "endswith":
 			return callable(func(args []ast.Expr) (any, error) {
 				if len(args) != 1 {
-					return nil, errors.New("`endwith` function support exactly one argument")
+					return nil, errors.New("`endswith` function support exactly one argument")
 				}
-				c, ok := parsedNode.(clause.Column)
-				if !ok {
-					return nil, errors.New("unsupported node type. has to be clause.Column")
-				}
-
 				arg, ok := args[0].(*ast.Str)
 				if !ok {
 					return nil, errors.New("unsupported argument type. has to be `string` only")
 				}
-				return clause.Like{
-					Value: fmt.Sprintf("%%%s", arg.S),
-					Column: clause.Column{
-						Table: c.Table,
-						Name:  c.Name,
-					},
-				}, nil
+				value := fmt.Sprintf("%%%s", arg.S)
+				switch c := parsedNode.(type) {
+				case clause.Column:
+					return clause.Like{
+						Value: value,
+						Column: clause.Column{
+							Table: c.Table,
+							Name:  c.Name,
+						},
+					}, nil
+				case Json:
+					return JsonLike{
+						Value: value,
+						Json:  c,
+					}, nil
+				default:
+					return nil, errors.New("unsupported node type. has to be clause.Column or Json")
+				}
 			}), nil
 		case "startswith":
 			return callable(func(args []ast.Expr) (any, error) {
 				if len(args) != 1 {
-					return nil, errors.New("`startwith` function support exactly one argument")
+					return nil, errors.New("`startswith` function support exactly one argument")
 				}
-				c, ok := parsedNode.(clause.Column)
-				if !ok {
-					return nil, errors.New("unsupported node type. has to be clause.Column")
-				}
-
 				arg, ok := args[0].(*ast.Str)
 				if !ok {
 					return nil, errors.New("unsupported argument type. has to be `string` only")
 				}
-				return clause.Like{
-					Value: fmt.Sprintf("%s%%", arg.S),
-					Column: clause.Column{
-						Table: c.Table,
-						Name:  c.Name,
-					},
-				}, nil
+				value := fmt.Sprintf("%s%%", arg.S)
+				switch c := parsedNode.(type) {
+				case clause.Column:
+					return clause.Like{
+						Value: value,
+						Column: clause.Column{
+							Table: c.Table,
+							Name:  c.Name,
+						},
+					}, nil
+				case Json:
+					return JsonLike{
+						Value: value,
+						Json:  c,
+					}, nil
+				default:
+					return nil, errors.New("unsupported node type. has to be clause.Column or Json")
+				}
 			}), nil
 		}
 

diff --git a/pkg/api/aim/query/query_test.go b/pkg/api/aim/query/query_test.go
@@ -209,6 +209,24 @@ func (s *QueryTestSuite) TestPostgresDialector_Ok() {
 				`AND "runs"."lifecycle_stage" <> $3`,
 			expectedVars: []interface{}{"{key1}", "%val1%", models.LifecycleStageDeleted},
 		},
+		{
+			name:          "TestMetricContextStartsWith",
+			query:         `metric.context.key1.startswith("va")`,
+			selectMetrics: true,
+			expectedSQL: `SELECT ID FROM "metrics" ` +
+				`WHERE "contexts"."json"#>>$1 LIKE $2 ` +
+				`AND "runs"."lifecycle_stage" <> $3`,
+			expectedVars: []interface{}{"{key1}", "va%", models.LifecycleStageDeleted},
+		},
+		{
+			name:          "TestMetricContextEndsWith",
+			query:         `metric.context.key1.endswith("va")`,
+			selectMetrics: true,
+			expectedSQL: `SELECT ID FROM "metrics" ` +
+				`WHERE "contexts"."json"#>>$1 LIKE $2 ` +
+				`AND "runs"."lifecycle_stage" <> $3`,
+			expectedVars: []interface{}{"{key1}", "%va", models.LifecycleStageDeleted},
+		},
 	}
 
 	for _, tt := range tests {
@@ -430,6 +448,24 @@ func (s *QueryTestSuite) TestSqliteDialector_Ok() {
 				`AND "runs"."lifecycle_stage" <> $3`,
 			expectedVars: []interface{}{"$.key1", "%val1%", models.LifecycleStageDeleted},
 		},
+		{
+			name:          "TestMetricContextStartsWith",
+			query:         `metric.context.key1.startswith("va")`,
+			selectMetrics: true,
+			expectedSQL: `SELECT ID FROM "metrics" ` +
+				`WHERE IFNULL("contexts"."json", JSON('{}'))->>$1 LIKE $2 ` +
+				`AND "runs"."lifecycle_stage" <> $3`,
+			expectedVars: []interface{}{"$.key1", "va%", models.LifecycleStageDeleted},
+		},
+		{
+			name:          "TestMetricContextEndsWith",
+			query:         `metric.context.key1.endswith("va")`,
+			selectMetrics: true,
+			expectedSQL: `SELECT ID FROM "metrics" ` +
+				`WHERE IFNULL("contexts"."json", JSON('{}'))->>$1 LIKE $2 ` +
+				`AND "runs"."lifecycle_stage" <> $3`,
+			expectedVars: []interface{}{"$.key1", "%va", models.LifecycleStageDeleted},
+		},
 	}
 
 	for _, tt := range tests {