diff --git a/chromadb/db/impl/grpc/client.py b/chromadb/db/impl/grpc/client.py index 0a0c7107e1fe..70818c55e11f 100644 --- a/chromadb/db/impl/grpc/client.py +++ b/chromadb/db/impl/grpc/client.py @@ -1,4 +1,3 @@ -import logging from typing import List, Optional, Sequence, Tuple, Union, cast from uuid import UUID from overrides import overrides @@ -27,6 +26,7 @@ GetDatabaseRequest, GetSegmentsRequest, GetTenantRequest, + ListDatabasesRequest, UpdateCollectionRequest, UpdateSegmentRequest, ) @@ -133,7 +133,26 @@ def list_databases( offset: Optional[int] = None, tenant: str = DEFAULT_TENANT, ) -> Sequence[Database]: - raise NotImplementedError() + try: + request = ListDatabasesRequest(limit=limit, offset=offset, tenant=tenant) + response = self._sys_db_stub.ListDatabases( + request, timeout=self._request_timeout_seconds + ) + results: List[Database] = [] + for proto_database in response.databases: + results.append( + Database( + id=UUID(hex=proto_database.id), + name=proto_database.name, + tenant=proto_database.tenant, + ) + ) + return results + except grpc.RpcError as e: + logger.info( + f"Failed to list databases for tenant {tenant} due to error: {e}" + ) + raise InternalError() @overrides def create_tenant(self, name: str) -> None: diff --git a/chromadb/test/api/test_list_databases.py b/chromadb/test/api/test_list_databases.py index 09e6dcdc7365..7ea1219288f5 100644 --- a/chromadb/test/api/test_list_databases.py +++ b/chromadb/test/api/test_list_databases.py @@ -1,14 +1,10 @@ from typing import Dict, List from hypothesis import given -import pytest -from chromadb.test.conftest import NOT_CLUSTER_ONLY, ClientFactories +from chromadb.test.conftest import ClientFactories import hypothesis.strategies as st def test_list_databases(client_factories: ClientFactories) -> None: - if not NOT_CLUSTER_ONLY: - pytest.skip("This API is not yet supported by distributed") - admin_client = client_factories.create_admin_client_from_system() for i in range(10): @@ -66,9 +62,6 @@ def test_list_databases_with_limit_offset( tenants_and_databases: Dict[str, List[str]], client_factories: ClientFactories, ) -> None: - if not NOT_CLUSTER_ONLY: - pytest.skip("This API is not yet supported by distributed") - client = client_factories.create_client() client.reset() diff --git a/go/pkg/sysdb/coordinator/coordinator.go b/go/pkg/sysdb/coordinator/coordinator.go index e440c0a07c9b..f6910fbc5bc5 100644 --- a/go/pkg/sysdb/coordinator/coordinator.go +++ b/go/pkg/sysdb/coordinator/coordinator.go @@ -66,6 +66,14 @@ func (s *Coordinator) GetDatabase(ctx context.Context, getDatabase *model.GetDat return database, nil } +func (s *Coordinator) ListDatabases(ctx context.Context, listDatabases *model.ListDatabases) ([]*model.Database, error) { + databases, err := s.catalog.ListDatabases(ctx, listDatabases, listDatabases.Ts) + if err != nil { + return nil, err + } + return databases, nil +} + func (s *Coordinator) CreateTenant(ctx context.Context, createTenant *model.CreateTenant) (*model.Tenant, error) { tenant, err := s.catalog.CreateTenant(ctx, createTenant, createTenant.Ts) if err != nil { diff --git a/go/pkg/sysdb/coordinator/model/database.go b/go/pkg/sysdb/coordinator/model/database.go index debbbecb21de..3780e128e4b4 100644 --- a/go/pkg/sysdb/coordinator/model/database.go +++ b/go/pkg/sysdb/coordinator/model/database.go @@ -22,3 +22,10 @@ type GetDatabase struct { Tenant string Ts types.Timestamp } + +type ListDatabases struct { + Limit *int32 + Offset *int32 + Tenant string + Ts types.Timestamp +} diff --git a/go/pkg/sysdb/coordinator/table_catalog.go b/go/pkg/sysdb/coordinator/table_catalog.go index 94d8ef2cd2c7..5d1427570c86 100644 --- a/go/pkg/sysdb/coordinator/table_catalog.go +++ b/go/pkg/sysdb/coordinator/table_catalog.go @@ -148,6 +148,18 @@ func (tc *Catalog) GetDatabases(ctx context.Context, getDatabase *model.GetDatab return result[0], nil } +func (tc *Catalog) ListDatabases(ctx context.Context, listDatabases *model.ListDatabases, ts types.Timestamp) ([]*model.Database, error) { + databases, err := tc.metaDomain.DatabaseDb(ctx).ListDatabases(listDatabases.Limit, listDatabases.Offset, listDatabases.Tenant) + if err != nil { + return nil, err + } + result := make([]*model.Database, 0, len(databases)) + for _, database := range databases { + result = append(result, convertDatabaseToModel(database)) + } + return result, nil +} + func (tc *Catalog) GetAllDatabases(ctx context.Context, ts types.Timestamp) ([]*model.Database, error) { databases, err := tc.metaDomain.DatabaseDb(ctx).GetAllDatabases() if err != nil { diff --git a/go/pkg/sysdb/grpc/tenant_database_service.go b/go/pkg/sysdb/grpc/tenant_database_service.go index 7f373204dbde..2b327c7bb73a 100644 --- a/go/pkg/sysdb/grpc/tenant_database_service.go +++ b/go/pkg/sysdb/grpc/tenant_database_service.go @@ -55,6 +55,31 @@ func (s *Server) GetDatabase(ctx context.Context, req *coordinatorpb.GetDatabase return res, nil } +func (s *Server) ListDatabases(ctx context.Context, req *coordinatorpb.ListDatabasesRequest) (*coordinatorpb.ListDatabasesResponse, error) { + res := &coordinatorpb.ListDatabasesResponse{} + listDatabases := &model.ListDatabases{ + Limit: req.Limit, + Offset: req.Offset, + Tenant: req.GetTenant(), + } + databases, err := s.coordinator.ListDatabases(ctx, listDatabases) + if err != nil { + log.Error("error ListDatabases", zap.String("request", req.String()), zap.Error(err)) + if err == common.ErrTenantNotFound { + return res, grpcutils.BuildNotFoundGrpcError(err.Error()) + } + return res, grpcutils.BuildInternalGrpcError(err.Error()) + } + for _, database := range databases { + res.Databases = append(res.Databases, &coordinatorpb.Database{ + Id: database.ID, + Name: database.Name, + Tenant: database.Tenant, + }) + } + return res, nil +} + func (s *Server) CreateTenant(ctx context.Context, req *coordinatorpb.CreateTenantRequest) (*coordinatorpb.CreateTenantResponse, error) { res := &coordinatorpb.CreateTenantResponse{} createTenant := &model.CreateTenant{ diff --git a/go/pkg/sysdb/metastore/db/dao/database.go b/go/pkg/sysdb/metastore/db/dao/database.go index 68e2071a2bef..1c21c3916af9 100644 --- a/go/pkg/sysdb/metastore/db/dao/database.go +++ b/go/pkg/sysdb/metastore/db/dao/database.go @@ -38,6 +38,28 @@ func (s *databaseDb) GetAllDatabases() ([]*dbmodel.Database, error) { return databases, nil } +func (s *databaseDb) ListDatabases(limit *int32, offset *int32, tenantID string) ([]*dbmodel.Database, error) { + var databases []*dbmodel.Database + query := s.db.Table("databases"). + Select("databases.id, databases.name, databases.tenant_id"). + Where("databases.tenant_id = ?", tenantID). + Order("databases.created_at ASC") + + if limit != nil { + query = query.Limit(int(*limit)) + } + + if offset != nil { + query = query.Offset(int(*offset)) + } + + if err := query.Find(&databases).Error; err != nil { + log.Error("ListDatabases", zap.Error(err)) + return nil, err + } + return databases, nil +} + func (s *databaseDb) GetDatabases(tenantID string, databaseName string) ([]*dbmodel.Database, error) { var databases []*dbmodel.Database query := s.db.Table("databases"). diff --git a/go/pkg/sysdb/metastore/db/dbmodel/database.go b/go/pkg/sysdb/metastore/db/dbmodel/database.go index cfe9a6a0c4ac..b1fca36b7165 100644 --- a/go/pkg/sysdb/metastore/db/dbmodel/database.go +++ b/go/pkg/sysdb/metastore/db/dbmodel/database.go @@ -24,6 +24,7 @@ func (v Database) TableName() string { type IDatabaseDb interface { GetAllDatabases() ([]*Database, error) GetDatabases(tenantID string, databaseName string) ([]*Database, error) + ListDatabases(limit *int32, offset *int32, tenantID string) ([]*Database, error) Insert(in *Database) error DeleteAll() error } diff --git a/idl/chromadb/proto/coordinator.proto b/idl/chromadb/proto/coordinator.proto index 1edb293fd934..a426f1a4b986 100644 --- a/idl/chromadb/proto/coordinator.proto +++ b/idl/chromadb/proto/coordinator.proto @@ -28,6 +28,17 @@ message GetDatabaseResponse { reserved "status"; } +message ListDatabasesRequest { + string tenant = 1; + optional int32 limit = 2; + optional int32 offset = 3; +} + +message ListDatabasesResponse { + repeated Database databases = 1; + reserved "status"; +} + message CreateTenantRequest { string name = 2; // Names are globally unique } @@ -346,6 +357,7 @@ message RestoreCollectionResponse { service SysDB { rpc CreateDatabase(CreateDatabaseRequest) returns (CreateDatabaseResponse) {} rpc GetDatabase(GetDatabaseRequest) returns (GetDatabaseResponse) {} + rpc ListDatabases(ListDatabasesRequest) returns (ListDatabasesResponse) {} rpc CreateTenant(CreateTenantRequest) returns (CreateTenantResponse) {} rpc GetTenant(GetTenantRequest) returns (GetTenantResponse) {} rpc CreateSegment(CreateSegmentRequest) returns (CreateSegmentResponse) {}