Spaces:
Running
Running
File size: 1,775 Bytes
287a0bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
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
def find_key_with_value_of_type(
type: AuthzResourceTypes, **kwargs: Any
) -> Dict[str, Any]:
from chromadb.server.fastapi.types import (
CreateCollection,
CreateDatabase,
CreateTenant,
)
for key, value in kwargs.items():
if type == AuthzResourceTypes.DB and isinstance(value, CreateDatabase):
return dict(value)
elif type == AuthzResourceTypes.COLLECTION and isinstance(
value, CreateCollection
):
return dict(value)
elif type == AuthzResourceTypes.TENANT and isinstance(value, CreateTenant):
return dict(value)
return {}
def attr_from_resource_object(
type: AuthzResourceTypes,
additional_attrs: Optional[Sequence[str]] = None,
**kwargs: Any,
) -> Callable[..., Dict[str, Any]]:
def _wrap(**wkwargs: Any) -> Dict[str, Any]:
obj = find_key_with_value_of_type(type, **wkwargs)
if additional_attrs:
obj.update({k: wkwargs["function_kwargs"][k]
for k in additional_attrs})
return obj
return partial(_wrap, **kwargs)
def attr_from_collection_lookup(
collection_id_arg: str, **kwargs: Any
) -> Callable[..., Dict[str, Any]]:
def _wrap(**kwargs: Any) -> Dict[str, Any]:
_api = cast(ServerAPI, kwargs["api"])
col = _api.get_collection(
id=string_to_uuid(kwargs["function_kwargs"][collection_id_arg]))
return {"tenant": col.tenant, "database": col.database}
return partial(_wrap, **kwargs)
|