Line data    Source code

       1              : use anyhow::Context;
       2              : use async_stream::{stream, try_stream};
       3              : use aws_sdk_s3::Client;
       4              : use tokio_stream::Stream;
       5              : 
       6              : use crate::{list_objects_with_retries, RootTarget, TenantId};
       7              : use utils::id::{TenantTimelineId, TimelineId};
       8              : 
       9              : /// Given an S3 bucket, output a stream of TenantIds discovered via ListObjectsv2
      10            0 : pub fn stream_tenants<'a>(
      11            0 :     s3_client: &'a Client,
      12            0 :     target: &'a RootTarget,
      13            0 : ) -> impl Stream<Item = anyhow::Result<TenantId>> + 'a {
      14            0 :     try_stream! {
      15            0 :         let mut continuation_token = None;
      16            0 :         loop {
      17            0 :             let tenants_target = target.tenants_root();
      18            0 :             let fetch_response =
      19            0 :                 list_objects_with_retries(s3_client, tenants_target, continuation_token.clone()).await?;
      20            0 : 
      21            0 :             let new_entry_ids = fetch_response
      22            0 :                 .common_prefixes()
      23            0 :                 .unwrap_or_default()
      24            0 :                 .iter()
      25            0 :                 .filter_map(|prefix| prefix.prefix())
      26            0 :                 .filter_map(|prefix| -> Option<&str> {
      27            0 :                     prefix
      28            0 :                         .strip_prefix(&tenants_target.prefix_in_bucket)?
      29            0 :                         .strip_suffix('/')
      30            0 :                 }).map(|entry_id_str| {
      31            0 :                 entry_id_str
      32            0 :                     .parse()
      33            0 :                     .with_context(|| format!("Incorrect entry id str: {entry_id_str}"))
      34            0 :             });
      35            0 : 
      36            0 :             for i in new_entry_ids {
      37            0 :                 yield i?;
      38            0 :             }
      39            0 : 
      40            0 :             match fetch_response.next_continuation_token {
      41            0 :                 Some(new_token) => continuation_token = Some(new_token),
      42            0 :                 None => break,
      43            0 :             }
      44            0 :         }
      45            0 :     }
      46            0 : }
      47              : 
      48              : /// Given a TenantId, output a stream of the timelines within that tenant, discovered
      49              : /// using ListObjectsv2.  The listing is done before the stream is built, so that this
      50              : /// function can be used to generate concurrency on a stream using buffer_unordered.
      51            0 : pub async fn stream_tenant_timelines<'a>(
      52            0 :     s3_client: &'a Client,
      53            0 :     target: &'a RootTarget,
      54            0 :     tenant: TenantId,
      55            0 : ) -> anyhow::Result<impl Stream<Item = Result<TenantTimelineId, anyhow::Error>> + 'a> {
      56            0 :     let mut timeline_ids: Vec<Result<TimelineId, anyhow::Error>> = Vec::new();
      57            0 :     let mut continuation_token = None;
      58            0 :     let timelines_target = target.timelines_root(&tenant);
      59              : 
      60              :     loop {
      61            0 :         tracing::info!("Listing in {}", tenant);
      62            0 :         let fetch_response =
      63            0 :             list_objects_with_retries(s3_client, &timelines_target, continuation_token.clone())
      64            0 :                 .await;
      65            0 :         let fetch_response = match fetch_response {
      66            0 :             Err(e) => {
      67            0 :                 timeline_ids.push(Err(e));
      68            0 :                 break;
      69              :             }
      70            0 :             Ok(r) => r,
      71            0 :         };
      72            0 : 
      73            0 :         let new_entry_ids = fetch_response
      74            0 :             .common_prefixes()
      75            0 :             .unwrap_or_default()
      76            0 :             .iter()
      77            0 :             .filter_map(|prefix| prefix.prefix())
      78            0 :             .filter_map(|prefix| -> Option<&str> {
      79            0 :                 prefix
      80            0 :                     .strip_prefix(&timelines_target.prefix_in_bucket)?
      81            0 :                     .strip_suffix('/')
      82            0 :             })
      83            0 :             .map(|entry_id_str| {
      84            0 :                 entry_id_str
      85            0 :                     .parse::<TimelineId>()
      86            0 :                     .with_context(|| format!("Incorrect entry id str: {entry_id_str}"))
      87            0 :             });
      88              : 
      89            0 :         for i in new_entry_ids {
      90            0 :             timeline_ids.push(i);
      91            0 :         }
      92              : 
      93            0 :         match fetch_response.next_continuation_token {
      94            0 :             Some(new_token) => continuation_token = Some(new_token),
      95            0 :             None => break,
      96              :         }
      97              :     }
      98              : 
      99            0 :     tracing::info!("Yielding for {}", tenant);
     100            0 :     Ok(stream! {
     101            0 :         for i in timeline_ids {
     102            0 :             let id = i?;
     103            0 :             yield Ok(TenantTimelineId::new(tenant, id));
     104              :         }
     105            0 :     })
     106            0 : }