From 22da192dcc42655dc865781b79fa1e013ebf682f Mon Sep 17 00:00:00 2001
From: Hammad Bashir <HammadB@users.noreply.github.com>
Date: Tue, 16 Jan 2024 14:08:15 -0800
Subject: [PATCH] [BUG] UUID should be uuid type not str (#1644)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- fastapi_utils needs to push the uuid as a UUID type not str to
get_collection.
	 - Refactored _uuid into a util
 - New functionality
	 - None
---
 chromadb/auth/fastapi_utils.py      |  3 ++-
 chromadb/server/fastapi/__init__.py | 10 +---------
 chromadb/server/fastapi/utils.py    |  9 ++++++++-
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/chromadb/auth/fastapi_utils.py b/chromadb/auth/fastapi_utils.py
index 881eee0d3ef..2612bf6716f 100644
--- a/chromadb/auth/fastapi_utils.py
+++ b/chromadb/auth/fastapi_utils.py
@@ -1,5 +1,6 @@
 from functools import partial
 from typing import Any, Callable, Dict, Optional, Sequence, cast
+from chromadb.server.fastapi.utils import string_to_uuid
 from chromadb.api import ServerAPI
 from chromadb.auth import AuthzResourceTypes
 
@@ -46,7 +47,7 @@ def attr_from_collection_lookup(
     def _wrap(**kwargs: Any) -> Dict[str, Any]:
         _api = cast(ServerAPI, kwargs["api"])
         col = _api.get_collection(
-            id=kwargs["function_kwargs"][collection_id_arg])
+            id=string_to_uuid(kwargs["function_kwargs"][collection_id_arg]))
         return {"tenant": col.tenant, "database": col.database}
 
     return partial(_wrap, **kwargs)
diff --git a/chromadb/server/fastapi/__init__.py b/chromadb/server/fastapi/__init__.py
index 826e842824f..f01f4137908 100644
--- a/chromadb/server/fastapi/__init__.py
+++ b/chromadb/server/fastapi/__init__.py
@@ -32,7 +32,6 @@
 from chromadb.api import ServerAPI
 from chromadb.errors import (
     ChromaError,
-    InvalidUUIDError,
     InvalidDimensionException,
     InvalidHTTPVersion,
 )
@@ -51,7 +50,7 @@
 
 import logging
 
-from chromadb.server.fastapi.utils import fastapi_json_response
+from chromadb.server.fastapi.utils import fastapi_json_response, string_to_uuid as _uuid
 from chromadb.telemetry.opentelemetry.fastapi import instrument_fastapi
 from chromadb.types import Database, Tenant
 from chromadb.telemetry.product import ServerContext, ProductTelemetryClient
@@ -96,13 +95,6 @@ async def check_http_version_middleware(
     return await call_next(request)
 
 
-def _uuid(uuid_str: str) -> UUID:
-    try:
-        return UUID(uuid_str)
-    except ValueError:
-        raise InvalidUUIDError(f"Could not parse {uuid_str} as a UUID")
-
-
 class ChromaAPIRouter(fastapi.APIRouter):  # type: ignore
     # A simple subclass of fastapi's APIRouter which treats URLs with a trailing "/" the
     # same as URLs without. Docs will only contain URLs without trailing "/"s.
diff --git a/chromadb/server/fastapi/utils.py b/chromadb/server/fastapi/utils.py
index ccbb68eee41..b7e781dae68 100644
--- a/chromadb/server/fastapi/utils.py
+++ b/chromadb/server/fastapi/utils.py
@@ -1,6 +1,7 @@
+from uuid import UUID
 from starlette.responses import JSONResponse
 
-from chromadb.errors import ChromaError
+from chromadb.errors import ChromaError, InvalidUUIDError
 
 
 def fastapi_json_response(error: ChromaError) -> JSONResponse:
@@ -8,3 +9,9 @@ def fastapi_json_response(error: ChromaError) -> JSONResponse:
         content={"error": error.name(), "message": error.message()},
         status_code=error.code(),
     )
+
+def string_to_uuid(uuid_str: str) -> UUID:
+    try:
+        return UUID(uuid_str)
+    except ValueError:
+        raise InvalidUUIDError(f"Could not parse {uuid_str} as a UUID")
\ No newline at end of file