Line data    Source code

       1              : use pageserver_api::models::{AuxFilePolicy, TenantConfig, TenantConfigRequest};
       2              : use pageserver_api::shard::TenantShardId;
       3              : use utils::id::TenantTimelineId;
       4              : use utils::lsn::Lsn;
       5              : 
       6              : use std::collections::HashMap;
       7              : use std::sync::Arc;
       8              : 
       9              : /// Ingest aux files into the pageserver.
      10            0 : #[derive(clap::Parser)]
      11              : pub(crate) struct Args {
      12              :     #[clap(long, default_value = "http://localhost:9898")]
      13            0 :     mgmt_api_endpoint: String,
      14              :     #[clap(long, default_value = "postgres://postgres@localhost:64000")]
      15            0 :     page_service_connstring: String,
      16              :     #[clap(long)]
      17              :     pageserver_jwt: Option<String>,
      18              : 
      19            0 :     targets: Option<Vec<TenantTimelineId>>,
      20              : }
      21              : 
      22            0 : pub(crate) fn main(args: Args) -> anyhow::Result<()> {
      23            0 :     let rt = tokio::runtime::Builder::new_multi_thread()
      24            0 :         .enable_all()
      25            0 :         .build()
      26            0 :         .unwrap();
      27            0 : 
      28            0 :     let main_task = rt.spawn(main_impl(args));
      29            0 :     rt.block_on(main_task).unwrap()
      30            0 : }
      31              : 
      32            0 : async fn main_impl(args: Args) -> anyhow::Result<()> {
      33            0 :     let args: &'static Args = Box::leak(Box::new(args));
      34            0 : 
      35            0 :     let mgmt_api_client = Arc::new(pageserver_client::mgmt_api::Client::new(
      36            0 :         args.mgmt_api_endpoint.clone(),
      37            0 :         args.pageserver_jwt.as_deref(),
      38            0 :     ));
      39              : 
      40              :     // discover targets
      41            0 :     let timelines: Vec<TenantTimelineId> = crate::util::cli::targets::discover(
      42            0 :         &mgmt_api_client,
      43            0 :         crate::util::cli::targets::Spec {
      44            0 :             limit_to_first_n_targets: None,
      45              :             targets: {
      46            0 :                 if let Some(targets) = &args.targets {
      47            0 :                     if targets.len() != 1 {
      48            0 :                         anyhow::bail!("must specify exactly one target");
      49            0 :                     }
      50            0 :                     Some(targets.clone())
      51              :                 } else {
      52            0 :                     None
      53              :                 }
      54              :             },
      55              :         },
      56              :     )
      57            0 :     .await?;
      58              : 
      59            0 :     let timeline = timelines[0];
      60            0 :     let tenant_shard_id = TenantShardId::unsharded(timeline.tenant_id);
      61            0 :     let timeline_id = timeline.timeline_id;
      62            0 : 
      63            0 :     println!("operating on timeline {}", timeline);
      64            0 : 
      65            0 :     mgmt_api_client
      66            0 :         .tenant_config(&TenantConfigRequest {
      67            0 :             tenant_id: timeline.tenant_id,
      68            0 :             config: TenantConfig {
      69            0 :                 switch_aux_file_policy: Some(AuxFilePolicy::V2),
      70            0 :                 ..Default::default()
      71            0 :             },
      72            0 :         })
      73            0 :         .await?;
      74              : 
      75            0 :     for batch in 0..100 {
      76            0 :         let items = (0..100)
      77            0 :             .map(|id| {
      78            0 :                 (
      79            0 :                     format!("pg_logical/mappings/{:03}.{:03}", batch, id),
      80            0 :                     format!("{:08}", id),
      81            0 :                 )
      82            0 :             })
      83            0 :             .collect::<HashMap<_, _>>();
      84            0 :         let file_cnt = items.len();
      85            0 :         mgmt_api_client
      86            0 :             .ingest_aux_files(tenant_shard_id, timeline_id, items)
      87            0 :             .await?;
      88            0 :         println!("ingested {file_cnt} files");
      89              :     }
      90              : 
      91            0 :     let files = mgmt_api_client
      92            0 :         .list_aux_files(tenant_shard_id, timeline_id, Lsn(Lsn::MAX.0 - 1))
      93            0 :         .await?;
      94              : 
      95            0 :     println!("{} files found", files.len());
      96            0 : 
      97            0 :     anyhow::Ok(())
      98            0 : }