codekingpro/portable-devtools

1from abc import abstractmethod
2from typing import Optional, Sequence, Tuple
3from uuid import UUID
4from chromadb.api.collection_configuration import (
5    CreateCollectionConfiguration,
6    UpdateCollectionConfiguration,
7)
8from chromadb.api.types import Schema
9from chromadb.types import (
10    Collection,
11    CollectionAndSegments,
12    Database,
13    Tenant,
14    Metadata,
15    Segment,
16    SegmentScope,
17    OptionalArgument,
18    Unspecified,
19    UpdateMetadata,
20)
21from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Component
22
23
24class SysDB(Component):
25    """Data interface for Chroma's System database"""
26
27    @abstractmethod
28    def create_database(
29        self, id: UUID, name: str, tenant: str = DEFAULT_TENANT
30    ) -> None:
31        """Create a new database in the System database. Raises an Error if the Database
32        already exists."""
33        pass
34
35    @abstractmethod
36    def get_database(self, name: str, tenant: str = DEFAULT_TENANT) -> Database:
37        """Get a database by name and tenant. Raises an Error if the Database does not
38        exist."""
39        pass
40
41    @abstractmethod
42    def delete_database(self, name: str, tenant: str = DEFAULT_TENANT) -> None:
43        """Delete a database."""
44        pass
45
46    @abstractmethod
47    def list_databases(
48        self,
49        limit: Optional[int] = None,
50        offset: Optional[int] = None,
51        tenant: str = DEFAULT_TENANT,
52    ) -> Sequence[Database]:
53        """List all databases for a tenant."""
54        pass
55
56    @abstractmethod
57    def create_tenant(self, name: str) -> None:
58        """Create a new tenant in the System database. The name must be unique.
59        Raises an Error if the Tenant already exists."""
60        pass
61
62    @abstractmethod
63    def get_tenant(self, name: str) -> Tenant:
64        """Get a tenant by name. Raises an Error if the Tenant does not exist."""
65        pass
66
67    # TODO: Investigate and remove this method, as segment creation is done as
68    # part of collection creation.
69    @abstractmethod
70    def create_segment(self, segment: Segment) -> None:
71        """Create a new segment in the System database. Raises an Error if the ID
72        already exists."""
73        pass
74
75    @abstractmethod
76    def delete_segment(self, collection: UUID, id: UUID) -> None:
77        """Delete a segment from the System database."""
78        pass
79
80    @abstractmethod
81    def get_segments(
82        self,
83        collection: UUID,
84        id: Optional[UUID] = None,
85        type: Optional[str] = None,
86        scope: Optional[SegmentScope] = None,
87    ) -> Sequence[Segment]:
88        """Find segments by id, type, scope or collection."""
89        pass
90
91    @abstractmethod
92    def update_segment(
93        self,
94        collection: UUID,
95        id: UUID,
96        metadata: OptionalArgument[Optional[UpdateMetadata]] = Unspecified(),
97    ) -> None:
98        """Update a segment. Unspecified fields will be left unchanged. For the
99        metadata, keys with None values will be removed and keys not present in the
100        UpdateMetadata dict will be left unchanged."""
101        pass
102
103    @abstractmethod
104    def create_collection(
105        self,
106        id: UUID,
107        name: str,
108        schema: Optional[Schema],
109        configuration: CreateCollectionConfiguration,
110        segments: Sequence[Segment],
111        metadata: Optional[Metadata] = None,
112        dimension: Optional[int] = None,
113        get_or_create: bool = False,
114        tenant: str = DEFAULT_TENANT,
115        database: str = DEFAULT_DATABASE,
116    ) -> Tuple[Collection, bool]:
117        """Create a new collection and associated resources
118        in the SysDB. If get_or_create is True, the
119        collection will be created if one with the same name does not exist.
120        The metadata will be updated using the same protocol as update_collection. If get_or_create
121        is False and the collection already exists, an error will be raised.
122
123        Returns a tuple of the created collection and a boolean indicating whether the
124        collection was created or not.
125        """
126        pass
127
128    @abstractmethod
129    def delete_collection(
130        self,
131        id: UUID,
132        tenant: str = DEFAULT_TENANT,
133        database: str = DEFAULT_DATABASE,
134    ) -> None:
135        """Delete a collection, all associated segments and any associate resources (log stream)
136        from the SysDB and the system at large."""
137        pass
138
139    @abstractmethod
140    def get_collections(
141        self,
142        id: Optional[UUID] = None,
143        name: Optional[str] = None,
144        tenant: str = DEFAULT_TENANT,
145        database: str = DEFAULT_DATABASE,
146        limit: Optional[int] = None,
147        offset: Optional[int] = None,
148    ) -> Sequence[Collection]:
149        """Find collections by id or name. If name is provided, tenant and database must also be provided."""
150        pass
151
152    @abstractmethod
153    def count_collections(
154        self,
155        tenant: str = DEFAULT_TENANT,
156        database: Optional[str] = None,
157    ) -> int:
158        """Gets the number of collections for the (tenant, database) combination."""
159        pass
160
161    @abstractmethod
162    def get_collection_with_segments(
163        self, collection_id: UUID
164    ) -> CollectionAndSegments:
165        """Get a consistent snapshot of a collection by id. This will return a collection with segment
166        information that matches the collection version and log position.
167        """
168        pass
169
170    @abstractmethod
171    def update_collection(
172        self,
173        id: UUID,
174        name: OptionalArgument[str] = Unspecified(),
175        dimension: OptionalArgument[Optional[int]] = Unspecified(),
176        metadata: OptionalArgument[Optional[UpdateMetadata]] = Unspecified(),
177        configuration: OptionalArgument[
178            Optional[UpdateCollectionConfiguration]
179        ] = Unspecified(),
180    ) -> None:
181        """Update a collection. Unspecified fields will be left unchanged. For metadata,
182        keys with None values will be removed and keys not present in the UpdateMetadata
183        dict will be left unchanged."""
184        pass
185
186    @abstractmethod
187    def get_collection_size(self, id: UUID) -> int:
188        """Returns the number of records in a collection."""
189        pass
190