From a15c78053d7afa7cc0c7bcb9d29fe7719000c8b7 Mon Sep 17 00:00:00 2001
From: Robert Escriva <robert@trychroma.com>
Date: Wed, 18 Dec 2024 14:54:46 -0800
Subject: [PATCH] [CHORE]  Add instrumentation to a slow Python call.

---
 chromadb/server/fastapi/__init__.py | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/chromadb/server/fastapi/__init__.py b/chromadb/server/fastapi/__init__.py
index be583c1206f..5a04c5c5ba7 100644
--- a/chromadb/server/fastapi/__init__.py
+++ b/chromadb/server/fastapi/__init__.py
@@ -710,17 +710,20 @@ async def get_collection(
 
         add_attributes_to_current_span({"tenant": tenant})
 
-        api_collection_model = cast(
-            CollectionModel,
-            await to_thread.run_sync(
-                self._api.get_collection,
-                collection_name,
-                tenant,
-                database_name,
-                limiter=self._capacity_limiter,
-            ),
-        )
-        return api_collection_model
+        @trace_method("Outside the pool", OpenTelemetryGranularity.OPERATION)
+        def inner():
+            api_collection_model = cast(
+                CollectionModel,
+                await to_thread.run_sync(
+                    self._api.get_collection,
+                    collection_name,
+                    tenant,
+                    database_name,
+                    limiter=self._capacity_limiter,
+                ),
+            )
+            return api_collection_model
+        return inner()
 
     @trace_method("FastAPI.update_collection", OpenTelemetryGranularity.OPERATION)
     async def update_collection(