Skip to content

Commit c7a8fa3

Browse files
authored
[ENH]: Metering for indexing_status (#6128)
## Description of changes As titled, - Improvements & Bug fixes - ... - New functionality - ... ## Test plan _How are these changes tested?_ - [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust ## Migration plan _Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?_ ## Observability plan _What is the plan to instrument and monitor this change?_ ## Documentation Changes _Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the_ [_docs section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_
1 parent 6af1747 commit c7a8fa3

File tree

2 files changed

+27
-3
lines changed

2 files changed

+27
-3
lines changed

rust/frontend/src/impls/service_based_frontend.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1448,6 +1448,11 @@ impl ServiceBasedFrontend {
14481448
num_indexed_ops as f32 / total_ops as f32
14491449
};
14501450

1451+
chroma_metering::with_current(|context| {
1452+
context.latest_collection_logical_size_bytes(0);
1453+
context.finish_request(Instant::now());
1454+
});
1455+
14511456
Ok(IndexStatusResponse {
14521457
op_indexing_progress,
14531458
num_unindexed_ops,

rust/frontend/src/server.rs

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1822,11 +1822,10 @@ async fn indexing_status(
18221822
format!("collection:{}", collection_id).as_str(),
18231823
])?;
18241824

1825-
// TODO(tanujnay112): Add authorization for indexing status
18261825
server
18271826
.authenticate_and_authorize_collection(
18281827
&headers,
1829-
AuthzAction::GetCollection,
1828+
AuthzAction::Count,
18301829
AuthzResource {
18311830
tenant: Some(tenant.clone()),
18321831
database: Some(database.clone()),
@@ -1836,10 +1835,30 @@ async fn indexing_status(
18361835
)
18371836
.await?;
18381837

1838+
let metering_context_container =
1839+
chroma_metering::create::<CollectionReadContext>(CollectionReadContext::new(
1840+
tenant.clone(),
1841+
database.clone(),
1842+
collection_id.clone(),
1843+
ReadAction::Query,
1844+
));
1845+
1846+
metering_context_container.enter();
1847+
1848+
chroma_metering::with_current(|context| {
1849+
context.start_request(Instant::now());
1850+
});
1851+
18391852
let collection_id =
18401853
CollectionUuid::from_str(&collection_id).map_err(|_| ValidationError::CollectionId)?;
18411854

1842-
Ok(Json(server.frontend.indexing_status(collection_id).await?))
1855+
Ok(Json(
1856+
server
1857+
.frontend
1858+
.indexing_status(collection_id)
1859+
.meter(metering_context_container)
1860+
.await?,
1861+
))
18431862
}
18441863

18451864
/// Retrieves records from a collection by ID or metadata filter.

0 commit comments

Comments
 (0)