Add tenant and database to the go coordinator (#1303)

## Description of changes *Summarize the changes made by this PR.* - Improvements & Bug fixes - ... - New functionality - ... ## Test plan *How are these changes tested?* - [ ] test_system.py - [ ] Property and example based unit tests in the go coordinator ## Documentation Changes *Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?* --------- Co-authored-by: Liquan Pei <[email protected]>
chroma-core · Oct 28, 2023 · f5ee581 · f5ee581
1 parent 833bb45
commit f5ee581
Show file tree

Hide file tree

Showing 38 changed files with 2,982 additions and 766 deletions.
diff --git a/chromadb/test/db/test_system.py b/chromadb/test/db/test_system.py
@@ -22,23 +22,23 @@
 
 sample_collections = [
     Collection(
-        id=uuid.UUID("93ffe3ec-0107-48d4-8695-51f978c509dc"),
+        id=uuid.UUID(int=1),
         name="test_collection_1",
-        topic="test_topic_1",
+        topic="persistent://test-tenant/test-topic/00000000-0000-0000-0000-000000000001",
         metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3},
         dimension=128,
     ),
     Collection(
-        id=uuid.UUID("f444f1d7-d06c-4357-ac22-5a4a1f92d761"),
+        id=uuid.UUID(int=2),
         name="test_collection_2",
-        topic="test_topic_2",
+        topic="persistent://test-tenant/test-topic/00000000-0000-0000-0000-000000000002",
         metadata={"test_str": "str2", "test_int": 2, "test_float": 2.3},
         dimension=None,
     ),
     Collection(
-        id=uuid.UUID("43babc1a-e403-4a50-91a9-16621ba29ab0"),
+        id=uuid.UUID(int=3),
         name="test_collection_3",
-        topic="test_topic_3",
+        topic="persistent://test-tenant/test-topic/00000000-0000-0000-0000-000000000003",
         metadata={"test_str": "str3", "test_int": 3, "test_float": 3.3},
         dimension=None,
     ),
@@ -105,17 +105,17 @@ def grpc_with_mock_server() -> Generator[SysDB, None, None]:
     yield client
 
 
-# def grpc_with_real_server() -> Generator[SysDB, None, None]:
-#     system = System(
-#         Settings(
-#             allow_reset=True,
-#             chroma_collection_assignment_policy_impl="chromadb.test.db.test_system.MockAssignmentPolicy",
-#         )
-#     )
-#     client = system.instance(GrpcSysDB)
-#     system.start()
-#     client.reset_and_wait_for_ready()
-#     yield client
+def grpc_with_real_server() -> Generator[SysDB, None, None]:
+    system = System(
+        Settings(
+            allow_reset=True,
+            chroma_collection_assignment_policy_impl="chromadb.test.db.test_system.MockAssignmentPolicy",
+        )
+    )
+    client = system.instance(GrpcSysDB)
+    system.start()
+    client.reset_and_wait_for_ready()
+    yield client
 
 
 def db_fixtures() -> List[Callable[[], Generator[SysDB, None, None]]]:

diff --git a/go/coordinator/go.mod b/go/coordinator/go.mod
@@ -3,7 +3,6 @@ module github.com/chroma/chroma-coordinator
 go 1.20
 
 require (
-	github.com/go-sql-driver/mysql v1.7.1
 	github.com/google/uuid v1.3.1
 	github.com/pingcap/log v1.1.0
 	github.com/rs/zerolog v1.31.0
@@ -21,6 +20,7 @@ require (
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/go-sql-driver/mysql v1.7.1 // indirect
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect

diff --git a/go/coordinator/internal/common/constants.go b/go/coordinator/internal/common/constants.go
@@ -0,0 +1,6 @@
+package common
+
+const (
+	DefaultTenant   = "default_tenant"
+	DefaultDatabase = "default_database"
+)
diff --git a/go/coordinator/internal/common/errors.go b/go/coordinator/internal/common/errors.go
@@ -5,11 +5,20 @@ import (
 )
 
 var (
+	// Tenant errors
+	ErrTenantNotFound                  = errors.New("tenant not found")
+	ErrTenantUniqueConstraintViolation = errors.New("tenant unique constraint violation")
+
+	// Database errors
+	ErrDatabaseNotFound                  = errors.New("database not found")
+	ErrDatabaseUniqueConstraintViolation = errors.New("database unique constraint violation")
+
 	// Collection errors
+	ErrCollectionNotFound                    = errors.New("collection not found")
 	ErrCollectionIDFormat                    = errors.New("collection id format error")
 	ErrCollectionNameEmpty                   = errors.New("collection name is empty")
 	ErrCollectionTopicEmpty                  = errors.New("collection topic is empty")
-	ErrCollectionUniqueConstraintViolation   = errors.New("unique constraint violation")
+	ErrCollectionUniqueConstraintViolation   = errors.New("collection unique constraint violation")
 	ErrCollectionDeleteNonExistingCollection = errors.New("delete non existing collection")
 
 	// Collection metadata errors

diff --git a/go/coordinator/internal/coordinator/apis.go b/go/coordinator/internal/coordinator/apis.go
@@ -7,6 +7,8 @@ import (
 	"github.com/chroma/chroma-coordinator/internal/common"
 	"github.com/chroma/chroma-coordinator/internal/model"
 	"github.com/chroma/chroma-coordinator/internal/types"
+	"github.com/pingcap/log"
+	"go.uber.org/zap"
 )
 
 // ICoordinator is an interface that defines the methods for interacting with the
@@ -15,37 +17,73 @@ import (
 type ICoordinator interface {
 	common.Component
 	ResetState(ctx context.Context) error
-	CreateCollection(ctx context.Context, collection *model.CreateCollection) (*model.Collection, error)
-	GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string) ([]*model.Collection, error)
-	DeleteCollection(ctx context.Context, collectionID types.UniqueID) error
-	UpdateCollection(ctx context.Context, collection *model.UpdateCollection) (*model.Collection, error)
-	CreateSegment(ctx context.Context, segment *model.CreateSegment) error
+	CreateCollection(ctx context.Context, createCollection *model.CreateCollection) (*model.Collection, error)
+	GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string, tenantID string, dataName string) ([]*model.Collection, error)
+	DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error
+	UpdateCollection(ctx context.Context, updateCollection *model.UpdateCollection) (*model.Collection, error)
+	CreateSegment(ctx context.Context, createSegment *model.CreateSegment) error
 	GetSegments(ctx context.Context, segmentID types.UniqueID, segmentType *string, scope *string, topic *string, collectionID types.UniqueID) ([]*model.Segment, error)
 	DeleteSegment(ctx context.Context, segmentID types.UniqueID) error
-	UpdateSegment(ctx context.Context, segment *model.UpdateSegment) (*model.Segment, error)
+	UpdateSegment(ctx context.Context, updateSegment *model.UpdateSegment) (*model.Segment, error)
+	CreateDatabase(ctx context.Context, createDatabase *model.CreateDatabase) (*model.Database, error)
+	GetDatabase(ctx context.Context, getDatabase *model.GetDatabase) (*model.Database, error)
+	CreateTenant(ctx context.Context, createTenant *model.CreateTenant) (*model.Tenant, error)
+	GetTenant(ctx context.Context, getTenant *model.GetTenant) (*model.Tenant, error)
 }
 
 func (s *Coordinator) ResetState(ctx context.Context) error {
 	return s.meta.ResetState(ctx)
 }
 
+func (s *Coordinator) CreateDatabase(ctx context.Context, createDatabase *model.CreateDatabase) (*model.Database, error) {
+	database, err := s.meta.CreateDatabase(ctx, createDatabase)
+	if err != nil {
+		return nil, err
+	}
+	return database, nil
+}
+
+func (s *Coordinator) GetDatabase(ctx context.Context, getDatabase *model.GetDatabase) (*model.Database, error) {
+	database, err := s.meta.GetDatabase(ctx, getDatabase)
+	if err != nil {
+		return nil, err
+	}
+	return database, nil
+}
+
+func (s *Coordinator) CreateTenant(ctx context.Context, createTenant *model.CreateTenant) (*model.Tenant, error) {
+	tenant, err := s.meta.CreateTenant(ctx, createTenant)
+	if err != nil {
+		return nil, err
+	}
+	return tenant, nil
+}
+
+func (s *Coordinator) GetTenant(ctx context.Context, getTenant *model.GetTenant) (*model.Tenant, error) {
+	tenant, err := s.meta.GetTenant(ctx, getTenant)
+	if err != nil {
+		return nil, err
+	}
+	return tenant, nil
+}
+
 func (s *Coordinator) CreateCollection(ctx context.Context, createCollection *model.CreateCollection) (*model.Collection, error) {
 	collectionTopic := s.assignCollection(createCollection.ID)
 	createCollection.Topic = collectionTopic
-
+	log.Info("apis create collection", zap.Any("collection", createCollection))
 	collection, err := s.meta.AddCollection(ctx, createCollection)
 	if err != nil {
 		return nil, err
 	}
 	return collection, nil
 }
 
-func (s *Coordinator) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string) ([]*model.Collection, error) {
-	return s.meta.GetCollections(ctx, collectionID, collectionName, collectionTopic)
+func (s *Coordinator) GetCollections(ctx context.Context, collectionID types.UniqueID, collectionName *string, collectionTopic *string, tenantID string, databaseName string) ([]*model.Collection, error) {
+	return s.meta.GetCollections(ctx, collectionID, collectionName, collectionTopic, tenantID, databaseName)
 }
 
-func (s *Coordinator) DeleteCollection(ctx context.Context, collectionID types.UniqueID) error {
-	return s.meta.DeleteCollection(ctx, collectionID)
+func (s *Coordinator) DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error {
+	return s.meta.DeleteCollection(ctx, deleteCollection)
 }
 
 func (s *Coordinator) UpdateCollection(ctx context.Context, collection *model.UpdateCollection) (*model.Collection, error) {
@@ -77,7 +115,6 @@ func (s *Coordinator) UpdateSegment(ctx context.Context, updateSegment *model.Up
 		return nil, err
 	}
 	return segment, nil
-
 }
 
 func verifyCreateCollection(collection *model.CreateCollection) error {