himanshud2611's picture
Upload folder using huggingface_hub
60e3a80 verified
import json
from typing import Optional, Sequence, Any, Tuple, cast, Dict, Union, Set
from uuid import UUID
from overrides import override
from pypika import Table, Column
from itertools import groupby
from chromadb.api.configuration import (
CollectionConfigurationInternal,
ConfigurationParameter,
HNSWConfigurationInternal,
InvalidConfigurationError,
)
from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, System
from chromadb.db.base import (
Cursor,
SqlDB,
ParameterValue,
get_sql,
UniqueConstraintError,
)
from chromadb.db.system import SysDB
from chromadb.errors import NotFoundError
from chromadb.telemetry.opentelemetry import (
add_attributes_to_current_span,
OpenTelemetryClient,
OpenTelemetryGranularity,
trace_method,
)
from chromadb.ingest import Producer
from chromadb.types import (
Database,
OptionalArgument,
Segment,
Metadata,
Collection,
SegmentScope,
Tenant,
Unspecified,
UpdateMetadata,
)
class SqlSysDB(SqlDB, SysDB):
# Used only to delete log streams on collection deletion.
# TODO: refactor to remove this dependency into a separate interface
_producer: Producer
def __init__(self, system: System):
super().__init__(system)
self._opentelemetry_client = system.require(OpenTelemetryClient)
@trace_method("SqlSysDB.create_segment", OpenTelemetryGranularity.ALL)
@override
def start(self) -> None:
super().start()
self._producer = self._system.instance(Producer)
@override
def create_database(
self, id: UUID, name: str, tenant: str = DEFAULT_TENANT
) -> None:
with self.tx() as cur:
# Get the tenant id for the tenant name and then insert the database with the id, name and tenant id
databases = Table("databases")
tenants = Table("tenants")
insert_database = (
self.querybuilder()
.into(databases)
.columns(databases.id, databases.name, databases.tenant_id)
.insert(
ParameterValue(self.uuid_to_db(id)),
ParameterValue(name),
self.querybuilder()
.select(tenants.id)
.from_(tenants)
.where(tenants.id == ParameterValue(tenant)),
)
)
sql, params = get_sql(insert_database, self.parameter_format())
try:
cur.execute(sql, params)
except self.unique_constraint_error() as e:
raise UniqueConstraintError(
f"Database {name} already exists for tenant {tenant}"
) from e
@override
def get_database(self, name: str, tenant: str = DEFAULT_TENANT) -> Database:
with self.tx() as cur:
databases = Table("databases")
q = (
self.querybuilder()
.from_(databases)
.select(databases.id, databases.name)
.where(databases.name == ParameterValue(name))
.where(databases.tenant_id == ParameterValue(tenant))
)
sql, params = get_sql(q, self.parameter_format())
row = cur.execute(sql, params).fetchone()
if not row:
raise NotFoundError(
f"Database {name} not found for tenant {tenant}. Are you sure it exists?"
)
if row[0] is None:
raise NotFoundError(
f"Database {name} not found for tenant {tenant}. Are you sure it exists?"
)
id: UUID = cast(UUID, self.uuid_from_db(row[0]))
return Database(
id=id,
name=row[1],
tenant=tenant,
)
@override
def create_tenant(self, name: str) -> None:
with self.tx() as cur:
tenants = Table("tenants")
insert_tenant = (
self.querybuilder()
.into(tenants)
.columns(tenants.id)
.insert(ParameterValue(name))
)
sql, params = get_sql(insert_tenant, self.parameter_format())
try:
cur.execute(sql, params)
except self.unique_constraint_error() as e:
raise UniqueConstraintError(f"Tenant {name} already exists") from e
@override
def get_tenant(self, name: str) -> Tenant:
with self.tx() as cur:
tenants = Table("tenants")
q = (
self.querybuilder()
.from_(tenants)
.select(tenants.id)
.where(tenants.id == ParameterValue(name))
)
sql, params = get_sql(q, self.parameter_format())
row = cur.execute(sql, params).fetchone()
if not row:
raise NotFoundError(f"Tenant {name} not found")
return Tenant(name=name)
@override
def create_segment(self, segment: Segment) -> None:
add_attributes_to_current_span(
{
"segment_id": str(segment["id"]),
"segment_type": segment["type"],
"segment_scope": segment["scope"].value,
"collection": str(segment["collection"]),
}
)
with self.tx() as cur:
segments = Table("segments")
insert_segment = (
self.querybuilder()
.into(segments)
.columns(
segments.id,
segments.type,
segments.scope,
segments.collection,
)
.insert(
ParameterValue(self.uuid_to_db(segment["id"])),
ParameterValue(segment["type"]),
ParameterValue(segment["scope"].value),
ParameterValue(self.uuid_to_db(segment["collection"])),
)
)
sql, params = get_sql(insert_segment, self.parameter_format())
try:
cur.execute(sql, params)
except self.unique_constraint_error() as e:
raise UniqueConstraintError(
f"Segment {segment['id']} already exists"
) from e
metadata_t = Table("segment_metadata")
if segment["metadata"]:
self._insert_metadata(
cur,
metadata_t,
metadata_t.segment_id,
segment["id"],
segment["metadata"],
)
@trace_method("SqlSysDB.create_collection", OpenTelemetryGranularity.ALL)
@override
def create_collection(
self,
id: UUID,
name: str,
configuration: CollectionConfigurationInternal,
metadata: Optional[Metadata] = None,
dimension: Optional[int] = None,
get_or_create: bool = False,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> Tuple[Collection, bool]:
if id is None and not get_or_create:
raise ValueError("id must be specified if get_or_create is False")
add_attributes_to_current_span(
{
"collection_id": str(id),
"collection_name": name,
}
)
existing = self.get_collections(name=name, tenant=tenant, database=database)
if existing:
if get_or_create:
collection = existing[0]
return (
self.get_collections(
id=collection.id, tenant=tenant, database=database
)[0],
False,
)
else:
raise UniqueConstraintError(f"Collection {name} already exists")
collection = Collection(
id=id,
name=name,
configuration=configuration,
metadata=metadata,
dimension=dimension,
tenant=tenant,
database=database,
version=0,
)
with self.tx() as cur:
collections = Table("collections")
databases = Table("databases")
insert_collection = (
self.querybuilder()
.into(collections)
.columns(
collections.id,
collections.name,
collections.config_json_str,
collections.dimension,
collections.database_id,
)
.insert(
ParameterValue(self.uuid_to_db(collection["id"])),
ParameterValue(collection["name"]),
ParameterValue(configuration.to_json_str()),
ParameterValue(collection["dimension"]),
# Get the database id for the database with the given name and tenant
self.querybuilder()
.select(databases.id)
.from_(databases)
.where(databases.name == ParameterValue(database))
.where(databases.tenant_id == ParameterValue(tenant)),
)
)
sql, params = get_sql(insert_collection, self.parameter_format())
try:
cur.execute(sql, params)
except self.unique_constraint_error() as e:
raise UniqueConstraintError(
f"Collection {collection['id']} already exists"
) from e
metadata_t = Table("collection_metadata")
if collection["metadata"]:
self._insert_metadata(
cur,
metadata_t,
metadata_t.collection_id,
collection.id,
collection["metadata"],
)
return collection, True
@trace_method("SqlSysDB.get_segments", OpenTelemetryGranularity.ALL)
@override
def get_segments(
self,
collection: UUID,
id: Optional[UUID] = None,
type: Optional[str] = None,
scope: Optional[SegmentScope] = None,
) -> Sequence[Segment]:
add_attributes_to_current_span(
{
"segment_id": str(id),
"segment_type": type if type else "",
"segment_scope": scope.value if scope else "",
"collection": str(collection),
}
)
segments_t = Table("segments")
metadata_t = Table("segment_metadata")
q = (
self.querybuilder()
.from_(segments_t)
.select(
segments_t.id,
segments_t.type,
segments_t.scope,
segments_t.collection,
metadata_t.key,
metadata_t.str_value,
metadata_t.int_value,
metadata_t.float_value,
metadata_t.bool_value,
)
.left_join(metadata_t)
.on(segments_t.id == metadata_t.segment_id)
.orderby(segments_t.id)
)
if id:
q = q.where(segments_t.id == ParameterValue(self.uuid_to_db(id)))
if type:
q = q.where(segments_t.type == ParameterValue(type))
if scope:
q = q.where(segments_t.scope == ParameterValue(scope.value))
if collection:
q = q.where(
segments_t.collection == ParameterValue(self.uuid_to_db(collection))
)
with self.tx() as cur:
sql, params = get_sql(q, self.parameter_format())
rows = cur.execute(sql, params).fetchall()
by_segment = groupby(rows, lambda r: cast(object, r[0]))
segments = []
for segment_id, segment_rows in by_segment:
id = self.uuid_from_db(str(segment_id))
rows = list(segment_rows)
type = str(rows[0][1])
scope = SegmentScope(str(rows[0][2]))
collection = self.uuid_from_db(rows[0][3]) # type: ignore[assignment]
metadata = self._metadata_from_rows(rows)
segments.append(
Segment(
id=cast(UUID, id),
type=type,
scope=scope,
collection=collection,
metadata=metadata,
)
)
return segments
@trace_method("SqlSysDB.get_collections", OpenTelemetryGranularity.ALL)
@override
def get_collections(
self,
id: Optional[UUID] = None,
name: Optional[str] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
limit: Optional[int] = None,
offset: Optional[int] = None,
) -> Sequence[Collection]:
"""Get collections by name, embedding function and/or metadata"""
if name is not None and (tenant is None or database is None):
raise ValueError(
"If name is specified, tenant and database must also be specified in order to uniquely identify the collection"
)
add_attributes_to_current_span(
{
"collection_id": str(id),
"collection_name": name if name else "",
}
)
collections_t = Table("collections")
metadata_t = Table("collection_metadata")
databases_t = Table("databases")
q = (
self.querybuilder()
.from_(collections_t)
.select(
collections_t.id,
collections_t.name,
collections_t.config_json_str,
collections_t.dimension,
databases_t.name,
databases_t.tenant_id,
metadata_t.key,
metadata_t.str_value,
metadata_t.int_value,
metadata_t.float_value,
metadata_t.bool_value,
)
.left_join(metadata_t)
.on(collections_t.id == metadata_t.collection_id)
.left_join(databases_t)
.on(collections_t.database_id == databases_t.id)
.orderby(collections_t.id)
)
if id:
q = q.where(collections_t.id == ParameterValue(self.uuid_to_db(id)))
if name:
q = q.where(collections_t.name == ParameterValue(name))
# Only if we have a name, tenant and database do we need to filter databases
# Given an id, we can uniquely identify the collection so we don't need to filter databases
if id is None and tenant and database:
databases_t = Table("databases")
q = q.where(
collections_t.database_id
== self.querybuilder()
.select(databases_t.id)
.from_(databases_t)
.where(databases_t.name == ParameterValue(database))
.where(databases_t.tenant_id == ParameterValue(tenant))
)
# cant set limit and offset here because this is metadata and we havent reduced yet
with self.tx() as cur:
sql, params = get_sql(q, self.parameter_format())
rows = cur.execute(sql, params).fetchall()
by_collection = groupby(rows, lambda r: cast(object, r[0]))
collections = []
for collection_id, collection_rows in by_collection:
id = self.uuid_from_db(str(collection_id))
rows = list(collection_rows)
name = str(rows[0][1])
metadata = self._metadata_from_rows(rows)
dimension = int(rows[0][3]) if rows[0][3] else None
if rows[0][2] is not None:
configuration = self._load_config_from_json_str_and_migrate(
str(collection_id), rows[0][2]
)
else:
# 07/2024: This is a legacy case where we don't have a collection
# configuration stored in the database. This non-destructively migrates
# the collection to have a configuration, and takes into account any
# HNSW params that might be in the existing metadata.
configuration = self._insert_config_from_legacy_params(
collection_id, metadata
)
collections.append(
Collection(
id=cast(UUID, id),
name=name,
configuration=configuration,
metadata=metadata,
dimension=dimension,
tenant=str(rows[0][5]),
database=str(rows[0][4]),
version=0,
)
)
# apply limit and offset
if limit is not None:
if offset is None:
offset = 0
collections = collections[offset : offset + limit]
else:
collections = collections[offset:]
return collections
@trace_method("SqlSysDB.delete_segment", OpenTelemetryGranularity.ALL)
@override
def delete_segment(self, collection: UUID, id: UUID) -> None:
"""Delete a segment from the SysDB"""
add_attributes_to_current_span(
{
"segment_id": str(id),
}
)
t = Table("segments")
q = (
self.querybuilder()
.from_(t)
.where(t.id == ParameterValue(self.uuid_to_db(id)))
.delete()
)
with self.tx() as cur:
# no need for explicit del from metadata table because of ON DELETE CASCADE
sql, params = get_sql(q, self.parameter_format())
sql = sql + " RETURNING id"
result = cur.execute(sql, params).fetchone()
if not result:
raise NotFoundError(f"Segment {id} not found")
@trace_method("SqlSysDB.delete_collection", OpenTelemetryGranularity.ALL)
@override
def delete_collection(
self,
id: UUID,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> None:
"""Delete a collection and all associated segments from the SysDB. Deletes
the log stream for this collection as well."""
add_attributes_to_current_span(
{
"collection_id": str(id),
}
)
t = Table("collections")
databases_t = Table("databases")
q = (
self.querybuilder()
.from_(t)
.where(t.id == ParameterValue(self.uuid_to_db(id)))
.where(
t.database_id
== self.querybuilder()
.select(databases_t.id)
.from_(databases_t)
.where(databases_t.name == ParameterValue(database))
.where(databases_t.tenant_id == ParameterValue(tenant))
)
.delete()
)
with self.tx() as cur:
# no need for explicit del from metadata table because of ON DELETE CASCADE
sql, params = get_sql(q, self.parameter_format())
sql = sql + " RETURNING id"
result = cur.execute(sql, params).fetchone()
if not result:
raise NotFoundError(f"Collection {id} not found")
self._producer.delete_log(result[0])
@trace_method("SqlSysDB.update_segment", OpenTelemetryGranularity.ALL)
@override
def update_segment(
self,
collection: UUID,
id: UUID,
metadata: OptionalArgument[Optional[UpdateMetadata]] = Unspecified(),
) -> None:
add_attributes_to_current_span(
{
"segment_id": str(id),
"collection": str(collection),
}
)
segments_t = Table("segments")
metadata_t = Table("segment_metadata")
q = (
self.querybuilder()
.update(segments_t)
.where(segments_t.id == ParameterValue(self.uuid_to_db(id)))
.set(segments_t.collection, ParameterValue(self.uuid_to_db(collection)))
)
with self.tx() as cur:
sql, params = get_sql(q, self.parameter_format())
if sql: # pypika emits a blank string if nothing to do
cur.execute(sql, params)
if metadata is None:
q = (
self.querybuilder()
.from_(metadata_t)
.where(metadata_t.segment_id == ParameterValue(self.uuid_to_db(id)))
.delete()
)
sql, params = get_sql(q, self.parameter_format())
cur.execute(sql, params)
elif metadata != Unspecified():
metadata = cast(UpdateMetadata, metadata)
metadata = cast(UpdateMetadata, metadata)
self._insert_metadata(
cur,
metadata_t,
metadata_t.segment_id,
id,
metadata,
set(metadata.keys()),
)
@trace_method("SqlSysDB.update_collection", OpenTelemetryGranularity.ALL)
@override
def update_collection(
self,
id: UUID,
name: OptionalArgument[str] = Unspecified(),
dimension: OptionalArgument[Optional[int]] = Unspecified(),
metadata: OptionalArgument[Optional[UpdateMetadata]] = Unspecified(),
) -> None:
add_attributes_to_current_span(
{
"collection_id": str(id),
}
)
collections_t = Table("collections")
metadata_t = Table("collection_metadata")
q = (
self.querybuilder()
.update(collections_t)
.where(collections_t.id == ParameterValue(self.uuid_to_db(id)))
)
if not name == Unspecified():
q = q.set(collections_t.name, ParameterValue(name))
if not dimension == Unspecified():
q = q.set(collections_t.dimension, ParameterValue(dimension))
with self.tx() as cur:
sql, params = get_sql(q, self.parameter_format())
if sql: # pypika emits a blank string if nothing to do
sql = sql + " RETURNING id"
result = cur.execute(sql, params)
if not result.fetchone():
raise NotFoundError(f"Collection {id} not found")
# TODO: Update to use better semantics where it's possible to update
# individual keys without wiping all the existing metadata.
# For now, follow current legancy semantics where metadata is fully reset
if metadata != Unspecified():
q = (
self.querybuilder()
.from_(metadata_t)
.where(
metadata_t.collection_id == ParameterValue(self.uuid_to_db(id))
)
.delete()
)
sql, params = get_sql(q, self.parameter_format())
cur.execute(sql, params)
if metadata is not None:
metadata = cast(UpdateMetadata, metadata)
self._insert_metadata(
cur,
metadata_t,
metadata_t.collection_id,
id,
metadata,
set(metadata.keys()),
)
@trace_method("SqlSysDB._metadata_from_rows", OpenTelemetryGranularity.ALL)
def _metadata_from_rows(
self, rows: Sequence[Tuple[Any, ...]]
) -> Optional[Metadata]:
"""Given SQL rows, return a metadata map (assuming that the last four columns
are the key, str_value, int_value & float_value)"""
add_attributes_to_current_span(
{
"num_rows": len(rows),
}
)
metadata: Dict[str, Union[str, int, float, bool]] = {}
for row in rows:
key = str(row[-5])
if row[-4] is not None:
metadata[key] = str(row[-4])
elif row[-3] is not None:
metadata[key] = int(row[-3])
elif row[-2] is not None:
metadata[key] = float(row[-2])
elif row[-1] is not None:
metadata[key] = bool(row[-1])
return metadata or None
@trace_method("SqlSysDB._insert_metadata", OpenTelemetryGranularity.ALL)
def _insert_metadata(
self,
cur: Cursor,
table: Table,
id_col: Column,
id: UUID,
metadata: UpdateMetadata,
clear_keys: Optional[Set[str]] = None,
) -> None:
# It would be cleaner to use something like ON CONFLICT UPDATE here But that is
# very difficult to do in a portable way (e.g sqlite and postgres have
# completely different sytnax)
add_attributes_to_current_span(
{
"num_keys": len(metadata),
}
)
if clear_keys:
q = (
self.querybuilder()
.from_(table)
.where(id_col == ParameterValue(self.uuid_to_db(id)))
.where(table.key.isin([ParameterValue(k) for k in clear_keys]))
.delete()
)
sql, params = get_sql(q, self.parameter_format())
cur.execute(sql, params)
q = (
self.querybuilder()
.into(table)
.columns(
id_col,
table.key,
table.str_value,
table.int_value,
table.float_value,
table.bool_value,
)
)
sql_id = self.uuid_to_db(id)
for k, v in metadata.items():
# Note: The order is important here because isinstance(v, bool)
# and isinstance(v, int) both are true for v of bool type.
if isinstance(v, bool):
q = q.insert(
ParameterValue(sql_id),
ParameterValue(k),
None,
None,
None,
ParameterValue(int(v)),
)
elif isinstance(v, str):
q = q.insert(
ParameterValue(sql_id),
ParameterValue(k),
ParameterValue(v),
None,
None,
None,
)
elif isinstance(v, int):
q = q.insert(
ParameterValue(sql_id),
ParameterValue(k),
None,
ParameterValue(v),
None,
None,
)
elif isinstance(v, float):
q = q.insert(
ParameterValue(sql_id),
ParameterValue(k),
None,
None,
ParameterValue(v),
None,
)
elif v is None:
continue
sql, params = get_sql(q, self.parameter_format())
if sql:
cur.execute(sql, params)
def _load_config_from_json_str_and_migrate(
self, collection_id: str, json_str: str
) -> CollectionConfigurationInternal:
try:
config_json = json.loads(json_str)
except json.JSONDecodeError:
raise ValueError(
f"Unable to decode configuration from JSON string: {json_str}"
)
try:
return CollectionConfigurationInternal.from_json_str(json_str)
except InvalidConfigurationError as error:
# 07/17/2024: the initial migration from the legacy metadata-based config to the new sysdb-based config had a bug where the batch_size and sync_threshold were swapped. Along with this migration, a validator was added to HNSWConfigurationInternal to ensure that batch_size <= sync_threshold.
hnsw_configuration = config_json.get("hnsw_configuration")
if hnsw_configuration:
batch_size = hnsw_configuration.get("batch_size")
sync_threshold = hnsw_configuration.get("sync_threshold")
if batch_size and sync_threshold and batch_size > sync_threshold:
# Allow new defaults to be set
hnsw_configuration = {
k: v
for k, v in hnsw_configuration.items()
if k not in ["batch_size", "sync_threshold"]
}
config_json.update({"hnsw_configuration": hnsw_configuration})
configuration = CollectionConfigurationInternal.from_json(
config_json
)
collections_t = Table("collections")
q = (
self.querybuilder()
.update(collections_t)
.set(
collections_t.config_json_str,
ParameterValue(configuration.to_json_str()),
)
.where(collections_t.id == ParameterValue(collection_id))
)
sql, params = get_sql(q, self.parameter_format())
with self.tx() as cur:
cur.execute(sql, params)
return configuration
raise error
def _insert_config_from_legacy_params(
self, collection_id: Any, metadata: Optional[Metadata]
) -> CollectionConfigurationInternal:
"""Insert the configuration from legacy metadata params into the collections table, and return the configuration object."""
# This is a legacy case where we don't have configuration stored in the database
# This is non-destructive, we don't delete or overwrite any keys in the metadata
from chromadb.segment.impl.vector.hnsw_params import PersistentHnswParams
collections_t = Table("collections")
# Get any existing HNSW params from the metadata (works regardless whether metadata has persistent params)
hnsw_metadata_params = PersistentHnswParams.extract(metadata or {})
hnsw_configuration = HNSWConfigurationInternal.from_legacy_params(
hnsw_metadata_params # type: ignore[arg-type]
)
configuration = CollectionConfigurationInternal(
parameters=[
ConfigurationParameter(
name="hnsw_configuration", value=hnsw_configuration
)
]
)
# Write the configuration into the database
configuration_json_str = configuration.to_json_str()
q = (
self.querybuilder()
.update(collections_t)
.set(
collections_t.config_json_str,
ParameterValue(configuration_json_str),
)
.where(collections_t.id == ParameterValue(collection_id))
)
sql, params = get_sql(q, self.parameter_format())
with self.tx() as cur:
cur.execute(sql, params)
return configuration