modify get_protobuf_or_bincode_cells to accept and return an iterator

This commit is contained in:
Edgar Xi
2022-03-21 18:48:51 -04:00
committed by Trent Nelson
parent fbcf6a0802
commit f717fda9a3
2 changed files with 20 additions and 20 deletions

View File

@ -680,17 +680,17 @@ impl<F: FnMut(Request<()>) -> InterceptedRequestResult> BigTable<F> {
deserialize_protobuf_or_bincode_cell_data(&row_data, table, key) deserialize_protobuf_or_bincode_cell_data(&row_data, table, key)
} }
pub async fn get_protobuf_or_bincode_cells<B, P>( pub async fn get_protobuf_or_bincode_cells<'a, B, P>(
&mut self, &mut self,
table: &str, table: &'a str,
row_keys: &[RowKey], row_keys: impl Iterator<Item = String>,
) -> Result<Vec<(RowKey, CellData<B, P>)>> ) -> Result<impl Iterator<Item = (RowKey, CellData<B, P>)> + 'a>
where where
B: serde::de::DeserializeOwned, B: serde::de::DeserializeOwned,
P: prost::Message + Default, P: prost::Message + Default,
{ {
Ok(self Ok(self
.get_multi_row_data(table, row_keys) .get_multi_row_data(table, row_keys.collect::<Vec<RowKey>>().as_slice())
.await? .await?
.into_iter() .into_iter()
.map(|(key, row_data)| { .map(|(key, row_data)| {
@ -699,8 +699,7 @@ impl<F: FnMut(Request<()>) -> InterceptedRequestResult> BigTable<F> {
key, key,
deserialize_protobuf_or_bincode_cell_data(&row_data, table, key_str).unwrap(), deserialize_protobuf_or_bincode_cell_data(&row_data, table, key_str).unwrap(),
) )
}) }))
.collect())
} }
pub async fn put_bincode_cells<T>( pub async fn put_bincode_cells<T>(

View File

@ -466,21 +466,22 @@ impl LedgerStorage {
); );
inc_new_counter_debug!("storage-bigtable-query", 1); inc_new_counter_debug!("storage-bigtable-query", 1);
let mut bigtable = self.connection.client(); let mut bigtable = self.connection.client();
let row_keys: Vec<RowKey> = slots.iter().copied().map(slot_to_blocks_key).collect(); let row_keys = slots.iter().copied().map(slot_to_blocks_key);
let data: Vec<(Slot, ConfirmedBlock)> = bigtable let data: Vec<(Slot, ConfirmedBlock)> = bigtable
.get_protobuf_or_bincode_cells::<StoredConfirmedBlock, generated::ConfirmedBlock>( .get_protobuf_or_bincode_cells("blocks", row_keys)
"blocks",
row_keys.as_slice(),
)
.await? .await?
.into_iter() .filter_map(
.filter_map(|(row_key, block_cell_data)| { |(row_key, block_cell_data): (
let block = match block_cell_data { RowKey,
bigtable::CellData::Bincode(block) => block.into(), bigtable::CellData<StoredConfirmedBlock, generated::ConfirmedBlock>,
bigtable::CellData::Protobuf(block) => block.try_into().ok()?, )| {
}; let block = match block_cell_data {
Some((key_to_slot(&row_key).unwrap(), block)) bigtable::CellData::Bincode(block) => block.into(),
}) bigtable::CellData::Protobuf(block) => block.try_into().ok()?,
};
Some((key_to_slot(&row_key).unwrap(), block))
},
)
.collect(); .collect();
Ok(data) Ok(data)
} }