fix(trace): misc bugs with objs_query when deleting objects

Author: bcsherma

tests/trace/test_objects_query_builder.py

@@ -177,6 +177,11 @@ def test_object_query_builder_sort():
         ) AS row_num,
         if (row_num = 1, 1, 0) AS is_latest
     FROM (
+        --
+        -- Object versions are uniquely identified by (kind, project_id, object_id, digest).
+        -- This subquery selects a row to represent each object version. There are multiple rows
+        -- for each object version if it has been deleted or recreated prior to a table merge.
+        --
         SELECT
             project_id,
             object_id,
@@ -193,7 +198,15 @@ def test_object_query_builder_sort():
                 kind,
                 object_id,
                 digest
-                ORDER BY created_at ASC
+                --
+                -- Prefer the most recent row. If there is a tie, prefer the row
+                -- with non-null deleted_at, which represents the deletion event.
+                --
+                -- Rows for the same object version may have the same created_at
+                -- because deletion events inherit the created_at of the last
+                -- non-deleted row for the object version.
+                --
+                ORDER BY created_at DESC, (deleted_at IS NULL) ASC
             ) AS rn
         FROM object_versions"""
 
@@ -289,7 +302,7 @@ def test_make_objects_val_query_and_parameters():
     )
 
     expected_query = """
-        SELECT object_id, digest, any(val_dump)
+        SELECT object_id, digest, argMax(val_dump, created_at)
         FROM object_versions
         WHERE project_id = {project_id: String} AND
             object_id IN {object_ids: Array(String)} AND

tests/trace/test_objs_query.py

@@ -185,3 +185,140 @@ def test_objs_query_wb_user_id(client: WeaveClient):
     res = client._objects()
     assert len(res) == 3
     assert all(obj.wb_user_id == correct_id for obj in res)
