From e71b78004dad0a4849b27ae18f3fc9a4048a6c97 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=AD=90=EF=B8=8FNINIKA=E2=AD=90=EF=B8=8F?= Date: Wed, 28 Aug 2024 10:31:30 +0300 Subject: [PATCH] feat(queries): expose the remaining number of items in rust client API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: ⭐️NINIKA⭐️ --- client/src/query.rs | 12 ++++---- client/tests/integration/pagination.rs | 30 ++++++++++++++++++- client_cli/src/main.rs | 4 +-- core/src/query/store.rs | 4 +-- data_model/src/query/builder.rs | 30 ++++++++++++++----- data_model/src/query/mod.rs | 4 +-- smart_contract/src/lib.rs | 20 +++++++++---- .../query_assets_and_save_cursor/src/lib.rs | 2 +- 8 files changed, 79 insertions(+), 27 deletions(-) diff --git a/client/src/query.rs b/client/src/query.rs index 532d685aa8e..c676b31e609 100644 --- a/client/src/query.rs +++ b/client/src/query.rs @@ -159,7 +159,7 @@ impl QueryExecutor for Client { fn start_query( &self, query: QueryWithParams, - ) -> Result<(QueryOutputBatchBox, Option), Self::Error> { + ) -> Result<(QueryOutputBatchBox, u64, Option), Self::Error> { let request_head = self.get_query_request_head(); let request = QueryRequest::Start(query); @@ -167,19 +167,19 @@ impl QueryExecutor for Client { let response = request_head.assemble(request).build()?.send()?; let response = decode_iterable_query_response(&response)?; - let (batch, cursor) = response.into_parts(); + let (batch, remaining_items, cursor) = response.into_parts(); let cursor = cursor.map(|cursor| QueryCursor { request_head, cursor, }); - Ok((batch, cursor)) + Ok((batch, remaining_items, cursor)) } fn continue_query( cursor: Self::Cursor, - ) -> Result<(QueryOutputBatchBox, Option), Self::Error> { + ) -> Result<(QueryOutputBatchBox, u64, Option), Self::Error> { let QueryCursor { request_head, cursor, @@ -190,14 +190,14 @@ impl QueryExecutor for Client { let response = request_head.assemble(request).build()?.send()?; let response = decode_iterable_query_response(&response)?; - let (batch, cursor) = response.into_parts(); + let (batch, remaining_items, cursor) = response.into_parts(); let cursor = cursor.map(|cursor| QueryCursor { request_head, cursor, }); - Ok((batch, cursor)) + Ok((batch, remaining_items, cursor)) } } diff --git a/client/tests/integration/pagination.rs b/client/tests/integration/pagination.rs index b741a4f11ac..2431f5abb42 100644 --- a/client/tests/integration/pagination.rs +++ b/client/tests/integration/pagination.rs @@ -24,6 +24,33 @@ fn limits_should_work() -> Result<()> { Ok(()) } +#[test] +fn reported_length_should_be_accurate() -> Result<()> { + let (_rt, _peer, client) = ::new().with_port(11_170).start_with_runtime(); + wait_for_genesis_committed(&vec![client.clone()], 0); + + register_assets(&client)?; + + let mut iter = client + .query(asset::all_definitions()) + .with_pagination(Pagination { + limit: Some(nonzero!(7_u64)), + offset: 1, + }) + .with_fetch_size(FetchSize::new(Some(nonzero!(3_u64)))) + .execute()?; + + assert_eq!(iter.len(), 7); + + for _ in 0..4 { + iter.next().unwrap().unwrap(); + } + + assert_eq!(iter.len(), 3); + + Ok(()) +} + #[test] fn fetch_size_should_work() -> Result<()> { // use the lower-level API to inspect the batch size @@ -49,9 +76,10 @@ fn fetch_size_should_work() -> Result<()> { FetchSize::new(Some(nonzero!(3_u64))), ), ); - let (first_batch, _continue_cursor) = client.start_query(query)?; + let (first_batch, remaining_items, _continue_cursor) = client.start_query(query)?; assert_eq!(first_batch.len(), 3); + assert_eq!(remaining_items, 4); Ok(()) } diff --git a/client_cli/src/main.rs b/client_cli/src/main.rs index 18f2a491cea..d7c4edd1c87 100644 --- a/client_cli/src/main.rs +++ b/client_cli/src/main.rs @@ -1176,11 +1176,11 @@ mod json { // we can't really do type-erased iterable queries in a nice way right now... use iroha::data_model::query::builder::QueryExecutor; - let (mut first_batch, mut continue_cursor) = + let (mut first_batch, _remaining_items, mut continue_cursor) = client.start_query(query)?; while let Some(cursor) = continue_cursor { - let (next_batch, next_continue_cursor) = + let (next_batch, _remaining_items, next_continue_cursor) = ::continue_query(cursor)?; first_batch.extend(next_batch); diff --git a/core/src/query/store.rs b/core/src/query/store.rs index 58203304e16..ca07126ac7e 100644 --- a/core/src/query/store.rs +++ b/core/src/query/store.rs @@ -347,7 +347,7 @@ mod tests { ) .unwrap(); - let (batch, mut current_cursor) = query_store_handle + let (batch, _remaining_items, mut current_cursor) = query_store_handle .handle_iter_start(query_output, &ALICE_ID) .unwrap() .into_parts(); @@ -359,7 +359,7 @@ mod tests { let Ok(batched) = query_store_handle.handle_iter_continue(cursor) else { break; }; - let (batch, cursor) = batched.into_parts(); + let (batch, _remaining_items, cursor) = batched.into_parts(); counter += batch.len(); current_cursor = cursor; diff --git a/data_model/src/query/builder.rs b/data_model/src/query/builder.rs index c704e818629..bafe5e1df1d 100644 --- a/data_model/src/query/builder.rs +++ b/data_model/src/query/builder.rs @@ -34,7 +34,7 @@ pub trait QueryExecutor { query: SingularQueryBox, ) -> Result; - /// Starts an iterable query and returns the first batch of results. + /// Starts an iterable query and returns the first batch of results, the remaining number of results and a cursor to continue the query. /// /// # Errors /// @@ -42,22 +42,23 @@ pub trait QueryExecutor { fn start_query( &self, query: QueryWithParams, - ) -> Result<(QueryOutputBatchBox, Option), Self::Error>; + ) -> Result<(QueryOutputBatchBox, u64, Option), Self::Error>; - /// Continues an iterable query from the given cursor and returns the next batch of results. + /// Continues an iterable query from the given cursor and returns the next batch of results, the remaining number of results and a cursor to continue the query. /// /// # Errors /// /// Returns an error if the query execution fails. fn continue_query( cursor: Self::Cursor, - ) -> Result<(QueryOutputBatchBox, Option), Self::Error>; + ) -> Result<(QueryOutputBatchBox, u64, Option), Self::Error>; } /// An iterator over results of an iterable query. #[derive(Debug)] pub struct QueryIterator { current_batch_iter: vec::IntoIter, + remaining_items: u64, continue_cursor: Option, } @@ -73,12 +74,14 @@ where /// Returns an error if the type of the batch does not match the expected type `T`. pub fn new( first_batch: QueryOutputBatchBox, + remaining_items: u64, continue_cursor: Option, ) -> Result as TryFrom>::Error> { let batch: Vec = first_batch.try_into()?; Ok(Self { current_batch_iter: batch.into_iter(), + remaining_items, continue_cursor, }) } @@ -102,7 +105,7 @@ where let cursor = self.continue_cursor.take()?; // get a next batch from iroha - let (batch, cursor) = match E::continue_query(cursor) { + let (batch, remaining_items, cursor) = match E::continue_query(cursor) { Ok(r) => r, Err(e) => return Some(Err(e)), }; @@ -115,11 +118,23 @@ where .expect("BUG: iroha returned unexpected type in iterable query"); self.current_batch_iter = batch.into_iter(); + self.remaining_items = remaining_items; self.next() } } +impl ExactSizeIterator for QueryIterator +where + E: QueryExecutor, + Vec: TryFrom, + as TryFrom>::Error: core::fmt::Debug, +{ + fn len(&self) -> usize { + self.current_batch_iter.len() + self.remaining_items as usize + } +} + /// An error that can occur when constraining the number of results of an iterable query to one. #[derive( Debug, @@ -256,9 +271,10 @@ where }, }; - let (first_batch, continue_cursor) = self.query_executor.start_query(query)?; + let (first_batch, remaining_items, continue_cursor) = + self.query_executor.start_query(query)?; - let iterator = QueryIterator::::new(first_batch, continue_cursor) + let iterator = QueryIterator::::new(first_batch, remaining_items, continue_cursor) .expect( "INTERNAL BUG: iroha returned unexpected type in iterable query. Is there a schema mismatch?", ); diff --git a/data_model/src/query/mod.rs b/data_model/src/query/mod.rs index 3aad85bea61..ae3776d84b5 100644 --- a/data_model/src/query/mod.rs +++ b/data_model/src/query/mod.rs @@ -323,8 +323,8 @@ impl QueryOutput { } /// Split this [`QueryOutput`] into its constituent parts. - pub fn into_parts(self) -> (QueryOutputBatchBox, Option) { - (self.batch, self.continue_cursor) + pub fn into_parts(self) -> (QueryOutputBatchBox, u64, Option) { + (self.batch, self.remaining_items, self.continue_cursor) } } diff --git a/smart_contract/src/lib.rs b/smart_contract/src/lib.rs index 52c3556bb6f..f0fc3d7ba6f 100644 --- a/smart_contract/src/lib.rs +++ b/smart_contract/src/lib.rs @@ -149,28 +149,36 @@ impl QueryExecutor for SmartContractQueryExecutor { fn start_query( &self, query: QueryWithParams, - ) -> Result<(QueryOutputBatchBox, Option), Self::Error> { + ) -> Result<(QueryOutputBatchBox, u64, Option), Self::Error> { let QueryResponse::Iterable(output) = execute_query(&QueryRequest::Start(query))? else { dbg_panic("BUG: iroha returned unexpected type in iterable query"); }; - let (batch, cursor) = output.into_parts(); + let (batch, remaining_items, cursor) = output.into_parts(); - Ok((batch, cursor.map(|cursor| QueryCursor { cursor }))) + Ok(( + batch, + remaining_items, + cursor.map(|cursor| QueryCursor { cursor }), + )) } fn continue_query( cursor: Self::Cursor, - ) -> Result<(QueryOutputBatchBox, Option), Self::Error> { + ) -> Result<(QueryOutputBatchBox, u64, Option), Self::Error> { let QueryResponse::Iterable(output) = execute_query(&QueryRequest::Continue(cursor.cursor))? else { dbg_panic("BUG: iroha returned unexpected type in iterable query"); }; - let (batch, cursor) = output.into_parts(); + let (batch, remaining_items, cursor) = output.into_parts(); - Ok((batch, cursor.map(|cursor| QueryCursor { cursor }))) + Ok(( + batch, + remaining_items, + cursor.map(|cursor| QueryCursor { cursor }), + )) } } diff --git a/wasm_samples/query_assets_and_save_cursor/src/lib.rs b/wasm_samples/query_assets_and_save_cursor/src/lib.rs index 8c719530978..e4a03191a6e 100644 --- a/wasm_samples/query_assets_and_save_cursor/src/lib.rs +++ b/wasm_samples/query_assets_and_save_cursor/src/lib.rs @@ -35,7 +35,7 @@ fn main(owner: AccountId) { pub cursor: ForwardCursor, } - let (_batch, cursor) = SmartContractQueryExecutor + let (_batch, _remaining_items, cursor) = SmartContractQueryExecutor .start_query(QueryWithParams::new( QueryWithFilter::new(FindAssets, CompoundPredicate::PASS).into(), QueryParams::new(