diff --git a/src/drives/list.rs b/src/drives/list.rs index 5c303db..513747e 100644 --- a/src/drives/list.rs +++ b/src/drives/list.rs @@ -11,13 +11,14 @@ use std::io; pub struct Config { pub skip_header: bool, pub field_separator: String, + pub max: Option, } pub async fn list(config: Config) -> Result<(), Error> { let hub = hub_helper::get_hub().await.map_err(Error::Hub)?; let delegate_config = UploadDelegateConfig::default(); - let drives = list_drives(&hub, delegate_config) + let drives = list_drives(&hub, delegate_config, config.max) .await .map_err(Error::ListDrives)?; @@ -55,18 +56,51 @@ fn print_drives_table(config: &Config, drives: Vec) { pub async fn list_drives( hub: &Hub, delegate_config: UploadDelegateConfig, + max: Option, ) -> Result, google_drive3::Error> { let mut delegate = UploadDelegate::new(delegate_config); + let mut all_drives = Vec::new(); + let mut page_token: Option = None; + + // Default to 10 if the user does not provide a max + let target_max = max.unwrap_or(10); + // API limit is 100; keep page size within target_max and the API cap + let page_size = target_max.min(100); - let (_, drives_list) = hub - .drives() - .list() - .add_scope(google_drive3::api::Scope::Full) - .delegate(&mut delegate) - .doit() - .await?; + loop { + let mut request = hub + .drives() + .list() + .page_size(page_size) + .add_scope(google_drive3::api::Scope::Full); + + if let Some(token) = page_token { + request = request.page_token(&token); + } + + let (_, drives_list) = request + .delegate(&mut delegate) + .doit() + .await?; + + if let Some(mut drives) = drives_list.drives { + all_drives.append(&mut drives); + } + + // Check if we've reached the target max or there are no more pages + if all_drives.len() >= target_max as usize || drives_list.next_page_token.is_none() { + break; + } + + page_token = drives_list.next_page_token; + } + + // Truncate to exact max if specified + if let Some(max_count) = max { + all_drives.truncate(max_count as usize); + } - Ok(drives_list.drives.unwrap_or_default()) + Ok(all_drives) } #[derive(Debug)] diff --git a/src/main.rs b/src/main.rs index b309066..4f29123 100644 --- a/src/main.rs +++ b/src/main.rs @@ -104,6 +104,10 @@ enum DriveCommand { /// Field separator #[arg(long, default_value_t = String::from("\t"))] field_separator: String, + + /// Max results to return + #[arg(long)] + max: Option, }, } @@ -438,9 +442,11 @@ async fn main() { DriveCommand::List { skip_header, field_separator, + max, } => drives::list(drives::list::Config { skip_header, field_separator, + max, }) .await .unwrap_or_else(handle_error),