From 66e74363161c6ea3071bd26d73ad4756b2fba054 Mon Sep 17 00:00:00 2001 From: gtarpenning Date: Thu, 12 Dec 2024 16:29:43 -0800 Subject: [PATCH] bugs in unit tests --- tests/trace/test_objects_query_builder.py | 8 ++++---- weave/trace_server/objects_query_builder.py | 19 ++++++++++++------- 2 files changed, 16 insertions(+), 11 deletions(-) diff --git a/tests/trace/test_objects_query_builder.py b/tests/trace/test_objects_query_builder.py index dbe645ea5a8..4aa6f5ed3a7 100644 --- a/tests/trace/test_objects_query_builder.py +++ b/tests/trace/test_objects_query_builder.py @@ -47,7 +47,7 @@ def test_make_conditions_part(): assert _make_conditions_part(["condition1"]) == "WHERE condition1" assert ( _make_conditions_part(["condition1", "condition2"]) - == "WHERE condition1 AND condition2" + == "WHERE ((condition1) AND (condition2))" ) @@ -56,7 +56,8 @@ def test_make_object_id_conditions_part(): assert _make_object_id_conditions_part([]) == "" assert _make_object_id_conditions_part(["id = 1"]) == "AND id = 1" assert ( - _make_object_id_conditions_part(["id = 1", "id = 2"]) == "AND id = 1 AND id = 2" + _make_object_id_conditions_part(["id = 1", "id = 2"]) + == "AND ((id = 1) AND (id = 2))" ) @@ -100,8 +101,7 @@ def test_object_query_builder_add_object_ids_condition(): def test_object_query_builder_add_is_op_condition(): builder = ObjectQueryBuilder(project_id="test_project") builder.add_is_op_condition(True) - assert "is_op = {is_op: Boolean}" in builder.conditions_part - assert builder.parameters["is_op"] is True + assert "is_op = 1" in builder.conditions_part def test_object_query_builder_limit_offset(): diff --git a/weave/trace_server/objects_query_builder.py b/weave/trace_server/objects_query_builder.py index ce3533215bb..adfdb0d540c 100644 --- a/weave/trace_server/objects_query_builder.py +++ b/weave/trace_server/objects_query_builder.py @@ -1,5 +1,5 @@ from collections.abc import Iterator -from typing import Any, Literal, Optional, cast +from typing import Any, Optional from weave.trace_server import trace_server_interface as tsi from weave.trace_server.clickhouse_schema import SelectableCHObjSchema @@ -31,10 +31,14 @@ def _make_optional_part(query_keyword: str, part: Optional[str]) -> str: def _make_limit_part(limit: Optional[int]) -> str: + if limit is None: + return "" return _make_optional_part("LIMIT", str(limit)) def _make_offset_part(offset: Optional[int]) -> str: + if offset is None: + return "" return _make_optional_part("OFFSET", str(offset)) @@ -92,7 +96,7 @@ def __init__( self.project_id = project_id self.parameters: dict[str, Any] = parameters or {} if not self.parameters.get(project_id): - self.parameters.update({project_id: project_id}) + self.parameters.update({"project_id": project_id}) self.metadata_only: bool = metadata_only or False self._conditions: list[str] = conditions or [] @@ -166,8 +170,10 @@ def add_is_latest_condition(self) -> None: self._conditions.append("is_latest = 1") def add_is_op_condition(self, is_op: bool) -> None: - self._conditions.append(f"is_op = {{{'is_op': Boolean}}}") - self.parameters.update({"is_op": is_op}) + if is_op: + self._conditions.append("is_op = 1") + else: + self._conditions.append("is_op = 0") def add_base_object_classes_condition(self, base_object_classes: list[str]) -> None: self._conditions.append( @@ -176,10 +182,9 @@ def add_base_object_classes_condition(self, base_object_classes: list[str]) -> N self.parameters.update({"base_object_classes": base_object_classes}) def add_order(self, field: str, direction: str) -> None: - direction = direction.upper() - if direction not in ("ASC", "DESC"): + direction = direction.lower() + if direction not in ("asc", "desc"): raise ValueError(f"Direction {direction} is not allowed") - direction = cast(Literal["ASC", "DESC"], direction) self._sort_by.append(tsi.SortBy(field=field, direction=direction)) def set_limit(self, limit: int) -> None: