chore: add $in to orm query builder (#2180)

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClientInterface/query.ts

@@ -40,6 +40,10 @@ type GtOperation = {
   $gt: [Operand, Operand];
 };
 
+type InOperation = {
+  $in: [Operand, Operand[]];
+};
+
 type GteOperation = {
   $gte: [Operand, Operand];
 };
@@ -61,6 +65,7 @@ type Operation =
   | EqOperation
   | GtOperation
   | GteOperation
+  | InOperation
   | ContainsOperation;
 
 type Operand =

weave/tests/trace/test_client_trace.py

@@ -2297,6 +2297,56 @@ def test_obj(val):
         assert inner_res.count == count
 
 
+def test_in_operation(client):
+    @weave.op()
+    def test_op(label, val):
+        return val
+
+    test_op(1, [1, 2, 3])
+    test_op(2, [1, 2, 3])
+    test_op(3, [5, 6, 7])
+    test_op(4, [8, 2, 3])
+
+    call_ids = [call.id for call in test_op.calls()]
+    assert len(call_ids) == 4
+
+    query = {
+        "$in": [
+            {"$getField": "id"},
+            [{"$literal": call_id} for call_id in call_ids[:2]],
+        ]
+    }
+
+    res = get_client_trace_server(client).calls_query_stats(
+        tsi.CallsQueryStatsReq.model_validate(
+            dict(
+                project_id=get_client_project_id(client),
+                query={"$expr": query},
+            )
+        )
+    )
+    assert res.count == 2
+
+    query = {
+        "$in": [
+            {"$getField": "id"},
+            [{"$literal": call_id} for call_id in call_ids],
+        ]
+    }
+    res = get_client_trace_server(client).calls_query_stream(
+        tsi.CallsQueryReq.model_validate(
+            dict(
+                project_id=get_client_project_id(client),
+                query={"$expr": query},
+            )
+        )
+    )
+    res = list(res)
+    assert len(res) == 4
+    for i in range(4):
+        assert res[i].id == call_ids[i]
+
+
 def test_call_has_client_version(client):
     @weave.op
     def test():

weave/trace_server/calls_query_builder.py

@@ -609,6 +609,10 @@ def process_operation(operation: tsi_query.Operation) -> str:
             lhs_part = process_operand(operation.gte_[0])
             rhs_part = process_operand(operation.gte_[1])
             cond = f"({lhs_part} >= {rhs_part})"
+        elif isinstance(operation, tsi_query.InOperation):
+            lhs_part = process_operand(operation.in_[0])
+            rhs_part = ",".join(process_operand(op) for op in operation.in_[1])
+            cond = f"({lhs_part} IN ({rhs_part}))"
         elif isinstance(operation, tsi_query.ContainsOperation):
             lhs_part = process_operand(operation.contains_.input)
             rhs_part = process_operand(operation.contains_.substr)
@@ -644,6 +648,7 @@ def process_operand(operand: "tsi_query.Operand") -> str:
                 tsi_query.EqOperation,
                 tsi_query.GtOperation,
                 tsi_query.GteOperation,
+                tsi_query.InOperation,
                 tsi_query.ContainsOperation,
             ),
         ):

weave/trace_server/interface/query.py

@@ -121,6 +121,11 @@ class GteOperation(BaseModel):
     gte_: typing.Tuple["Operand", "Operand"] = Field(alias="$gte")
 
 
+# https://www.mongodb.com/docs/manual/reference/operator/aggregation/in/
+class InOperation(BaseModel):
+    in_: typing.Tuple["Operand", list["Operand"]] = Field(alias="$in")
+
+
 # This is not technically in the Mongo spec. Mongo has:
 # https://www.mongodb.com/docs/manual/reference/operator/aggregation/regexMatch/,
 # however, rather than support a full regex match right now, we will
@@ -143,6 +148,7 @@ class ContainsSpec(BaseModel):
     EqOperation,
     GtOperation,
     GteOperation,
+    InOperation,
     ContainsOperation,
 ]
 Operand = typing.Union[
@@ -159,4 +165,5 @@ class ContainsSpec(BaseModel):
 EqOperation.model_rebuild()
 GtOperation.model_rebuild()
 GteOperation.model_rebuild()
+InOperation.model_rebuild()
 ContainsOperation.model_rebuild()

weave/trace_server/orm.py

@@ -626,6 +626,10 @@ def process_operation(operation: tsi_query.Operation) -> str:
             lhs_part = process_operand(operation.gte_[0])
             rhs_part = process_operand(operation.gte_[1])
             cond = f"({lhs_part} >= {rhs_part})"
+        elif isinstance(operation, tsi_query.InOperation):
+            lhs_part = process_operand(operation.in_[0])
+            rhs_part = ",".join(process_operand(op) for op in operation.in_[1])
+            cond = f"({lhs_part} IN ({rhs_part}))"
         elif isinstance(operation, tsi_query.ContainsOperation):
             lhs_part = process_operand(operation.contains_.input)
             rhs_part = process_operand(operation.contains_.substr)
@@ -665,6 +669,7 @@ def process_operand(operand: tsi_query.Operand) -> str:
                 tsi_query.EqOperation,
                 tsi_query.GtOperation,
                 tsi_query.GteOperation,
+                tsi_query.InOperation,
                 tsi_query.ContainsOperation,
             ),
         ):

weave/trace_server/sqlite_trace_server.py

@@ -325,6 +325,10 @@ def process_operation(operation: tsi_query.Operation) -> str:
                     lhs_part = process_operand(operation.gte_[0])
                     rhs_part = process_operand(operation.gte_[1])
                     cond = f"({lhs_part} >= {rhs_part})"
+                elif isinstance(operation, tsi_query.InOperation):
+                    lhs_part = process_operand(operation.in_[0])
+                    rhs_part = ",".join(process_operand(op) for op in operation.in_[1])
+                    cond = f"({lhs_part} IN ({rhs_part}))"
                 elif isinstance(operation, tsi_query.ContainsOperation):
                     lhs_part = process_operand(operation.contains_.input)
                     rhs_part = process_operand(operation.contains_.substr)
@@ -368,6 +372,7 @@ def process_operand(operand: tsi_query.Operand) -> str:
                         tsi_query.EqOperation,
                         tsi_query.GtOperation,
                         tsi_query.GteOperation,
+                        tsi_query.InOperation,
                         tsi_query.ContainsOperation,
                     ),
                 ):