Skip to content

Commit

Permalink
[ENH]: fix sysdb client grpc (#2068)
Browse files Browse the repository at this point in the history
## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
	 - Fix issue with sysdb grpc client
  • Loading branch information
nicolasgere authored Apr 26, 2024
1 parent 4d6999d commit 32509ee
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 21 deletions.
34 changes: 26 additions & 8 deletions chromadb/db/impl/grpc/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -246,14 +246,32 @@ def get_collections(
offset: Optional[int] = None,
) -> Sequence[Collection]:
# TODO: implement limit and offset in the gRPC service
request = GetCollectionsRequest(
id=id.hex if id else None,
name=name,
tenant=tenant,
database=database,
limit=limit,
offset=offset,
)
request = None
if id is not None:
request = GetCollectionsRequest(
id=id.hex,
limit=limit,
offset=offset,
)
if name is not None:
if tenant is None and database is None:
raise ValueError(
"If name is specified, tenant and database must also be specified in order to uniquely identify the collection"
)
request = GetCollectionsRequest(
name=name,
tenant=tenant,
database=database,
limit=limit,
offset=offset,
)
if id is None and name is None:
request = GetCollectionsRequest(
tenant=tenant,
database=database,
limit=limit,
offset=offset,
)
response: GetCollectionsResponse = self._sys_db_stub.GetCollections(request)
results: List[Collection] = []
for collection in response.collections:
Expand Down
25 changes: 12 additions & 13 deletions chromadb/db/impl/grpc/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -359,20 +359,19 @@ def GetCollections(
target_id = UUID(hex=request.id) if request.HasField("id") else None
target_name = request.name if request.HasField("name") else None

tenant = request.tenant
database = request.database
if tenant not in self._tenants_to_databases_to_collections:
return GetCollectionsResponse(
status=proto.Status(code=404, reason=f"Tenant {tenant} not found")
)
if database not in self._tenants_to_databases_to_collections[tenant]:
return GetCollectionsResponse(
status=proto.Status(code=404, reason=f"Database {database} not found")
)
collections = self._tenants_to_databases_to_collections[tenant][database]

allCollections = {}
for tenant, databases in self._tenants_to_databases_to_collections.items():
for database, collections in databases.items():
if request.tenant != "" and tenant != request.tenant:
continue
if request.database != "" and database != request.database:
continue
allCollections.update(collections)
print(
f"Tenant: {tenant}, Database: {database}, Collections: {collections}"
)
found_collections = []
for collection in collections.values():
for collection in allCollections.values():
if target_id and collection["id"] != target_id:
continue
if target_name and collection["name"] != target_name:
Expand Down

0 comments on commit 32509ee

Please sign in to comment.