Repository: G-Research/fasttrackml (go)

Issue:
Allow dot syntax when searching tags via Aim UI query syntax
Searching for runs or metrics via run tags can currently be done like this in the Aim UI:

```python
run.tags["my.tag"] == "something"
```

Given that the query field autocomplete suggests bracket syntax for tags containing `.` or `-`, this is totally fine. However, tags that don't contain characters like that get (rightfully) completed with dot syntax, for which we send back a syntax error. Here is an example of a query that _should_ work but doesn't yet:

```python
run.tags.my_tag == "something"
```

Pull request description:
Tags query dot syntax
Allows `runs.tags` query to accept either attribute or subscript style key name.

Fixes #306 

Patch:
diff --git a/pkg/api/aim2/query/query.go b/pkg/api/aim2/query/query.go
index 5ada839..356f0b6 100644
--- a/pkg/api/aim2/query/query.go
+++ b/pkg/api/aim2/query/query.go
@@ -51,6 +51,8 @@ type attributeGetter func(attr string) (any, error)
 
 type subscriptSlicer func(index ast.Slicer) (any, error)
 
+type attributeOrSubscript func(v any) (any, error)
+
 type join struct {
 	key   string
 	alias string
@@ -301,6 +303,8 @@ func (pq *parsedQuery) parseAttribute(node *ast.Attribute) (any, error) {
 		switch value := parsedNode.(type) {
 		case attributeGetter:
 			return value(attribute)
+		case attributeOrSubscript:
+			return value(attribute)
 		default:
 			return nil, fmt.Errorf("unsupported attribute value %#v", value)
 		}
@@ -583,38 +587,19 @@ func (pq *parsedQuery) parseName(node *ast.Name) (any, error) {
 							}
 						}), nil
 					case "tags":
-						return subscriptSlicer(func(s ast.Slicer) (any, error) {
-							switch s := s.(type) {
+						// handle dot (attribute) or dict (subscriptSlicer) syntax
+						return attributeOrSubscript(func(v any) (any, error) {
+							switch v := v.(type) {
+							case string:
+								return pq.tagJoin(v, table)
 							case *ast.Index:
-								v, err := pq.parseNode(s.Value)
+								val, err := pq.parseNode(v.Value)
 								if err != nil {
 									return nil, err
 								}
-								switch v := v.(type) {
-								case string:
-									joinKey := fmt.Sprintf("tags:%s", v)
-									j, ok := pq.joins[joinKey]
-									if !ok {
-										alias := fmt.Sprintf("tags_%d", len(pq.joins))
-										j = join{
-											alias: alias,
-											query: fmt.Sprintf(
-												"LEFT JOIN tags %s ON %s.run_uuid = %s.run_uuid AND %s.key = ?",
-												alias, table, alias, alias,
-											),
-											args: []any{v},
-										}
-										pq.AddJoin(joinKey, j)
-									}
-									return clause.Column{
-										Table: j.alias,
-										Name:  "value",
-									}, nil
-								default:
-									return nil, fmt.Errorf("unsupported index value type %T", v)
-								}
+								return pq.tagsSubscriptSlicer(val, table)
 							default:
-								return nil, fmt.Errorf("unsupported slicer %q", ast.Dump(s))
+								return nil, fmt.Errorf("unsupported slicer or attribute %v", v)
 							}
 						}), nil
 					default:
@@ -762,6 +747,38 @@ func (pq *parsedQuery) metricSubscriptSlicer(v any) (any, error) {
 	}
 }
 
+// tagsSubscriptSlicer will join the tags table using the index key.
+func (pq *parsedQuery) tagsSubscriptSlicer(key any, table string) (any, error) {
+	switch v := key.(type) {
+	case string:
+		return pq.tagJoin(v, table)
+	default:
+		return nil, fmt.Errorf("unsupported index value type %T", v)
+	}
+}
+
+// tagJoin joins the tags table using provided key.
+func (pq *parsedQuery) tagJoin(key string, table string) (any, error) {
+	joinKey := fmt.Sprintf("tags:%s", key)
+	j, ok := pq.joins[joinKey]
+	if !ok {
+		alias := fmt.Sprintf("tags_%d", len(pq.joins))
+		j = join{
+			alias: alias,
+			query: fmt.Sprintf(
+				"LEFT JOIN tags %s ON %s.run_uuid = %s.run_uuid AND %s.key = ?",
+				alias, table, alias, alias,
+			),
+			args: []any{key},
+		}
+		pq.AddJoin(joinKey, j)
+	}
+	return clause.Column{
+		Table: j.alias,
+		Name:  "value",
+	}, nil
+}
+
 // latestMetricsKeyJoin joins the latest_metrics table by run_uuid and metric key, returning the join struct.
 func (pq *parsedQuery) latestMetricsKeyJoin(key, table string) join {
 	joinsKey := fmt.Sprintf("metrics:%s", key)
@@ -883,6 +900,8 @@ func (pq *parsedQuery) parseSubscript(node *ast.Subscript) (any, error) {
 		switch v := v.(type) {
 		case subscriptSlicer:
 			return v(node.Slice)
+		case attributeOrSubscript:
+			return v(node.Slice)
 		default:
 			return nil, fmt.Errorf("unsupported attribute value %#v", v)
 		}


Test patch:
diff --git a/pkg/api/aim2/query/query_test.go b/pkg/api/aim2/query/query_test.go
index 626945a..a6c72fe 100644
--- a/pkg/api/aim2/query/query_test.go
+++ b/pkg/api/aim2/query/query_test.go
@@ -142,6 +142,20 @@ func (s *QueryTestSuite) TestPostgresDialector_Ok() {
 				`AND ("metrics_0"."value" < $4 AND "runs"."lifecycle_stage" <> $5)`,
 			expectedVars: []interface{}{"my_metric", "{key1}", "value1", -1, models.LifecycleStageDeleted},
 		},
+		{
+			name:  "TestTagsSubscript",
+			query: `(run.tags["foo"] == "bar")`,
+			expectedSQL: `SELECT "run_uuid" FROM "runs" LEFT JOIN tags tags_0 ON runs.run_uuid = tags_0.run_uuid ` +
+				`AND tags_0.key = $1 WHERE "tags_0"."value" = $2 AND "runs"."lifecycle_stage" <> $3`,
+			expectedVars: []interface{}{"foo", "bar", models.LifecycleStageDeleted},
+		},
+		{
+			name:  "TestTagsAttribute",
+			query: `(run.tags.foo == "bar")`,
+			expectedSQL: `SELECT "run_uuid" FROM "runs" LEFT JOIN tags tags_0 ON runs.run_uuid = tags_0.run_uuid ` +
+				`AND tags_0.key = $1 WHERE "tags_0"."value" = $2 AND "runs"."lifecycle_stage" <> $3`,
+			expectedVars: []interface{}{"foo", "bar", models.LifecycleStageDeleted},
+		},
 	}
 
 	for _, tt := range tests {