Repository: G-Research/fasttrackml
Pull request #215 (2023-08-07 17:17:01)

Implement `IN` and `NOT IN` string comparison function in Aim query


Patch:
diff --git a/pkg/api/aim/query/query.go b/pkg/api/aim/query/query.go
index 607bc29..00677c8 100644
--- a/pkg/api/aim/query/query.go
+++ b/pkg/api/aim/query/query.go
@@ -38,10 +38,10 @@ type parsedQuery struct {
 	conditions []clause.Expression
 }
 
-type attributeGetter func(attr string) (any, error)
-
 type callable func(args []ast.Expr) (any, error)
 
+type attributeGetter func(attr string) (any, error)
+
 type subscriptSlicer func(index ast.Slicer) (any, error)
 
 type join struct {
@@ -234,28 +234,6 @@ func (pq *parsedQuery) parseAttribute(node *ast.Attribute) (any, error) {
 					},
 				}, nil
 			}), nil
-		case "contains":
-			return callable(func(args []ast.Expr) (any, error) {
-				if len(args) != 1 {
-					return nil, errors.New("`contains` 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
-			}), nil
 		case "startswith":
 			return callable(func(args []ast.Expr) (any, error) {
 				if len(args) != 1 {
@@ -364,13 +342,34 @@ func (pq *parsedQuery) parseCompare(node *ast.Compare) (any, error) {
 		default:
 			switch right := right.(type) {
 			case clause.Column:
-				o, l, r, err := reverseComparison(op, left, right)
-				if err != nil {
-					return nil, err
-				}
-				exprs[i], err = newSqlComparison(o, l, r)
-				if err != nil {
-					return nil, err
+				switch op {
+				case ast.In:
+					// for `IN` statement, left parameter has to be always `string`.
+					if _, ok := left.(string); !ok {
+						return nil, errors.New("left parameter has to be a string")
+					}
+					return clause.Like{
+						Value:  fmt.Sprintf("%%%s%%", left),
+						Column: right,
+					}, nil
+				case ast.NotIn:
+					// for `NOT IN` statement, left parameter has to be always `string`.
+					if _, ok := left.(string); !ok {
+						return nil, errors.New("left parameter has to be a string")
+					}
+					return negativeClause(clause.Like{
+						Value:  fmt.Sprintf("%%%s%%", left),
+						Column: right,
+					}), nil
+				default:
+					o, l, r, err := reverseComparison(op, left, right)
+					if err != nil {
+						return nil, err
+					}
+					exprs[i], err = newSqlComparison(o, l, r)
+					if err != nil {
+						return nil, err
+					}
 				}
 			case clause.Eq:
 				switch left := left.(type) {
@@ -802,14 +801,10 @@ func newSqlComparison(op ast.CmpOp, left clause.Column, right any) (clause.Expre
 		if !ok {
 			return nil, fmt.Errorf("right value in \"not in\" comparison is not a list: %#v", right)
 		}
-		return clause.NotConditions{
-			Exprs: []clause.Expression{
-				clause.IN{
-					Column: left,
-					Values: r,
-				},
-			},
-		}, nil
+		return negativeClause(clause.IN{
+			Column: left,
+			Values: r,
+		}), nil
 	default:
 		return nil, fmt.Errorf("unsupported comparison operation %q", op)
 	}
@@ -831,3 +826,11 @@ func reverseComparison(op ast.CmpOp, left any, right clause.Column) (ast.CmpOp,
 		return op, right, left, fmt.Errorf("unable to reverse comparison operator %q", op)
 	}
 }
+
+func negativeClause(expression clause.Expression) clause.Expression {
+	return clause.NotConditions{
+		Exprs: []clause.Expression{
+			expression,
+		},
+	}
+}


Test patch:
diff --git a/pkg/api/aim/query/query_test.go b/pkg/api/aim/query/query_test.go
index 86be249..4fe774e 100644
--- a/pkg/api/aim/query/query_test.go
+++ b/pkg/api/aim/query/query_test.go
@@ -48,11 +48,17 @@ func (s *QueryTestSuite) TestPostgresDialector_Ok() {
 			expectedVars: []interface{}{"run", models.LifecycleStageDeleted},
 		},
 		{
-			name:         "TestRunNameWithContainsFunction",
-			query:        `(run.name.contains('run'))`,
+			name:         "TestRunNameWithInFunction",
+			query:        `('run' in run.name)`,
 			expectedSQL:  `SELECT * FROM "runs" WHERE ("runs"."name" LIKE $1 AND "runs"."lifecycle_stage" <> $2) ORDER BY "runs"."run_uuid" LIMIT 1`,
 			expectedVars: []interface{}{"%run%", models.LifecycleStageDeleted},
 		},
+		{
+			name:         "TestRunNameWithNotInFunction",
+			query:        `('run' not in run.name)`,
+			expectedSQL:  `SELECT * FROM "runs" WHERE ("runs"."name" NOT LIKE $1 AND "runs"."lifecycle_stage" <> $2) ORDER BY "runs"."run_uuid" LIMIT 1`,
+			expectedVars: []interface{}{"%run%", models.LifecycleStageDeleted},
+		},
 		{
 			name:         "TestRunNameWithStartWithFunction",
 			query:        `(run.name.startswith('run'))`,
@@ -118,11 +124,17 @@ func (s *QueryTestSuite) TestSqliteDialector_Ok() {
 			expectedVars: []interface{}{"run", models.LifecycleStageDeleted},
 		},
 		{
-			name:         "TestRunNameWithContainsFunction",
-			query:        `(run.name.contains('run'))`,
+			name:         "TestRunNameWithInFunction",
+			query:        `('run' in run.name)`,
 			expectedSQL:  `SELECT * FROM "runs" WHERE ("runs"."name" LIKE $1 AND "runs"."lifecycle_stage" <> $2) ORDER BY "runs"."run_uuid" LIMIT 1`,
 			expectedVars: []interface{}{"%run%", models.LifecycleStageDeleted},
 		},
+		{
+			name:         "TestRunNameWithNotInFunction",
+			query:        `('run' not in run.name)`,
+			expectedSQL:  `SELECT * FROM "runs" WHERE ("runs"."name" NOT LIKE $1 AND "runs"."lifecycle_stage" <> $2) ORDER BY "runs"."run_uuid" LIMIT 1`,
+			expectedVars: []interface{}{"%run%", models.LifecycleStageDeleted},
+		},
 		{
 			name:         "TestRunNameWithStartWithFunction",
 			query:        `(run.name.startswith('run'))`,
diff --git a/tests/integration/golang/aim/metric/search_test.go b/tests/integration/golang/aim/metric/search_test.go
index d73ef3a..6418d50 100644
--- a/tests/integration/golang/aim/metric/search_test.go
+++ b/tests/integration/golang/aim/metric/search_test.go
@@ -88,10 +88,6 @@ func (s *SearchMetricsTestSuite) Test_Ok() {
 		name  string
 		query string
 	}{
-		{
-			name:  "TestContainsFunction",
-			query: `q=(metric.name=='key1' and run.name.contains("chill"))&p=500&report_progress=false`,
-		},
 		{
 			name:  "TestStartWithFunction",
 			query: `q=(metric.name=='key1' and run.name.startswith("chill"))&p=500&report_progress=false`,
@@ -108,6 +104,14 @@ func (s *SearchMetricsTestSuite) Test_Ok() {
 			name:  "TestRegexpSearchFunction",
 			query: `q=(metric.name=='key1' and re.search("run", run.name))&p=500&report_progress=false`,
 		},
+		{
+			name:  "TestInFunction",
+			query: `q=(metric.name=='key1' and 'chill' in run.name)&p=500&report_progress=false`,
+		},
+		{
+			name:  "TestNotInFunction",
+			query: `q=(metric.name=='key1' and 'grill' not in run.name)&p=500&report_progress=false`,
+		},
 	}
 	for _, tt := range tests {
 		s.T().Run(tt.name, func(T *testing.T) {