Skip to content

Commit c7c257f

Browse files
authored
[ENH]: enforce maximum get_collections limit as 100 (#4929)
## Description of changes Enforces maximum get collections limit as 100 as a temporary patch. ## Test plan - [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust ## Documentation Changes We should update docs somewhere saying this is enforced.
1 parent e66ea23 commit c7c257f

File tree

2 files changed

+20
-2
lines changed

2 files changed

+20
-2
lines changed

rust/sysdb/src/sysdb.rs

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -139,10 +139,25 @@ impl SysDb {
139139

140140
pub async fn get_collections(
141141
&mut self,
142-
options: GetCollectionsOptions,
142+
mut options: GetCollectionsOptions,
143143
) -> Result<Vec<Collection>, GetCollectionsError> {
144144
match self {
145-
SysDb::Grpc(grpc) => grpc.get_collections(options).await,
145+
SysDb::Grpc(grpc) => {
146+
// TODO(c-gamble): Move this to config
147+
let max_get_collections_limit = 100u32;
148+
match options.limit {
149+
Some(limit) => {
150+
if limit > max_get_collections_limit {
151+
return Err(GetCollectionsError::MaximumLimitExceeded(
152+
limit,
153+
max_get_collections_limit,
154+
));
155+
}
156+
}
157+
None => options.limit = Some(max_get_collections_limit),
158+
}
159+
grpc.get_collections(options).await
160+
}
146161
SysDb::Sqlite(sqlite) => sqlite.get_collections(options).await,
147162
SysDb::Test(test) => test.get_collections(options).await,
148163
}

rust/types/src/api_types.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -694,6 +694,8 @@ pub enum GetCollectionsError {
694694
CollectionId(#[from] uuid::Error),
695695
#[error("Could not deserialize database ID")]
696696
DatabaseId,
697+
#[error("Provided limit `{0}` exceeds maximum allowable limit `{1}`")]
698+
MaximumLimitExceeded(u32, u32),
697699
}
698700

699701
impl ChromaError for GetCollectionsError {
@@ -703,6 +705,7 @@ impl ChromaError for GetCollectionsError {
703705
GetCollectionsError::Configuration(_) => ErrorCodes::Internal,
704706
GetCollectionsError::CollectionId(_) => ErrorCodes::Internal,
705707
GetCollectionsError::DatabaseId => ErrorCodes::Internal,
708+
GetCollectionsError::MaximumLimitExceeded(_, _) => ErrorCodes::InvalidArgument,
706709
}
707710
}
708711
}

0 commit comments

Comments
 (0)