Line data    Source code

       1              : use anyhow::Context;
       2              : use async_stream::{stream, try_stream};
       3              : use aws_sdk_s3::{types::ObjectIdentifier, Client};
       4              : use tokio_stream::Stream;
       5              : 
       6              : use crate::{list_objects_with_retries, RootTarget, S3Target, TenantShardTimelineId};
       7              : use pageserver_api::shard::TenantShardId;
       8              : use utils::id::TimelineId;
       9              : 
      10              : /// Given an S3 bucket, output a stream of TenantIds discovered via ListObjectsv2
      11            0 : pub fn stream_tenants<'a>(
      12            0 :     s3_client: &'a Client,
      13            0 :     target: &'a RootTarget,
      14            0 : ) -> impl Stream<Item = anyhow::Result<TenantShardId>> + 'a {
      15            0 :     try_stream! {
      16            0 :         let mut continuation_token = None;
      17            0 :         let tenants_target = target.tenants_root();
      18              :         loop {
      19            0 :             let fetch_response =
      20            0 :                 list_objects_with_retries(s3_client, &tenants_target, continuation_token.clone()).await?;
      21              : 
      22            0 :             let new_entry_ids = fetch_response
      23            0 :                 .common_prefixes()
      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              : 
      36            0 :             for i in new_entry_ids {
      37            0 :                 yield i?;
      38              :             }
      39              : 
      40            0 :             match fetch_response.next_continuation_token {
      41            0 :                 Some(new_token) => continuation_token = Some(new_token),
      42            0 :                 None => break,
      43              :             }
      44              :         }
      45              :     }
      46            0 : }
      47              : 
      48              : /// Given a TenantShardId, 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: TenantShardId,
      55            0 : ) -> anyhow::Result<impl Stream<Item = Result<TenantShardTimelineId, 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 :             .iter()
      76            0 :             .filter_map(|prefix| prefix.prefix())
      77            0 :             .filter_map(|prefix| -> Option<&str> {
      78            0 :                 prefix
      79            0 :                     .strip_prefix(&timelines_target.prefix_in_bucket)?
      80            0 :                     .strip_suffix('/')
      81            0 :             })
      82            0 :             .map(|entry_id_str| {
      83            0 :                 entry_id_str
      84            0 :                     .parse::<TimelineId>()
      85            0 :                     .with_context(|| format!("Incorrect entry id str: {entry_id_str}"))
      86            0 :             });
      87              : 
      88            0 :         for i in new_entry_ids {
      89            0 :             timeline_ids.push(i);
      90            0 :         }
      91              : 
      92            0 :         match fetch_response.next_continuation_token {
      93            0 :             Some(new_token) => continuation_token = Some(new_token),
      94            0 :             None => break,
      95              :         }
      96              :     }
      97              : 
      98            0 :     tracing::info!("Yielding for {}", tenant);
      99            0 :     Ok(stream! {
     100            0 :         for i in timeline_ids {
     101            0 :             let id = i?;
     102            0 :             yield Ok(TenantShardTimelineId::new(tenant, id));
     103              :         }
     104            0 :     })
     105            0 : }
     106              : 
     107            0 : pub(crate) fn stream_listing<'a>(
     108            0 :     s3_client: &'a Client,
     109            0 :     target: &'a S3Target,
     110            0 : ) -> impl Stream<Item = anyhow::Result<ObjectIdentifier>> + 'a {
     111            0 :     try_stream! {
     112            0 :         let mut continuation_token = None;
     113              :         loop {
     114            0 :             let fetch_response =
     115            0 :                 list_objects_with_retries(s3_client, target, continuation_token.clone()).await?;
     116              : 
     117            0 :             if target.delimiter.is_empty() {
     118            0 :                 for object_key in fetch_response.contents().iter().filter_map(|object| object.key())
     119              :                 {
     120            0 :                     let object_id = ObjectIdentifier::builder().key(object_key).build()?;
     121            0 :                     yield object_id;
     122              :                 }
     123              :             } else {
     124            0 :                 for prefix in fetch_response.common_prefixes().iter().filter_map(|p| p.prefix()) {
     125            0 :                     let object_id = ObjectIdentifier::builder().key(prefix).build()?;
     126            0 :                     yield object_id;
     127              :                 }
     128              :             }
     129              : 
     130            0 :             match fetch_response.next_continuation_token {
     131            0 :                 Some(new_token) => continuation_token = Some(new_token),
     132            0 :                 None => break,
     133              :             }
     134              :         }
     135              :     }
     136            0 : }