From 8d5f2080b692e28b7f97d02a6d7a8c300b7d8285 Mon Sep 17 00:00:00 2001
From: Tim Sweeney <tim@wandb.com>
Date: Tue, 26 Mar 2024 11:43:12 -0700
Subject: [PATCH] chore(weave): Fix version index lookup

* init

* init
---
 .../clickhouse_trace_server_batched.py        | 35 +++++++++++++------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/weave/trace_server/clickhouse_trace_server_batched.py b/weave/trace_server/clickhouse_trace_server_batched.py
index bcd620b63e7..b9f513a3fb4 100644
--- a/weave/trace_server/clickhouse_trace_server_batched.py
+++ b/weave/trace_server/clickhouse_trace_server_batched.py
@@ -358,12 +358,17 @@ def obj_create(self, req: tsi.ObjCreateReq) -> tsi.ObjCreateRes:
 
     def obj_read(self, req: tsi.ObjReadReq) -> tsi.ObjReadRes:
         conds = ["name = {name: String}"]
-        parameters = {"name": req.name}
+        parameters: typing.Dict[str, typing.Union[str, int]] = {"name": req.name}
         if req.version_digest == "latest":
             conds.append("is_latest = 1")
         else:
-            conds.append("digest = {version_digest: String}")
-            parameters["version_digest"] = req.version_digest
+            (is_version, version_index) = _digest_is_version_like(req.version_digest)
+            if is_version:
+                conds.append("version_index = {version_index: UInt64}")
+                parameters["version_index"] = version_index
+            else:
+                conds.append("digest = {version_digest: String}")
+                parameters["version_digest"] = req.version_digest
         objs = self._select_objs_query(
             req.entity, req.project, conditions=conds, parameters=parameters
         )
@@ -517,13 +522,14 @@ def get_object_ref_root_val(r: refs.ObjectRef) -> typing.Any:
                 val = root_val_cache[cache_key]
             else:
                 if r.version == "latest":
-                    conds = [
-                        "name = {name: String}",
-                        "is_latest = 1",
-                    ]
-                    parameters = {
-                        "name": r.name,
-                    }
+                    raise NotFoundError("Reading refs with `latest` is not supported")
+                    # conds = [
+                    #     "name = {name: String}",
+                    #     "is_latest = 1",
+                    # ]
+                    # parameters = {
+                    #     "name": r.name,
+                    # }
                 else:
                     conds = [
                         "name = {name: String}",
@@ -1179,3 +1185,12 @@ def get_type(val: typing.Any) -> str:
     elif isinstance(val, list):
         return "list"
     return "unknown"
+
+
+def _digest_is_version_like(digest: str) -> typing.Tuple[bool, int]:
+    if not digest.startswith("v"):
+        return (False, -1)
+    try:
+        return (True, int(digest[1:]))
+    except ValueError:
+        return (False, -1)