+
+
+def test_objs_query_deleted_interaction(client: WeaveClient):
+    weave.publish({"i": 1}, name="obj_1")
+    weave.publish({"i": 2}, name="obj_1")
+    weave.publish({"i": 3}, name="obj_1")
+
+    res = client.server.objs_query(
+        tsi.ObjQueryReq(
+            project_id=client._project_id(),
+            filter=tsi.ObjectVersionFilter(latest_only=False),
+        )
+    )
+    assert len(res.objs) == 3
+    assert all(obj.val["i"] in [1, 2, 3] for obj in res.objs)
+
+    res = client.server.obj_delete(
+        tsi.ObjDeleteReq(
+            project_id=client._project_id(),
+            object_id="obj_1",
+            digests=[res.objs[0].digest],
+        )
+    )
+
+    assert res.num_deleted == 1
+
+    res = client.server.objs_query(
+        tsi.ObjQueryReq(
+            project_id=client._project_id(),
+            filter=tsi.ObjectVersionFilter(latest_only=False),
+        )
+    )
+    assert len(res.objs) == 2
+    assert all(obj.val["i"] in [2, 3] for obj in res.objs)
+
+    # Delete the remaining objects
+    res = client.server.obj_delete(
+        tsi.ObjDeleteReq(
+            project_id=client._project_id(),
+            object_id="obj_1",
+            digests=[res.objs[0].digest, res.objs[1].digest],
+        )
+    )
+    assert res.num_deleted == 2
+
+    res = client.server.objs_query(
+        tsi.ObjQueryReq(
+            project_id=client._project_id(),
+            filter=tsi.ObjectVersionFilter(latest_only=False),
+        )
+    )
+    assert len(res.objs) == 0
+
+
+def test_objs_query_delete_and_recreate(client: WeaveClient):
+    weave.publish({"i": 1}, name="obj_1")
+    weave.publish({"i": 2}, name="obj_1")
+    weave.publish({"i": 3}, name="obj_1")
+
+    res = client.server.objs_query(
+        tsi.ObjQueryReq(
+            project_id=client._project_id(),
+            filter=tsi.ObjectVersionFilter(latest_only=False),
+        )
+    )
+    assert len(res.objs) == 3
+
+    original_created_at = res.objs[0].created_at
+
+    res = client.server.obj_delete(
+        tsi.ObjDeleteReq(
+            project_id=client._project_id(),
+            object_id="obj_1",
+        )
+    )
+    assert res.num_deleted == 3
+
+    weave.publish({"i": 1}, name="obj_1")
+
+    res = client.server.objs_query(
+        tsi.ObjQueryReq(
+            project_id=client._project_id(),
+            filter=tsi.ObjectVersionFilter(latest_only=False),
+        )
+    )
+    assert len(res.objs) == 1
+    assert res.objs[0].val["i"] == 1
+    assert res.objs[0].created_at > original_created_at
+
+    weave.publish({"i": 2}, name="obj_1")
+    weave.publish({"i": 3}, name="obj_1")
+
+    res = client.server.objs_query(
+        tsi.ObjQueryReq(
+            project_id=client._project_id(),
+            filter=tsi.ObjectVersionFilter(latest_only=False),
+        )
+    )
+    assert len(res.objs) == 3
+
+    for i in range(3):
+        print("res.objs[i].val", res.objs[i].val)
+        assert res.objs[i].val["i"] == i + 1
+
+
+def test_objs_query_delete_and_add_new_versions(client: WeaveClient):
+    weave.publish({"i": 1}, name="obj_1")
+    weave.publish({"i": 2}, name="obj_1")
+    weave.publish({"i": 3}, name="obj_1")
+
+    res = client.server.objs_query(
+        tsi.ObjQueryReq(
+            project_id=client._project_id(),
+            filter=tsi.ObjectVersionFilter(latest_only=False),
+        )
+    )
+    assert len(res.objs) == 3
+
+    res = client.server.obj_delete(
+        tsi.ObjDeleteReq(
+            project_id=client._project_id(),
+            object_id="obj_1",
+        )
+    )
+
+    weave.publish({"i": 4}, name="obj_1")
+    weave.publish({"i": 5}, name="obj_1")
+    weave.publish({"i": 6}, name="obj_1")
+
+    res = client.server.objs_query(
+        tsi.ObjQueryReq(
+            project_id=client._project_id(),
+            filter=tsi.ObjectVersionFilter(latest_only=False),
+        )
+    )
+    assert len(res.objs) == 3
+    assert all(obj.val["i"] in [4, 5, 6] for obj in res.objs)

weave/trace_server/objects_query_builder.py

@@ -273,6 +273,11 @@ def make_metadata_query(self) -> str:
         ) AS row_num,
         if (row_num = 1, 1, 0) AS is_latest
     FROM (
+        --
+        -- Object versions are uniquely identified by (kind, project_id, object_id, digest).
+        -- This subquery selects a row to represent each object version. There are multiple rows
+        -- for each object version if it has been deleted or recreated prior to a table merge.
+        --
         SELECT
             project_id,
             object_id,
@@ -289,7 +294,15 @@ def make_metadata_query(self) -> str:
                 kind,
                 object_id,
                 digest
-                ORDER BY created_at ASC
+                --
+                -- Prefer the most recent row. If there is a tie, prefer the row
+                -- with non-null deleted_at, which represents the deletion event.
+                --
+                -- Rows for the same object version may have the same created_at
+                -- because deletion events inherit the created_at of the last
+                -- non-deleted row for the object version.
+                --
+                ORDER BY created_at DESC, (deleted_at IS NULL) ASC
             ) AS rn
         FROM object_versions
         WHERE project_id = {{project_id: String}}{self.object_id_conditions_part}
@@ -311,7 +324,7 @@ def make_objects_val_query_and_parameters(
     project_id: str, object_ids: list[str], digests: list[str]
 ) -> tuple[str, dict[str, Any]]:
     query = """
-        SELECT object_id, digest, any(val_dump)
+        SELECT object_id, digest, argMax(val_dump, created_at)
         FROM object_versions
         WHERE project_id = {project_id: String} AND
             object_id IN {object_ids: Array(String)} AND