LCOV - code coverage report
Current view: top level - control_plane/src/bin - neon_local.rs (source / functions) Coverage Total Hit
Test: 32f4a56327bc9da697706839ed4836b2a00a408f.info Lines: 80.0 % 1319 1055
Test Date: 2024-02-07 07:37:29 Functions: 68.9 % 90 62

            Line data    Source code
       1              : //!
       2              : //! `neon_local` is an executable that can be used to create a local
       3              : //! Neon environment, for testing purposes. The local environment is
       4              : //! quite different from the cloud environment with Kubernetes, but it
       5              : //! easier to work with locally. The python tests in `test_runner`
       6              : //! rely on `neon_local` to set up the environment for each test.
       7              : //!
       8              : use anyhow::{anyhow, bail, Context, Result};
       9              : use clap::{value_parser, Arg, ArgAction, ArgMatches, Command, ValueEnum};
      10              : use compute_api::spec::ComputeMode;
      11              : use control_plane::attachment_service::{
      12              :     AttachmentService, NodeAvailability, NodeConfigureRequest, NodeSchedulingPolicy,
      13              : };
      14              : use control_plane::endpoint::ComputeControlPlane;
      15              : use control_plane::local_env::{InitForceMode, LocalEnv};
      16              : use control_plane::pageserver::{PageServerNode, PAGESERVER_REMOTE_STORAGE_DIR};
      17              : use control_plane::safekeeper::SafekeeperNode;
      18              : use control_plane::{broker, local_env};
      19              : use pageserver_api::models::{
      20              :     ShardParameters, TenantCreateRequest, TimelineCreateRequest, TimelineInfo,
      21              : };
      22              : use pageserver_api::shard::{ShardCount, ShardStripeSize, TenantShardId};
      23              : use pageserver_api::{
      24              :     DEFAULT_HTTP_LISTEN_PORT as DEFAULT_PAGESERVER_HTTP_PORT,
      25              :     DEFAULT_PG_LISTEN_PORT as DEFAULT_PAGESERVER_PG_PORT,
      26              : };
      27              : use postgres_backend::AuthType;
      28              : use postgres_connection::parse_host_port;
      29              : use safekeeper_api::{
      30              :     DEFAULT_HTTP_LISTEN_PORT as DEFAULT_SAFEKEEPER_HTTP_PORT,
      31              :     DEFAULT_PG_LISTEN_PORT as DEFAULT_SAFEKEEPER_PG_PORT,
      32              : };
      33              : use std::collections::{BTreeSet, HashMap};
      34              : use std::path::PathBuf;
      35              : use std::process::exit;
      36              : use std::str::FromStr;
      37              : use storage_broker::DEFAULT_LISTEN_ADDR as DEFAULT_BROKER_ADDR;
      38              : use url::Host;
      39              : use utils::{
      40              :     auth::{Claims, Scope},
      41              :     id::{NodeId, TenantId, TenantTimelineId, TimelineId},
      42              :     lsn::Lsn,
      43              :     project_git_version,
      44              : };
      45              : 
      46              : // Default id of a safekeeper node, if not specified on the command line.
      47              : const DEFAULT_SAFEKEEPER_ID: NodeId = NodeId(1);
      48              : const DEFAULT_PAGESERVER_ID: NodeId = NodeId(1);
      49              : const DEFAULT_BRANCH_NAME: &str = "main";
      50              : project_git_version!(GIT_VERSION);
      51              : 
      52              : const DEFAULT_PG_VERSION: &str = "15";
      53              : 
      54              : const DEFAULT_PAGESERVER_CONTROL_PLANE_API: &str = "http://127.0.0.1:1234/upcall/v1/";
      55              : 
      56            0 : fn default_conf(num_pageservers: u16) -> String {
      57            0 :     let mut template = format!(
      58            0 :         r#"
      59            0 : # Default built-in configuration, defined in main.rs
      60            0 : control_plane_api = '{DEFAULT_PAGESERVER_CONTROL_PLANE_API}'
      61            0 : 
      62            0 : [broker]
      63            0 : listen_addr = '{DEFAULT_BROKER_ADDR}'
      64            0 : 
      65            0 : [[safekeepers]]
      66            0 : id = {DEFAULT_SAFEKEEPER_ID}
      67            0 : pg_port = {DEFAULT_SAFEKEEPER_PG_PORT}
      68            0 : http_port = {DEFAULT_SAFEKEEPER_HTTP_PORT}
      69            0 : 
      70            0 : "#,
      71            0 :     );
      72              : 
      73            0 :     for i in 0..num_pageservers {
      74            0 :         let pageserver_id = NodeId(DEFAULT_PAGESERVER_ID.0 + i as u64);
      75            0 :         let pg_port = DEFAULT_PAGESERVER_PG_PORT + i;
      76            0 :         let http_port = DEFAULT_PAGESERVER_HTTP_PORT + i;
      77            0 : 
      78            0 :         template += &format!(
      79            0 :             r#"
      80            0 : [[pageservers]]
      81            0 : id = {pageserver_id}
      82            0 : listen_pg_addr = '127.0.0.1:{pg_port}'
      83            0 : listen_http_addr = '127.0.0.1:{http_port}'
      84            0 : pg_auth_type = '{trust_auth}'
      85            0 : http_auth_type = '{trust_auth}'
      86            0 : "#,
      87            0 :             trust_auth = AuthType::Trust,
      88            0 :         )
      89              :     }
      90              : 
      91            0 :     template
      92            0 : }
      93              : 
      94              : ///
      95              : /// Timelines tree element used as a value in the HashMap.
      96              : ///
      97              : struct TimelineTreeEl {
      98              :     /// `TimelineInfo` received from the `pageserver` via the `timeline_list` http API call.
      99              :     pub info: TimelineInfo,
     100              :     /// Name, recovered from neon config mappings
     101              :     pub name: Option<String>,
     102              :     /// Holds all direct children of this timeline referenced using `timeline_id`.
     103              :     pub children: BTreeSet<TimelineId>,
     104              : }
     105              : 
     106              : // Main entry point for the 'neon_local' CLI utility
     107              : //
     108              : // This utility helps to manage neon installation. That includes following:
     109              : //   * Management of local postgres installations running on top of the
     110              : //     pageserver.
     111              : //   * Providing CLI api to the pageserver
     112              : //   * TODO: export/import to/from usual postgres
     113         6106 : fn main() -> Result<()> {
     114         6106 :     let matches = cli().get_matches();
     115              : 
     116         6106 :     let (sub_name, sub_args) = match matches.subcommand() {
     117         6106 :         Some(subcommand_data) => subcommand_data,
     118            0 :         None => bail!("no subcommand provided"),
     119              :     };
     120              : 
     121              :     // Check for 'neon init' command first.
     122         6106 :     let subcommand_result = if sub_name == "init" {
     123          354 :         handle_init(sub_args).map(Some)
     124              :     } else {
     125              :         // all other commands need an existing config
     126         5752 :         let mut env = LocalEnv::load_config().context("Error loading config")?;
     127         5752 :         let original_env = env.clone();
     128         5752 : 
     129         5752 :         let rt = tokio::runtime::Builder::new_current_thread()
     130         5752 :             .enable_all()
     131         5752 :             .build()
     132         5752 :             .unwrap();
     133              : 
     134         5752 :         let subcommand_result = match sub_name {
     135         5752 :             "tenant" => rt.block_on(handle_tenant(sub_args, &mut env)),
     136         5273 :             "timeline" => rt.block_on(handle_timeline(sub_args, &mut env)),
     137         4909 :             "start" => rt.block_on(handle_start_all(sub_args, &env)),
     138         4906 :             "stop" => rt.block_on(handle_stop_all(sub_args, &env)),
     139         4903 :             "pageserver" => rt.block_on(handle_pageserver(sub_args, &env)),
     140         3700 :             "attachment_service" => rt.block_on(handle_attachment_service(sub_args, &env)),
     141         2982 :             "safekeeper" => rt.block_on(handle_safekeeper(sub_args, &env)),
     142         1927 :             "endpoint" => rt.block_on(handle_endpoint(sub_args, &env)),
     143            1 :             "mappings" => handle_mappings(sub_args, &mut env),
     144            0 :             "pg" => bail!("'pg' subcommand has been renamed to 'endpoint'"),
     145            0 :             _ => bail!("unexpected subcommand {sub_name}"),
     146              :         };
     147              : 
     148         5752 :         if original_env != env {
     149          796 :             subcommand_result.map(|()| Some(env))
     150              :         } else {
     151         4956 :             subcommand_result.map(|()| None)
     152              :         }
     153              :     };
     154              : 
     155         6060 :     match subcommand_result {
     156         1150 :         Ok(Some(updated_env)) => updated_env.persist_config(&updated_env.base_data_dir)?,
     157         4910 :         Ok(None) => (),
     158           46 :         Err(e) => {
     159           46 :             eprintln!("command failed: {e:?}");
     160           46 :             exit(1);
     161              :         }
     162              :     }
     163         6060 :     Ok(())
     164         6060 : }
     165              : 
     166              : ///
     167              : /// Prints timelines list as a tree-like structure.
     168              : ///
     169           19 : fn print_timelines_tree(
     170           19 :     timelines: Vec<TimelineInfo>,
     171           19 :     mut timeline_name_mappings: HashMap<TenantTimelineId, String>,
     172           19 : ) -> Result<()> {
     173           19 :     let mut timelines_hash = timelines
     174           19 :         .iter()
     175           34 :         .map(|t| {
     176           34 :             (
     177           34 :                 t.timeline_id,
     178           34 :                 TimelineTreeEl {
     179           34 :                     info: t.clone(),
     180           34 :                     children: BTreeSet::new(),
     181           34 :                     name: timeline_name_mappings
     182           34 :                         .remove(&TenantTimelineId::new(t.tenant_id.tenant_id, t.timeline_id)),
     183           34 :                 },
     184           34 :             )
     185           34 :         })
     186           19 :         .collect::<HashMap<_, _>>();
     187              : 
     188              :     // Memorize all direct children of each timeline.
     189           34 :     for timeline in timelines.iter() {
     190           34 :         if let Some(ancestor_timeline_id) = timeline.ancestor_timeline_id {
     191           15 :             timelines_hash
     192           15 :                 .get_mut(&ancestor_timeline_id)
     193           15 :                 .context("missing timeline info in the HashMap")?
     194              :                 .children
     195           15 :                 .insert(timeline.timeline_id);
     196           19 :         }
     197              :     }
     198              : 
     199           34 :     for timeline in timelines_hash.values() {
     200              :         // Start with root local timelines (no ancestors) first.
     201           34 :         if timeline.info.ancestor_timeline_id.is_none() {
     202           19 :             print_timeline(0, &Vec::from([true]), timeline, &timelines_hash)?;
     203           15 :         }
     204              :     }
     205              : 
     206           19 :     Ok(())
     207           19 : }
     208              : 
     209              : ///
     210              : /// Recursively prints timeline info with all its children.
     211              : ///
     212           34 : fn print_timeline(
     213           34 :     nesting_level: usize,
     214           34 :     is_last: &[bool],
     215           34 :     timeline: &TimelineTreeEl,
     216           34 :     timelines: &HashMap<TimelineId, TimelineTreeEl>,
     217           34 : ) -> Result<()> {
     218           34 :     if nesting_level > 0 {
     219           15 :         let ancestor_lsn = match timeline.info.ancestor_lsn {
     220           15 :             Some(lsn) => lsn.to_string(),
     221            0 :             None => "Unknown Lsn".to_string(),
     222              :         };
     223              : 
     224           15 :         let mut br_sym = "┣━";
     225           15 : 
     226           15 :         // Draw each nesting padding with proper style
     227           15 :         // depending on whether its timeline ended or not.
     228           15 :         if nesting_level > 1 {
     229            3 :             for l in &is_last[1..is_last.len() - 1] {
     230            3 :                 if *l {
     231            0 :                     print!("   ");
     232            3 :                 } else {
     233            3 :                     print!("┃  ");
     234            3 :                 }
     235              :             }
     236           12 :         }
     237              : 
     238              :         // We are the last in this sub-timeline
     239           15 :         if *is_last.last().unwrap() {
     240           10 :             br_sym = "┗━";
     241           10 :         }
     242              : 
     243           15 :         print!("{} @{}: ", br_sym, ancestor_lsn);
     244           19 :     }
     245              : 
     246              :     // Finally print a timeline id and name with new line
     247           34 :     println!(
     248           34 :         "{} [{}]",
     249           34 :         timeline.name.as_deref().unwrap_or("_no_name_"),
     250           34 :         timeline.info.timeline_id
     251           34 :     );
     252           34 : 
     253           34 :     let len = timeline.children.len();
     254           34 :     let mut i: usize = 0;
     255           34 :     let mut is_last_new = Vec::from(is_last);
     256           34 :     is_last_new.push(false);
     257              : 
     258           49 :     for child in &timeline.children {
     259           15 :         i += 1;
     260           15 : 
     261           15 :         // Mark that the last padding is the end of the timeline
     262           15 :         if i == len {
     263           10 :             if let Some(last) = is_last_new.last_mut() {
     264           10 :                 *last = true;
     265           10 :             }
     266            5 :         }
     267              : 
     268              :         print_timeline(
     269           15 :             nesting_level + 1,
     270           15 :             &is_last_new,
     271           15 :             timelines
     272           15 :                 .get(child)
     273           15 :                 .context("missing timeline info in the HashMap")?,
     274           15 :             timelines,
     275            0 :         )?;
     276              :     }
     277              : 
     278           34 :     Ok(())
     279           34 : }
     280              : 
     281              : /// Returns a map of timeline IDs to timeline_id@lsn strings.
     282              : /// Connects to the pageserver to query this information.
     283            0 : async fn get_timeline_infos(
     284            0 :     env: &local_env::LocalEnv,
     285            0 :     tenant_shard_id: &TenantShardId,
     286            0 : ) -> Result<HashMap<TimelineId, TimelineInfo>> {
     287            0 :     Ok(get_default_pageserver(env)
     288            0 :         .timeline_list(tenant_shard_id)
     289            0 :         .await?
     290            0 :         .into_iter()
     291            0 :         .map(|timeline_info| (timeline_info.timeline_id, timeline_info))
     292            0 :         .collect())
     293            0 : }
     294              : 
     295              : // Helper function to parse --tenant_id option, or get the default from config file
     296          905 : fn get_tenant_id(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> anyhow::Result<TenantId> {
     297          905 :     if let Some(tenant_id_from_arguments) = parse_tenant_id(sub_match).transpose() {
     298          905 :         tenant_id_from_arguments
     299            0 :     } else if let Some(default_id) = env.default_tenant_id {
     300            0 :         Ok(default_id)
     301              :     } else {
     302            0 :         anyhow::bail!("No tenant id. Use --tenant-id, or set a default tenant");
     303              :     }
     304          905 : }
     305              : 
     306              : // Helper function to parse --tenant_id option, for commands that accept a shard suffix
     307           19 : fn get_tenant_shard_id(
     308           19 :     sub_match: &ArgMatches,
     309           19 :     env: &local_env::LocalEnv,
     310           19 : ) -> anyhow::Result<TenantShardId> {
     311           19 :     if let Some(tenant_id_from_arguments) = parse_tenant_shard_id(sub_match).transpose() {
     312           19 :         tenant_id_from_arguments
     313            0 :     } else if let Some(default_id) = env.default_tenant_id {
     314            0 :         Ok(TenantShardId::unsharded(default_id))
     315              :     } else {
     316            0 :         anyhow::bail!("No tenant shard id. Use --tenant-id, or set a default tenant");
     317              :     }
     318           19 : }
     319              : 
     320         1363 : fn parse_tenant_id(sub_match: &ArgMatches) -> anyhow::Result<Option<TenantId>> {
     321         1363 :     sub_match
     322         1363 :         .get_one::<String>("tenant-id")
     323         1363 :         .map(|tenant_id| TenantId::from_str(tenant_id))
     324         1363 :         .transpose()
     325         1363 :         .context("Failed to parse tenant id from the argument string")
     326         1363 : }
     327              : 
     328           19 : fn parse_tenant_shard_id(sub_match: &ArgMatches) -> anyhow::Result<Option<TenantShardId>> {
     329           19 :     sub_match
     330           19 :         .get_one::<String>("tenant-id")
     331           19 :         .map(|id_str| TenantShardId::from_str(id_str))
     332           19 :         .transpose()
     333           19 :         .context("Failed to parse tenant shard id from the argument string")
     334           19 : }
     335              : 
     336          554 : fn parse_timeline_id(sub_match: &ArgMatches) -> anyhow::Result<Option<TimelineId>> {
     337          554 :     sub_match
     338          554 :         .get_one::<String>("timeline-id")
     339          554 :         .map(|timeline_id| TimelineId::from_str(timeline_id))
     340          554 :         .transpose()
     341          554 :         .context("Failed to parse timeline id from the argument string")
     342          554 : }
     343              : 
     344          354 : fn handle_init(init_match: &ArgMatches) -> anyhow::Result<LocalEnv> {
     345          354 :     let num_pageservers = init_match
     346          354 :         .get_one::<u16>("num-pageservers")
     347          354 :         .expect("num-pageservers arg has a default");
     348              :     // Create config file
     349          354 :     let toml_file: String = if let Some(config_path) = init_match.get_one::<PathBuf>("config") {
     350              :         // load and parse the file
     351          354 :         std::fs::read_to_string(config_path).with_context(|| {
     352            0 :             format!(
     353            0 :                 "Could not read configuration file '{}'",
     354            0 :                 config_path.display()
     355            0 :             )
     356          354 :         })?
     357              :     } else {
     358              :         // Built-in default config
     359            0 :         default_conf(*num_pageservers)
     360              :     };
     361              : 
     362          354 :     let pg_version = init_match
     363          354 :         .get_one::<u32>("pg-version")
     364          354 :         .copied()
     365          354 :         .context("Failed to parse postgres version from the argument string")?;
     366              : 
     367          354 :     let mut env =
     368          354 :         LocalEnv::parse_config(&toml_file).context("Failed to create neon configuration")?;
     369          354 :     let force = init_match.get_one("force").expect("we set a default value");
     370          354 :     env.init(pg_version, force)
     371          354 :         .context("Failed to initialize neon repository")?;
     372              : 
     373              :     // Create remote storage location for default LocalFs remote storage
     374          354 :     std::fs::create_dir_all(env.base_data_dir.join(PAGESERVER_REMOTE_STORAGE_DIR))?;
     375              : 
     376              :     // Initialize pageserver, create initial tenant and timeline.
     377          742 :     for ps_conf in &env.pageservers {
     378          388 :         PageServerNode::from_env(&env, ps_conf)
     379          388 :             .initialize(&pageserver_config_overrides(init_match))
     380          388 :             .unwrap_or_else(|e| {
     381            0 :                 eprintln!("pageserver init failed: {e:?}");
     382            0 :                 exit(1);
     383          388 :             });
     384          388 :     }
     385              : 
     386          354 :     Ok(env)
     387          354 : }
     388              : 
     389              : /// The default pageserver is the one where CLI tenant/timeline operations are sent by default.
     390              : /// For typical interactive use, one would just run with a single pageserver.  Scenarios with
     391              : /// tenant/timeline placement across multiple pageservers are managed by python test code rather
     392              : /// than this CLI.
     393          843 : fn get_default_pageserver(env: &local_env::LocalEnv) -> PageServerNode {
     394          843 :     let ps_conf = env
     395          843 :         .pageservers
     396          843 :         .first()
     397          843 :         .expect("Config is validated to contain at least one pageserver");
     398          843 :     PageServerNode::from_env(env, ps_conf)
     399          843 : }
     400              : 
     401          993 : fn pageserver_config_overrides(init_match: &ArgMatches) -> Vec<&str> {
     402          993 :     init_match
     403          993 :         .get_many::<String>("pageserver-config-override")
     404          993 :         .into_iter()
     405          993 :         .flatten()
     406          993 :         .map(String::as_str)
     407          993 :         .collect()
     408          993 : }
     409              : 
     410          479 : async fn handle_tenant(
     411          479 :     tenant_match: &ArgMatches,
     412          479 :     env: &mut local_env::LocalEnv,
     413          479 : ) -> anyhow::Result<()> {
     414          479 :     let pageserver = get_default_pageserver(env);
     415          479 :     match tenant_match.subcommand() {
     416          479 :         Some(("list", _)) => {
     417           24 :             for t in pageserver.tenant_list().await? {
     418           11 :                 println!("{} {:?}", t.id, t.state);
     419           11 :             }
     420              :         }
     421          473 :         Some(("create", create_match)) => {
     422          459 :             let tenant_conf: HashMap<_, _> = create_match
     423          459 :                 .get_many::<String>("config")
     424          459 :                 .map(|vals: clap::parser::ValuesRef<'_, String>| {
     425          737 :                     vals.flat_map(|c| c.split_once(':')).collect()
     426          459 :                 })
     427          459 :                 .unwrap_or_default();
     428          459 : 
     429          459 :             let shard_count: u8 = create_match
     430          459 :                 .get_one::<u8>("shard-count")
     431          459 :                 .cloned()
     432          459 :                 .unwrap_or(0);
     433          459 : 
     434          459 :             let shard_stripe_size: Option<u32> =
     435          459 :                 create_match.get_one::<u32>("shard-stripe-size").cloned();
     436              : 
     437          459 :             let tenant_conf = PageServerNode::parse_config(tenant_conf)?;
     438              : 
     439              :             // If tenant ID was not specified, generate one
     440          458 :             let tenant_id = parse_tenant_id(create_match)?.unwrap_or_else(TenantId::generate);
     441          458 : 
     442          458 :             // We must register the tenant with the attachment service, so
     443          458 :             // that when the pageserver restarts, it will be re-attached.
     444          458 :             let attachment_service = AttachmentService::from_env(env);
     445          458 :             attachment_service
     446          458 :                 .tenant_create(TenantCreateRequest {
     447          458 :                     // Note that ::unsharded here isn't actually because the tenant is unsharded, its because the
     448          458 :                     // attachment service expecfs a shard-naive tenant_id in this attribute, and the TenantCreateRequest
     449          458 :                     // type is used both in attachment service (for creating tenants) and in pageserver (for creating shards)
     450          458 :                     new_tenant_id: TenantShardId::unsharded(tenant_id),
     451          458 :                     generation: None,
     452          458 :                     shard_parameters: ShardParameters {
     453          458 :                         count: ShardCount(shard_count),
     454          458 :                         stripe_size: shard_stripe_size
     455          458 :                             .map(ShardStripeSize)
     456          458 :                             .unwrap_or(ShardParameters::DEFAULT_STRIPE_SIZE),
     457          458 :                     },
     458          458 :                     config: tenant_conf,
     459          458 :                 })
     460         1374 :                 .await?;
     461          457 :             println!("tenant {tenant_id} successfully created on the pageserver");
     462              : 
     463              :             // Create an initial timeline for the new tenant
     464          457 :             let new_timeline_id =
     465          457 :                 parse_timeline_id(create_match)?.unwrap_or(TimelineId::generate());
     466          457 :             let pg_version = create_match
     467          457 :                 .get_one::<u32>("pg-version")
     468          457 :                 .copied()
     469          457 :                 .context("Failed to parse postgres version from the argument string")?;
     470              : 
     471              :             // FIXME: passing None for ancestor_start_lsn is not kosher in a sharded world: we can't have
     472              :             // different shards picking different start lsns.  Maybe we have to teach attachment service
     473              :             // to let shard 0 branch first and then propagate the chosen LSN to other shards.
     474          457 :             attachment_service
     475          457 :                 .tenant_timeline_create(
     476          457 :                     tenant_id,
     477          457 :                     TimelineCreateRequest {
     478          457 :                         new_timeline_id,
     479          457 :                         ancestor_timeline_id: None,
     480          457 :                         ancestor_start_lsn: None,
     481          457 :                         existing_initdb_timeline_id: None,
     482          457 :                         pg_version: Some(pg_version),
     483          457 :                     },
     484          457 :                 )
     485          457 :                 .await?;
     486              : 
     487          457 :             env.register_branch_mapping(
     488          457 :                 DEFAULT_BRANCH_NAME.to_string(),
     489          457 :                 tenant_id,
     490          457 :                 new_timeline_id,
     491          457 :             )?;
     492              : 
     493          457 :             println!("Created an initial timeline '{new_timeline_id}' for tenant: {tenant_id}",);
     494          457 : 
     495          457 :             if create_match.get_flag("set-default") {
     496            1 :                 println!("Setting tenant {tenant_id} as a default one");
     497            1 :                 env.default_tenant_id = Some(tenant_id);
     498          456 :             }
     499              :         }
     500           14 :         Some(("set-default", set_default_match)) => {
     501            0 :             let tenant_id =
     502            0 :                 parse_tenant_id(set_default_match)?.context("No tenant id specified")?;
     503            0 :             println!("Setting tenant {tenant_id} as a default one");
     504            0 :             env.default_tenant_id = Some(tenant_id);
     505              :         }
     506           14 :         Some(("config", create_match)) => {
     507           14 :             let tenant_id = get_tenant_id(create_match, env)?;
     508           14 :             let tenant_conf: HashMap<_, _> = create_match
     509           14 :                 .get_many::<String>("config")
     510           47 :                 .map(|vals| vals.flat_map(|c| c.split_once(':')).collect())
     511           14 :                 .unwrap_or_default();
     512           14 : 
     513           14 :             pageserver
     514           14 :                 .tenant_config(tenant_id, tenant_conf)
     515           56 :                 .await
     516           14 :                 .with_context(|| format!("Tenant config failed for tenant with id {tenant_id}"))?;
     517           14 :             println!("tenant {tenant_id} successfully configured on the pageserver");
     518              :         }
     519            0 :         Some(("migrate", matches)) => {
     520            0 :             let tenant_shard_id = get_tenant_shard_id(matches, env)?;
     521            0 :             let new_pageserver = get_pageserver(env, matches)?;
     522            0 :             let new_pageserver_id = new_pageserver.conf.id;
     523            0 : 
     524            0 :             let attachment_service = AttachmentService::from_env(env);
     525            0 :             attachment_service
     526            0 :                 .tenant_migrate(tenant_shard_id, new_pageserver_id)
     527            0 :                 .await?;
     528              : 
     529            0 :             println!("tenant {tenant_shard_id} migrated to {}", new_pageserver_id);
     530              :         }
     531            0 :         Some(("status", matches)) => {
     532            0 :             let tenant_id = get_tenant_id(matches, env)?;
     533              : 
     534            0 :             let mut shard_table = comfy_table::Table::new();
     535            0 :             shard_table.set_header(["Shard", "Pageserver", "Physical Size"]);
     536            0 : 
     537            0 :             let mut tenant_synthetic_size = None;
     538            0 : 
     539            0 :             let attachment_service = AttachmentService::from_env(env);
     540            0 :             for shard in attachment_service.tenant_locate(tenant_id).await?.shards {
     541            0 :                 let pageserver =
     542            0 :                     PageServerNode::from_env(env, env.get_pageserver_conf(shard.node_id)?);
     543              : 
     544            0 :                 let size = pageserver
     545            0 :                     .http_client
     546            0 :                     .tenant_details(shard.shard_id)
     547            0 :                     .await?
     548              :                     .tenant_info
     549              :                     .current_physical_size
     550            0 :                     .unwrap();
     551            0 : 
     552            0 :                 shard_table.add_row([
     553            0 :                     format!("{}", shard.shard_id.shard_slug()),
     554            0 :                     format!("{}", shard.node_id.0),
     555            0 :                     format!("{} MiB", size / (1024 * 1024)),
     556            0 :                 ]);
     557            0 : 
     558            0 :                 if shard.shard_id.is_zero() {
     559              :                     tenant_synthetic_size =
     560            0 :                         Some(pageserver.tenant_synthetic_size(shard.shard_id).await?);
     561            0 :                 }
     562              :             }
     563              : 
     564            0 :             let Some(synthetic_size) = tenant_synthetic_size else {
     565            0 :                 bail!("Shard 0 not found")
     566              :             };
     567              : 
     568            0 :             let mut tenant_table = comfy_table::Table::new();
     569            0 :             tenant_table.add_row(["Tenant ID".to_string(), tenant_id.to_string()]);
     570            0 :             tenant_table.add_row([
     571            0 :                 "Synthetic size".to_string(),
     572            0 :                 format!("{} MiB", synthetic_size.size.unwrap_or(0) / (1024 * 1024)),
     573            0 :             ]);
     574            0 : 
     575            0 :             println!("{tenant_table}");
     576            0 :             println!("{shard_table}");
     577              :         }
     578            0 :         Some((sub_name, _)) => bail!("Unexpected tenant subcommand '{}'", sub_name),
     579            0 :         None => bail!("no tenant subcommand provided"),
     580              :     }
     581          477 :     Ok(())
     582          479 : }
     583              : 
     584          364 : async fn handle_timeline(timeline_match: &ArgMatches, env: &mut local_env::LocalEnv) -> Result<()> {
     585          364 :     let pageserver = get_default_pageserver(env);
     586          364 : 
     587          364 :     match timeline_match.subcommand() {
     588          364 :         Some(("list", list_match)) => {
     589              :             // TODO(sharding): this command shouldn't have to specify a shard ID: we should ask the attachment service
     590              :             // where shard 0 is attached, and query there.
     591           19 :             let tenant_shard_id = get_tenant_shard_id(list_match, env)?;
     592           76 :             let timelines = pageserver.timeline_list(&tenant_shard_id).await?;
     593           19 :             print_timelines_tree(timelines, env.timeline_name_mappings())?;
     594              :         }
     595          345 :         Some(("create", create_match)) => {
     596           91 :             let tenant_id = get_tenant_id(create_match, env)?;
     597           91 :             let new_branch_name = create_match
     598           91 :                 .get_one::<String>("branch-name")
     599           91 :                 .ok_or_else(|| anyhow!("No branch name provided"))?;
     600              : 
     601           91 :             let pg_version = create_match
     602           91 :                 .get_one::<u32>("pg-version")
     603           91 :                 .copied()
     604           91 :                 .context("Failed to parse postgres version from the argument string")?;
     605              : 
     606           91 :             let new_timeline_id_opt = parse_timeline_id(create_match)?;
     607           91 :             let new_timeline_id = new_timeline_id_opt.unwrap_or(TimelineId::generate());
     608           91 : 
     609           91 :             let attachment_service = AttachmentService::from_env(env);
     610           91 :             let create_req = TimelineCreateRequest {
     611           91 :                 new_timeline_id,
     612           91 :                 ancestor_timeline_id: None,
     613           91 :                 existing_initdb_timeline_id: None,
     614           91 :                 ancestor_start_lsn: None,
     615           91 :                 pg_version: Some(pg_version),
     616           91 :             };
     617           91 :             let timeline_info = attachment_service
     618           91 :                 .tenant_timeline_create(tenant_id, create_req)
     619          273 :                 .await?;
     620              : 
     621           91 :             let last_record_lsn = timeline_info.last_record_lsn;
     622           91 :             env.register_branch_mapping(new_branch_name.to_string(), tenant_id, new_timeline_id)?;
     623              : 
     624           91 :             println!(
     625           91 :                 "Created timeline '{}' at Lsn {last_record_lsn} for tenant: {tenant_id}",
     626           91 :                 timeline_info.timeline_id
     627           91 :             );
     628              :         }
     629          254 :         Some(("import", import_match)) => {
     630            6 :             let tenant_id = get_tenant_id(import_match, env)?;
     631            6 :             let timeline_id = parse_timeline_id(import_match)?.expect("No timeline id provided");
     632            6 :             let name = import_match
     633            6 :                 .get_one::<String>("node-name")
     634            6 :                 .ok_or_else(|| anyhow!("No node name provided"))?;
     635              : 
     636              :             // Parse base inputs
     637            6 :             let base_tarfile = import_match
     638            6 :                 .get_one::<PathBuf>("base-tarfile")
     639            6 :                 .ok_or_else(|| anyhow!("No base-tarfile provided"))?
     640            6 :                 .to_owned();
     641            6 :             let base_lsn = Lsn::from_str(
     642            6 :                 import_match
     643            6 :                     .get_one::<String>("base-lsn")
     644            6 :                     .ok_or_else(|| anyhow!("No base-lsn provided"))?,
     645            0 :             )?;
     646            6 :             let base = (base_lsn, base_tarfile);
     647            6 : 
     648            6 :             // Parse pg_wal inputs
     649            6 :             let wal_tarfile = import_match.get_one::<PathBuf>("wal-tarfile").cloned();
     650            6 :             let end_lsn = import_match
     651            6 :                 .get_one::<String>("end-lsn")
     652            6 :                 .map(|s| Lsn::from_str(s).unwrap());
     653            6 :             // TODO validate both or none are provided
     654            6 :             let pg_wal = end_lsn.zip(wal_tarfile);
     655              : 
     656            6 :             let pg_version = import_match
     657            6 :                 .get_one::<u32>("pg-version")
     658            6 :                 .copied()
     659            6 :                 .context("Failed to parse postgres version from the argument string")?;
     660              : 
     661            6 :             let mut cplane = ComputeControlPlane::load(env.clone())?;
     662            6 :             println!("Importing timeline into pageserver ...");
     663            6 :             pageserver
     664            6 :                 .timeline_import(tenant_id, timeline_id, base, pg_wal, pg_version)
     665        65461 :                 .await?;
     666            4 :             env.register_branch_mapping(name.to_string(), tenant_id, timeline_id)?;
     667              : 
     668            4 :             println!("Creating endpoint for imported timeline ...");
     669            4 :             cplane.new_endpoint(
     670            4 :                 name,
     671            4 :                 tenant_id,
     672            4 :                 timeline_id,
     673            4 :                 None,
     674            4 :                 None,
     675            4 :                 pg_version,
     676            4 :                 ComputeMode::Primary,
     677            4 :             )?;
     678            4 :             println!("Done");
     679              :         }
     680          248 :         Some(("branch", branch_match)) => {
     681          248 :             let tenant_id = get_tenant_id(branch_match, env)?;
     682          248 :             let new_branch_name = branch_match
     683          248 :                 .get_one::<String>("branch-name")
     684          248 :                 .ok_or_else(|| anyhow!("No branch name provided"))?;
     685          248 :             let ancestor_branch_name = branch_match
     686          248 :                 .get_one::<String>("ancestor-branch-name")
     687          248 :                 .map(|s| s.as_str())
     688          248 :                 .unwrap_or(DEFAULT_BRANCH_NAME);
     689          248 :             let ancestor_timeline_id = env
     690          248 :                 .get_branch_timeline_id(ancestor_branch_name, tenant_id)
     691          248 :                 .ok_or_else(|| {
     692            0 :                     anyhow!("Found no timeline id for branch name '{ancestor_branch_name}'")
     693          248 :                 })?;
     694              : 
     695          248 :             let start_lsn = branch_match
     696          248 :                 .get_one::<String>("ancestor-start-lsn")
     697          248 :                 .map(|lsn_str| Lsn::from_str(lsn_str))
     698          248 :                 .transpose()
     699          248 :                 .context("Failed to parse ancestor start Lsn from the request")?;
     700          248 :             let new_timeline_id = TimelineId::generate();
     701          248 :             let attachment_service = AttachmentService::from_env(env);
     702          248 :             let create_req = TimelineCreateRequest {
     703          248 :                 new_timeline_id,
     704          248 :                 ancestor_timeline_id: Some(ancestor_timeline_id),
     705          248 :                 existing_initdb_timeline_id: None,
     706          248 :                 ancestor_start_lsn: start_lsn,
     707          248 :                 pg_version: None,
     708          248 :             };
     709          248 :             let timeline_info = attachment_service
     710          248 :                 .tenant_timeline_create(tenant_id, create_req)
     711          744 :                 .await?;
     712              : 
     713          244 :             let last_record_lsn = timeline_info.last_record_lsn;
     714          244 : 
     715          244 :             env.register_branch_mapping(new_branch_name.to_string(), tenant_id, new_timeline_id)?;
     716              : 
     717          244 :             println!(
     718          244 :                 "Created timeline '{}' at Lsn {last_record_lsn} for tenant: {tenant_id}. Ancestor timeline: '{ancestor_branch_name}'",
     719          244 :                 timeline_info.timeline_id
     720          244 :             );
     721              :         }
     722            0 :         Some((sub_name, _)) => bail!("Unexpected tenant subcommand '{sub_name}'"),
     723            0 :         None => bail!("no tenant subcommand provided"),
     724              :     }
     725              : 
     726          358 :     Ok(())
     727          364 : }
     728              : 
     729         1926 : async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
     730         1926 :     let (sub_name, sub_args) = match ep_match.subcommand() {
     731         1926 :         Some(ep_subcommand_data) => ep_subcommand_data,
     732            0 :         None => bail!("no endpoint subcommand provided"),
     733              :     };
     734         1926 :     let mut cplane = ComputeControlPlane::load(env.clone())?;
     735              : 
     736         1926 :     match sub_name {
     737         1926 :         "list" => {
     738              :             // TODO(sharding): this command shouldn't have to specify a shard ID: we should ask the attachment service
     739              :             // where shard 0 is attached, and query there.
     740            0 :             let tenant_shard_id = get_tenant_shard_id(sub_args, env)?;
     741            0 :             let timeline_infos = get_timeline_infos(env, &tenant_shard_id)
     742            0 :                 .await
     743            0 :                 .unwrap_or_else(|e| {
     744            0 :                     eprintln!("Failed to load timeline info: {}", e);
     745            0 :                     HashMap::new()
     746            0 :                 });
     747            0 : 
     748            0 :             let timeline_name_mappings = env.timeline_name_mappings();
     749            0 : 
     750            0 :             let mut table = comfy_table::Table::new();
     751            0 : 
     752            0 :             table.load_preset(comfy_table::presets::NOTHING);
     753            0 : 
     754            0 :             table.set_header([
     755            0 :                 "ENDPOINT",
     756            0 :                 "ADDRESS",
     757            0 :                 "TIMELINE",
     758            0 :                 "BRANCH NAME",
     759            0 :                 "LSN",
     760            0 :                 "STATUS",
     761            0 :             ]);
     762              : 
     763            0 :             for (endpoint_id, endpoint) in cplane
     764            0 :                 .endpoints
     765            0 :                 .iter()
     766            0 :                 .filter(|(_, endpoint)| endpoint.tenant_id == tenant_shard_id.tenant_id)
     767            0 :             {
     768            0 :                 let lsn_str = match endpoint.mode {
     769            0 :                     ComputeMode::Static(lsn) => {
     770            0 :                         // -> read-only endpoint
     771            0 :                         // Use the node's LSN.
     772            0 :                         lsn.to_string()
     773              :                     }
     774              :                     _ => {
     775              :                         // -> primary endpoint or hot replica
     776              :                         // Use the LSN at the end of the timeline.
     777            0 :                         timeline_infos
     778            0 :                             .get(&endpoint.timeline_id)
     779            0 :                             .map(|bi| bi.last_record_lsn.to_string())
     780            0 :                             .unwrap_or_else(|| "?".to_string())
     781              :                     }
     782              :                 };
     783              : 
     784            0 :                 let branch_name = timeline_name_mappings
     785            0 :                     .get(&TenantTimelineId::new(
     786            0 :                         tenant_shard_id.tenant_id,
     787            0 :                         endpoint.timeline_id,
     788            0 :                     ))
     789            0 :                     .map(|name| name.as_str())
     790            0 :                     .unwrap_or("?");
     791            0 : 
     792            0 :                 table.add_row([
     793            0 :                     endpoint_id.as_str(),
     794            0 :                     &endpoint.pg_address.to_string(),
     795            0 :                     &endpoint.timeline_id.to_string(),
     796            0 :                     branch_name,
     797            0 :                     lsn_str.as_str(),
     798            0 :                     &format!("{}", endpoint.status()),
     799            0 :                 ]);
     800              :             }
     801              : 
     802            0 :             println!("{table}");
     803              :         }
     804         1926 :         "create" => {
     805          546 :             let tenant_id = get_tenant_id(sub_args, env)?;
     806          546 :             let branch_name = sub_args
     807          546 :                 .get_one::<String>("branch-name")
     808          546 :                 .map(|s| s.as_str())
     809          546 :                 .unwrap_or(DEFAULT_BRANCH_NAME);
     810          546 :             let endpoint_id = sub_args
     811          546 :                 .get_one::<String>("endpoint_id")
     812          546 :                 .map(String::to_string)
     813          546 :                 .unwrap_or_else(|| format!("ep-{branch_name}"));
     814              : 
     815          546 :             let lsn = sub_args
     816          546 :                 .get_one::<String>("lsn")
     817          546 :                 .map(|lsn_str| Lsn::from_str(lsn_str))
     818          546 :                 .transpose()
     819          546 :                 .context("Failed to parse Lsn from the request")?;
     820          546 :             let timeline_id = env
     821          546 :                 .get_branch_timeline_id(branch_name, tenant_id)
     822          546 :                 .ok_or_else(|| anyhow!("Found no timeline id for branch name '{branch_name}'"))?;
     823              : 
     824          546 :             let pg_port: Option<u16> = sub_args.get_one::<u16>("pg-port").copied();
     825          546 :             let http_port: Option<u16> = sub_args.get_one::<u16>("http-port").copied();
     826          546 :             let pg_version = sub_args
     827          546 :                 .get_one::<u32>("pg-version")
     828          546 :                 .copied()
     829          546 :                 .context("Failed to parse postgres version from the argument string")?;
     830              : 
     831          546 :             let hot_standby = sub_args
     832          546 :                 .get_one::<bool>("hot-standby")
     833          546 :                 .copied()
     834          546 :                 .unwrap_or(false);
     835              : 
     836          546 :             let mode = match (lsn, hot_standby) {
     837           49 :                 (Some(lsn), false) => ComputeMode::Static(lsn),
     838            5 :                 (None, true) => ComputeMode::Replica,
     839          492 :                 (None, false) => ComputeMode::Primary,
     840            0 :                 (Some(_), true) => anyhow::bail!("cannot specify both lsn and hot-standby"),
     841              :             };
     842              : 
     843          546 :             match (mode, hot_standby) {
     844              :                 (ComputeMode::Static(_), true) => {
     845            0 :                     bail!("Cannot start a node in hot standby mode when it is already configured as a static replica")
     846              :                 }
     847              :                 (ComputeMode::Primary, true) => {
     848            0 :                     bail!("Cannot start a node as a hot standby replica, it is already configured as primary node")
     849              :                 }
     850          546 :                 _ => {}
     851          546 :             }
     852          546 : 
     853          546 :             cplane.check_conflicting_endpoints(mode, tenant_id, timeline_id)?;
     854              : 
     855          518 :             cplane.new_endpoint(
     856          518 :                 &endpoint_id,
     857          518 :                 tenant_id,
     858          518 :                 timeline_id,
     859          518 :                 pg_port,
     860          518 :                 http_port,
     861          518 :                 pg_version,
     862          518 :                 mode,
     863          518 :             )?;
     864              :         }
     865         1380 :         "start" => {
     866          585 :             let endpoint_id = sub_args
     867          585 :                 .get_one::<String>("endpoint_id")
     868          585 :                 .ok_or_else(|| anyhow!("No endpoint ID was provided to start"))?;
     869              : 
     870          585 :             let pageserver_id =
     871          585 :                 if let Some(id_str) = sub_args.get_one::<String>("endpoint-pageserver-id") {
     872              :                     Some(NodeId(
     873           22 :                         id_str.parse().context("while parsing pageserver id")?,
     874              :                     ))
     875              :                 } else {
     876          563 :                     None
     877              :                 };
     878              : 
     879          585 :             let remote_ext_config = sub_args.get_one::<String>("remote-ext-config");
     880              : 
     881              :             // If --safekeepers argument is given, use only the listed safekeeper nodes.
     882          585 :             let safekeepers =
     883          585 :                 if let Some(safekeepers_str) = sub_args.get_one::<String>("safekeepers") {
     884          580 :                     let mut safekeepers: Vec<NodeId> = Vec::new();
     885          727 :                     for sk_id in safekeepers_str.split(',').map(str::trim) {
     886          727 :                         let sk_id = NodeId(u64::from_str(sk_id).map_err(|_| {
     887            0 :                             anyhow!("invalid node ID \"{sk_id}\" in --safekeepers list")
     888          727 :                         })?);
     889          727 :                         safekeepers.push(sk_id);
     890              :                     }
     891          580 :                     safekeepers
     892              :                 } else {
     893            5 :                     env.safekeepers.iter().map(|sk| sk.id).collect()
     894              :                 };
     895              : 
     896          585 :             let endpoint = cplane
     897          585 :                 .endpoints
     898          585 :                 .get(endpoint_id.as_str())
     899          585 :                 .ok_or_else(|| anyhow::anyhow!("endpoint {endpoint_id} not found"))?;
     900              : 
     901          585 :             cplane.check_conflicting_endpoints(
     902          585 :                 endpoint.mode,
     903          585 :                 endpoint.tenant_id,
     904          585 :                 endpoint.timeline_id,
     905          585 :             )?;
     906              : 
     907          584 :             let (pageservers, stripe_size) = if let Some(pageserver_id) = pageserver_id {
     908           22 :                 let conf = env.get_pageserver_conf(pageserver_id).unwrap();
     909           22 :                 let parsed = parse_host_port(&conf.listen_pg_addr).expect("Bad config");
     910           22 :                 (
     911           22 :                     vec![(parsed.0, parsed.1.unwrap_or(5432))],
     912           22 :                     // If caller is telling us what pageserver to use, this is not a tenant which is
     913           22 :                     // full managed by attachment service, therefore not sharded.
     914           22 :                     ShardParameters::DEFAULT_STRIPE_SIZE,
     915           22 :                 )
     916              :             } else {
     917              :                 // Look up the currently attached location of the tenant, and its striping metadata,
     918              :                 // to pass these on to postgres.
     919          562 :                 let attachment_service = AttachmentService::from_env(env);
     920         1686 :                 let locate_result = attachment_service.tenant_locate(endpoint.tenant_id).await?;
     921          562 :                 let pageservers = locate_result
     922          562 :                     .shards
     923          562 :                     .into_iter()
     924          574 :                     .map(|shard| {
     925          574 :                         (
     926          574 :                             Host::parse(&shard.listen_pg_addr)
     927          574 :                                 .expect("Attachment service reported bad hostname"),
     928          574 :                             shard.listen_pg_port,
     929          574 :                         )
     930          574 :                     })
     931          562 :                     .collect::<Vec<_>>();
     932          562 :                 let stripe_size = locate_result.shard_params.stripe_size;
     933          562 : 
     934          562 :                 (pageservers, stripe_size)
     935              :             };
     936          584 :             assert!(!pageservers.is_empty());
     937              : 
     938          584 :             let ps_conf = env.get_pageserver_conf(DEFAULT_PAGESERVER_ID)?;
     939          584 :             let auth_token = if matches!(ps_conf.pg_auth_type, AuthType::NeonJWT) {
     940           15 :                 let claims = Claims::new(Some(endpoint.tenant_id), Scope::Tenant);
     941           15 : 
     942           15 :                 Some(env.generate_auth_token(&claims)?)
     943              :             } else {
     944          569 :                 None
     945              :             };
     946              : 
     947          584 :             println!("Starting existing endpoint {endpoint_id}...");
     948          584 :             endpoint
     949          584 :                 .start(
     950          584 :                     &auth_token,
     951          584 :                     safekeepers,
     952          584 :                     pageservers,
     953          584 :                     remote_ext_config,
     954          584 :                     stripe_size.0 as usize,
     955          584 :                 )
     956        10997 :                 .await?;
     957              :         }
     958          795 :         "reconfigure" => {
     959          221 :             let endpoint_id = sub_args
     960          221 :                 .get_one::<String>("endpoint_id")
     961          221 :                 .ok_or_else(|| anyhow!("No endpoint ID provided to reconfigure"))?;
     962          221 :             let endpoint = cplane
     963          221 :                 .endpoints
     964          221 :                 .get(endpoint_id.as_str())
     965          221 :                 .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?;
     966          221 :             let pageservers =
     967          221 :                 if let Some(id_str) = sub_args.get_one::<String>("endpoint-pageserver-id") {
     968          219 :                     let ps_id = NodeId(id_str.parse().context("while parsing pageserver id")?);
     969          219 :                     let pageserver = PageServerNode::from_env(env, env.get_pageserver_conf(ps_id)?);
     970          219 :                     vec![(
     971          219 :                         pageserver.pg_connection_config.host().clone(),
     972          219 :                         pageserver.pg_connection_config.port(),
     973          219 :                     )]
     974              :                 } else {
     975            2 :                     let attachment_service = AttachmentService::from_env(env);
     976            2 :                     attachment_service
     977            2 :                         .tenant_locate(endpoint.tenant_id)
     978            6 :                         .await?
     979              :                         .shards
     980            2 :                         .into_iter()
     981            8 :                         .map(|shard| {
     982            8 :                             (
     983            8 :                                 Host::parse(&shard.listen_pg_addr)
     984            8 :                                     .expect("Attachment service reported malformed host"),
     985            8 :                                 shard.listen_pg_port,
     986            8 :                             )
     987            8 :                         })
     988            2 :                         .collect::<Vec<_>>()
     989              :                 };
     990          663 :             endpoint.reconfigure(pageservers).await?;
     991              :         }
     992          574 :         "stop" => {
     993          574 :             let endpoint_id = sub_args
     994          574 :                 .get_one::<String>("endpoint_id")
     995          574 :                 .ok_or_else(|| anyhow!("No endpoint ID was provided to stop"))?;
     996          574 :             let destroy = sub_args.get_flag("destroy");
     997              : 
     998          574 :             let endpoint = cplane
     999          574 :                 .endpoints
    1000          574 :                 .get(endpoint_id.as_str())
    1001          574 :                 .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?;
    1002          574 :             endpoint.stop(destroy)?;
    1003              :         }
    1004              : 
    1005            0 :         _ => bail!("Unexpected endpoint subcommand '{sub_name}'"),
    1006              :     }
    1007              : 
    1008         1889 :     Ok(())
    1009         1926 : }
    1010              : 
    1011            1 : fn handle_mappings(sub_match: &ArgMatches, env: &mut local_env::LocalEnv) -> Result<()> {
    1012            1 :     let (sub_name, sub_args) = match sub_match.subcommand() {
    1013            1 :         Some(ep_subcommand_data) => ep_subcommand_data,
    1014            0 :         None => bail!("no mappings subcommand provided"),
    1015              :     };
    1016              : 
    1017            1 :     match sub_name {
    1018            1 :         "map" => {
    1019            1 :             let branch_name = sub_args
    1020            1 :                 .get_one::<String>("branch-name")
    1021            1 :                 .expect("branch-name argument missing");
    1022            1 : 
    1023            1 :             let tenant_id = sub_args
    1024            1 :                 .get_one::<String>("tenant-id")
    1025            1 :                 .map(|x| TenantId::from_str(x))
    1026            1 :                 .expect("tenant-id argument missing")
    1027            1 :                 .expect("malformed tenant-id arg");
    1028            1 : 
    1029            1 :             let timeline_id = sub_args
    1030            1 :                 .get_one::<String>("timeline-id")
    1031            1 :                 .map(|x| TimelineId::from_str(x))
    1032            1 :                 .expect("timeline-id argument missing")
    1033            1 :                 .expect("malformed timeline-id arg");
    1034            1 : 
    1035            1 :             env.register_branch_mapping(branch_name.to_owned(), tenant_id, timeline_id)?;
    1036              : 
    1037            1 :             Ok(())
    1038              :         }
    1039            0 :         other => unimplemented!("mappings subcommand {other}"),
    1040              :     }
    1041            1 : }
    1042              : 
    1043         1203 : fn get_pageserver(env: &local_env::LocalEnv, args: &ArgMatches) -> Result<PageServerNode> {
    1044         1203 :     let node_id = if let Some(id_str) = args.get_one::<String>("pageserver-id") {
    1045         1203 :         NodeId(id_str.parse().context("while parsing pageserver id")?)
    1046              :     } else {
    1047            0 :         DEFAULT_PAGESERVER_ID
    1048              :     };
    1049              : 
    1050              :     Ok(PageServerNode::from_env(
    1051         1203 :         env,
    1052         1203 :         env.get_pageserver_conf(node_id)?,
    1053              :     ))
    1054         1203 : }
    1055              : 
    1056         1203 : async fn handle_pageserver(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
    1057         1203 :     match sub_match.subcommand() {
    1058         1203 :         Some(("start", subcommand_args)) => {
    1059          601 :             let register = subcommand_args.get_one::<bool>("register").unwrap_or(&true);
    1060          601 :             if let Err(e) = get_pageserver(env, subcommand_args)?
    1061          601 :                 .start(&pageserver_config_overrides(subcommand_args), *register)
    1062         5449 :                 .await
    1063              :             {
    1064            0 :                 eprintln!("pageserver start failed: {e}");
    1065            0 :                 exit(1);
    1066          601 :             }
    1067              :         }
    1068              : 
    1069          602 :         Some(("stop", subcommand_args)) => {
    1070          602 :             let immediate = subcommand_args
    1071          602 :                 .get_one::<String>("stop-mode")
    1072          602 :                 .map(|s| s.as_str())
    1073          602 :                 == Some("immediate");
    1074              : 
    1075          602 :             if let Err(e) = get_pageserver(env, subcommand_args)?.stop(immediate) {
    1076            0 :                 eprintln!("pageserver stop failed: {}", e);
    1077            0 :                 exit(1);
    1078          601 :             }
    1079              :         }
    1080              : 
    1081            0 :         Some(("restart", subcommand_args)) => {
    1082            0 :             let pageserver = get_pageserver(env, subcommand_args)?;
    1083              :             //TODO what shutdown strategy should we use here?
    1084            0 :             if let Err(e) = pageserver.stop(false) {
    1085            0 :                 eprintln!("pageserver stop failed: {}", e);
    1086            0 :                 exit(1);
    1087            0 :             }
    1088              : 
    1089            0 :             if let Err(e) = pageserver
    1090            0 :                 .start(&pageserver_config_overrides(subcommand_args), false)
    1091            0 :                 .await
    1092              :             {
    1093            0 :                 eprintln!("pageserver start failed: {e}");
    1094            0 :                 exit(1);
    1095            0 :             }
    1096              :         }
    1097              : 
    1098            0 :         Some(("set-state", subcommand_args)) => {
    1099            0 :             let pageserver = get_pageserver(env, subcommand_args)?;
    1100            0 :             let scheduling = subcommand_args.get_one("scheduling");
    1101            0 :             let availability = subcommand_args.get_one("availability");
    1102            0 : 
    1103            0 :             let attachment_service = AttachmentService::from_env(env);
    1104            0 :             attachment_service
    1105            0 :                 .node_configure(NodeConfigureRequest {
    1106            0 :                     node_id: pageserver.conf.id,
    1107            0 :                     scheduling: scheduling.cloned(),
    1108            0 :                     availability: availability.cloned(),
    1109            0 :                 })
    1110            0 :                 .await?;
    1111              :         }
    1112              : 
    1113            0 :         Some(("status", subcommand_args)) => {
    1114            0 :             match get_pageserver(env, subcommand_args)?.check_status().await {
    1115            0 :                 Ok(_) => println!("Page server is up and running"),
    1116            0 :                 Err(err) => {
    1117            0 :                     eprintln!("Page server is not available: {}", err);
    1118            0 :                     exit(1);
    1119              :                 }
    1120              :             }
    1121              :         }
    1122              : 
    1123            0 :         Some((sub_name, _)) => bail!("Unexpected pageserver subcommand '{}'", sub_name),
    1124            0 :         None => bail!("no pageserver subcommand provided"),
    1125              :     }
    1126         1202 :     Ok(())
    1127         1203 : }
    1128              : 
    1129          718 : async fn handle_attachment_service(
    1130          718 :     sub_match: &ArgMatches,
    1131          718 :     env: &local_env::LocalEnv,
    1132          718 : ) -> Result<()> {
    1133          718 :     let svc = AttachmentService::from_env(env);
    1134          718 :     match sub_match.subcommand() {
    1135          718 :         Some(("start", _start_match)) => {
    1136         4291 :             if let Err(e) = svc.start().await {
    1137            0 :                 eprintln!("start failed: {e}");
    1138            0 :                 exit(1);
    1139          358 :             }
    1140              :         }
    1141              : 
    1142          360 :         Some(("stop", stop_match)) => {
    1143          360 :             let immediate = stop_match
    1144          360 :                 .get_one::<String>("stop-mode")
    1145          360 :                 .map(|s| s.as_str())
    1146          360 :                 == Some("immediate");
    1147              : 
    1148          724 :             if let Err(e) = svc.stop(immediate).await {
    1149            0 :                 eprintln!("stop failed: {}", e);
    1150            0 :                 exit(1);
    1151          360 :             }
    1152              :         }
    1153            0 :         Some((sub_name, _)) => bail!("Unexpected attachment_service subcommand '{}'", sub_name),
    1154            0 :         None => bail!("no attachment_service subcommand provided"),
    1155              :     }
    1156          718 :     Ok(())
    1157          718 : }
    1158              : 
    1159         1055 : fn get_safekeeper(env: &local_env::LocalEnv, id: NodeId) -> Result<SafekeeperNode> {
    1160         1406 :     if let Some(node) = env.safekeepers.iter().find(|node| node.id == id) {
    1161         1055 :         Ok(SafekeeperNode::from_env(env, node))
    1162              :     } else {
    1163            0 :         bail!("could not find safekeeper {id}")
    1164              :     }
    1165         1055 : }
    1166              : 
    1167              : // Get list of options to append to safekeeper command invocation.
    1168          509 : fn safekeeper_extra_opts(init_match: &ArgMatches) -> Vec<String> {
    1169          509 :     init_match
    1170          509 :         .get_many::<String>("safekeeper-extra-opt")
    1171          509 :         .into_iter()
    1172          509 :         .flatten()
    1173          509 :         .map(|s| s.to_owned())
    1174          509 :         .collect()
    1175          509 : }
    1176              : 
    1177         1055 : async fn handle_safekeeper(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
    1178         1055 :     let (sub_name, sub_args) = match sub_match.subcommand() {
    1179         1055 :         Some(safekeeper_command_data) => safekeeper_command_data,
    1180            0 :         None => bail!("no safekeeper subcommand provided"),
    1181              :     };
    1182              : 
    1183              :     // All the commands take an optional safekeeper name argument
    1184         1055 :     let sk_id = if let Some(id_str) = sub_args.get_one::<String>("id") {
    1185         1037 :         NodeId(id_str.parse().context("while parsing safekeeper id")?)
    1186              :     } else {
    1187           18 :         DEFAULT_SAFEKEEPER_ID
    1188              :     };
    1189         1055 :     let safekeeper = get_safekeeper(env, sk_id)?;
    1190              : 
    1191         1055 :     match sub_name {
    1192         1055 :         "start" => {
    1193          509 :             let extra_opts = safekeeper_extra_opts(sub_args);
    1194              : 
    1195         2041 :             if let Err(e) = safekeeper.start(extra_opts).await {
    1196            0 :                 eprintln!("safekeeper start failed: {}", e);
    1197            0 :                 exit(1);
    1198          509 :             }
    1199              :         }
    1200              : 
    1201          546 :         "stop" => {
    1202          546 :             let immediate =
    1203          546 :                 sub_args.get_one::<String>("stop-mode").map(|s| s.as_str()) == Some("immediate");
    1204              : 
    1205          546 :             if let Err(e) = safekeeper.stop(immediate) {
    1206            0 :                 eprintln!("safekeeper stop failed: {}", e);
    1207            0 :                 exit(1);
    1208          546 :             }
    1209              :         }
    1210              : 
    1211            0 :         "restart" => {
    1212            0 :             let immediate =
    1213            0 :                 sub_args.get_one::<String>("stop-mode").map(|s| s.as_str()) == Some("immediate");
    1214              : 
    1215            0 :             if let Err(e) = safekeeper.stop(immediate) {
    1216            0 :                 eprintln!("safekeeper stop failed: {}", e);
    1217            0 :                 exit(1);
    1218            0 :             }
    1219            0 : 
    1220            0 :             let extra_opts = safekeeper_extra_opts(sub_args);
    1221            0 :             if let Err(e) = safekeeper.start(extra_opts).await {
    1222            0 :                 eprintln!("safekeeper start failed: {}", e);
    1223            0 :                 exit(1);
    1224            0 :             }
    1225              :         }
    1226              : 
    1227              :         _ => {
    1228            0 :             bail!("Unexpected safekeeper subcommand '{}'", sub_name)
    1229              :         }
    1230              :     }
    1231         1055 :     Ok(())
    1232         1055 : }
    1233              : 
    1234            3 : async fn handle_start_all(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> anyhow::Result<()> {
    1235            3 :     // Endpoints are not started automatically
    1236            3 : 
    1237           10 :     broker::start_broker_process(env).await?;
    1238              : 
    1239              :     // Only start the attachment service if the pageserver is configured to need it
    1240            3 :     if env.control_plane_api.is_some() {
    1241            3 :         let attachment_service = AttachmentService::from_env(env);
    1242           34 :         if let Err(e) = attachment_service.start().await {
    1243            0 :             eprintln!("attachment_service start failed: {:#}", e);
    1244            0 :             try_stop_all(env, true).await;
    1245            0 :             exit(1);
    1246            3 :         }
    1247            0 :     }
    1248              : 
    1249            7 :     for ps_conf in &env.pageservers {
    1250            4 :         let pageserver = PageServerNode::from_env(env, ps_conf);
    1251            4 :         if let Err(e) = pageserver
    1252            4 :             .start(&pageserver_config_overrides(sub_match), true)
    1253           35 :             .await
    1254              :         {
    1255            0 :             eprintln!("pageserver {} start failed: {:#}", ps_conf.id, e);
    1256            0 :             try_stop_all(env, true).await;
    1257            0 :             exit(1);
    1258            4 :         }
    1259              :     }
    1260              : 
    1261            4 :     for node in env.safekeepers.iter() {
    1262            4 :         let safekeeper = SafekeeperNode::from_env(env, node);
    1263           16 :         if let Err(e) = safekeeper.start(vec![]).await {
    1264            0 :             eprintln!("safekeeper {} start failed: {:#}", safekeeper.id, e);
    1265            0 :             try_stop_all(env, false).await;
    1266            0 :             exit(1);
    1267            4 :         }
    1268              :     }
    1269            3 :     Ok(())
    1270            3 : }
    1271              : 
    1272            3 : async fn handle_stop_all(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
    1273            3 :     let immediate =
    1274            3 :         sub_match.get_one::<String>("stop-mode").map(|s| s.as_str()) == Some("immediate");
    1275            3 : 
    1276            6 :     try_stop_all(env, immediate).await;
    1277              : 
    1278            3 :     Ok(())
    1279            3 : }
    1280              : 
    1281            3 : async fn try_stop_all(env: &local_env::LocalEnv, immediate: bool) {
    1282            3 :     // Stop all endpoints
    1283            3 :     match ComputeControlPlane::load(env.clone()) {
    1284            3 :         Ok(cplane) => {
    1285            5 :             for (_k, node) in cplane.endpoints {
    1286            2 :                 if let Err(e) = node.stop(false) {
    1287            0 :                     eprintln!("postgres stop failed: {e:#}");
    1288            2 :                 }
    1289              :             }
    1290              :         }
    1291            0 :         Err(e) => {
    1292            0 :             eprintln!("postgres stop failed, could not restore control plane data from env: {e:#}")
    1293              :         }
    1294              :     }
    1295              : 
    1296            7 :     for ps_conf in &env.pageservers {
    1297            4 :         let pageserver = PageServerNode::from_env(env, ps_conf);
    1298            4 :         if let Err(e) = pageserver.stop(immediate) {
    1299            0 :             eprintln!("pageserver {} stop failed: {:#}", ps_conf.id, e);
    1300            4 :         }
    1301              :     }
    1302              : 
    1303            4 :     for node in env.safekeepers.iter() {
    1304            4 :         let safekeeper = SafekeeperNode::from_env(env, node);
    1305            4 :         if let Err(e) = safekeeper.stop(immediate) {
    1306            0 :             eprintln!("safekeeper {} stop failed: {:#}", safekeeper.id, e);
    1307            4 :         }
    1308              :     }
    1309              : 
    1310            3 :     if let Err(e) = broker::stop_broker_process(env) {
    1311            0 :         eprintln!("neon broker stop failed: {e:#}");
    1312            3 :     }
    1313              : 
    1314            3 :     if env.control_plane_api.is_some() {
    1315            3 :         let attachment_service = AttachmentService::from_env(env);
    1316            6 :         if let Err(e) = attachment_service.stop(immediate).await {
    1317            0 :             eprintln!("attachment service stop failed: {e:#}");
    1318            3 :         }
    1319            0 :     }
    1320            3 : }
    1321              : 
    1322         6108 : fn cli() -> Command {
    1323         6108 :     let branch_name_arg = Arg::new("branch-name")
    1324         6108 :         .long("branch-name")
    1325         6108 :         .help("Name of the branch to be created or used as an alias for other services")
    1326         6108 :         .required(false);
    1327         6108 : 
    1328         6108 :     let endpoint_id_arg = Arg::new("endpoint_id")
    1329         6108 :         .help("Postgres endpoint id")
    1330         6108 :         .required(false);
    1331         6108 : 
    1332         6108 :     let safekeeper_id_arg = Arg::new("id").help("safekeeper id").required(false);
    1333         6108 : 
    1334         6108 :     // --id, when using a pageserver command
    1335         6108 :     let pageserver_id_arg = Arg::new("pageserver-id")
    1336         6108 :         .long("id")
    1337         6108 :         .global(true)
    1338         6108 :         .help("pageserver id")
    1339         6108 :         .required(false);
    1340         6108 :     // --pageserver-id when using a non-pageserver command
    1341         6108 :     let endpoint_pageserver_id_arg = Arg::new("endpoint-pageserver-id")
    1342         6108 :         .long("pageserver-id")
    1343         6108 :         .required(false);
    1344         6108 : 
    1345         6108 :     let safekeeper_extra_opt_arg = Arg::new("safekeeper-extra-opt")
    1346         6108 :         .short('e')
    1347         6108 :         .long("safekeeper-extra-opt")
    1348         6108 :         .num_args(1)
    1349         6108 :         .action(ArgAction::Append)
    1350         6108 :         .help("Additional safekeeper invocation options, e.g. -e=--http-auth-public-key-path=foo")
    1351         6108 :         .required(false);
    1352         6108 : 
    1353         6108 :     let tenant_id_arg = Arg::new("tenant-id")
    1354         6108 :         .long("tenant-id")
    1355         6108 :         .help("Tenant id. Represented as a hexadecimal string 32 symbols length")
    1356         6108 :         .required(false);
    1357         6108 : 
    1358         6108 :     let timeline_id_arg = Arg::new("timeline-id")
    1359         6108 :         .long("timeline-id")
    1360         6108 :         .help("Timeline id. Represented as a hexadecimal string 32 symbols length")
    1361         6108 :         .required(false);
    1362         6108 : 
    1363         6108 :     let pg_version_arg = Arg::new("pg-version")
    1364         6108 :         .long("pg-version")
    1365         6108 :         .help("Postgres version to use for the initial tenant")
    1366         6108 :         .required(false)
    1367         6108 :         .value_parser(value_parser!(u32))
    1368         6108 :         .default_value(DEFAULT_PG_VERSION);
    1369         6108 : 
    1370         6108 :     let pg_port_arg = Arg::new("pg-port")
    1371         6108 :         .long("pg-port")
    1372         6108 :         .required(false)
    1373         6108 :         .value_parser(value_parser!(u16))
    1374         6108 :         .value_name("pg-port");
    1375         6108 : 
    1376         6108 :     let http_port_arg = Arg::new("http-port")
    1377         6108 :         .long("http-port")
    1378         6108 :         .required(false)
    1379         6108 :         .value_parser(value_parser!(u16))
    1380         6108 :         .value_name("http-port");
    1381         6108 : 
    1382         6108 :     let safekeepers_arg = Arg::new("safekeepers")
    1383         6108 :         .long("safekeepers")
    1384         6108 :         .required(false)
    1385         6108 :         .value_name("safekeepers");
    1386         6108 : 
    1387         6108 :     let stop_mode_arg = Arg::new("stop-mode")
    1388         6108 :         .short('m')
    1389         6108 :         .value_parser(["fast", "immediate"])
    1390         6108 :         .default_value("fast")
    1391         6108 :         .help("If 'immediate', don't flush repository data at shutdown")
    1392         6108 :         .required(false)
    1393         6108 :         .value_name("stop-mode");
    1394         6108 : 
    1395         6108 :     let pageserver_config_args = Arg::new("pageserver-config-override")
    1396         6108 :         .long("pageserver-config-override")
    1397         6108 :         .num_args(1)
    1398         6108 :         .action(ArgAction::Append)
    1399         6108 :         .help("Additional pageserver's configuration options or overrides, refer to pageserver's 'config-override' CLI parameter docs for more")
    1400         6108 :         .required(false);
    1401         6108 : 
    1402         6108 :     let remote_ext_config_args = Arg::new("remote-ext-config")
    1403         6108 :         .long("remote-ext-config")
    1404         6108 :         .num_args(1)
    1405         6108 :         .help("Configure the remote extensions storage proxy gateway to request for extensions.")
    1406         6108 :         .required(false);
    1407         6108 : 
    1408         6108 :     let lsn_arg = Arg::new("lsn")
    1409         6108 :         .long("lsn")
    1410         6108 :         .help("Specify Lsn on the timeline to start from. By default, end of the timeline would be used.")
    1411         6108 :         .required(false);
    1412         6108 : 
    1413         6108 :     let hot_standby_arg = Arg::new("hot-standby")
    1414         6108 :         .value_parser(value_parser!(bool))
    1415         6108 :         .long("hot-standby")
    1416         6108 :         .help("If set, the node will be a hot replica on the specified timeline")
    1417         6108 :         .required(false);
    1418         6108 : 
    1419         6108 :     let force_arg = Arg::new("force")
    1420         6108 :         .value_parser(value_parser!(InitForceMode))
    1421         6108 :         .long("force")
    1422         6108 :         .default_value(
    1423         6108 :             InitForceMode::MustNotExist
    1424         6108 :                 .to_possible_value()
    1425         6108 :                 .unwrap()
    1426         6108 :                 .get_name()
    1427         6108 :                 .to_owned(),
    1428         6108 :         )
    1429         6108 :         .help("Force initialization even if the repository is not empty")
    1430         6108 :         .required(false);
    1431         6108 : 
    1432         6108 :     let num_pageservers_arg = Arg::new("num-pageservers")
    1433         6108 :         .value_parser(value_parser!(u16))
    1434         6108 :         .long("num-pageservers")
    1435         6108 :         .help("How many pageservers to create (default 1)")
    1436         6108 :         .required(false)
    1437         6108 :         .default_value("1");
    1438         6108 : 
    1439         6108 :     Command::new("Neon CLI")
    1440         6108 :         .arg_required_else_help(true)
    1441         6108 :         .version(GIT_VERSION)
    1442         6108 :         .subcommand(
    1443         6108 :             Command::new("init")
    1444         6108 :                 .about("Initialize a new Neon repository, preparing configs for services to start with")
    1445         6108 :                 .arg(pageserver_config_args.clone())
    1446         6108 :                 .arg(num_pageservers_arg.clone())
    1447         6108 :                 .arg(
    1448         6108 :                     Arg::new("config")
    1449         6108 :                         .long("config")
    1450         6108 :                         .required(false)
    1451         6108 :                         .value_parser(value_parser!(PathBuf))
    1452         6108 :                         .value_name("config"),
    1453         6108 :                 )
    1454         6108 :                 .arg(pg_version_arg.clone())
    1455         6108 :                 .arg(force_arg)
    1456         6108 :         )
    1457         6108 :         .subcommand(
    1458         6108 :             Command::new("timeline")
    1459         6108 :             .about("Manage timelines")
    1460         6108 :             .subcommand(Command::new("list")
    1461         6108 :                 .about("List all timelines, available to this pageserver")
    1462         6108 :                 .arg(tenant_id_arg.clone()))
    1463         6108 :             .subcommand(Command::new("branch")
    1464         6108 :                 .about("Create a new timeline, using another timeline as a base, copying its data")
    1465         6108 :                 .arg(tenant_id_arg.clone())
    1466         6108 :                 .arg(branch_name_arg.clone())
    1467         6108 :                 .arg(Arg::new("ancestor-branch-name").long("ancestor-branch-name")
    1468         6108 :                     .help("Use last Lsn of another timeline (and its data) as base when creating the new timeline. The timeline gets resolved by its branch name.").required(false))
    1469         6108 :                 .arg(Arg::new("ancestor-start-lsn").long("ancestor-start-lsn")
    1470         6108 :                     .help("When using another timeline as base, use a specific Lsn in it instead of the latest one").required(false)))
    1471         6108 :             .subcommand(Command::new("create")
    1472         6108 :                 .about("Create a new blank timeline")
    1473         6108 :                 .arg(tenant_id_arg.clone())
    1474         6108 :                 .arg(timeline_id_arg.clone())
    1475         6108 :                 .arg(branch_name_arg.clone())
    1476         6108 :                 .arg(pg_version_arg.clone())
    1477         6108 :             )
    1478         6108 :             .subcommand(Command::new("import")
    1479         6108 :                 .about("Import timeline from basebackup directory")
    1480         6108 :                 .arg(tenant_id_arg.clone())
    1481         6108 :                 .arg(timeline_id_arg.clone())
    1482         6108 :                 .arg(Arg::new("node-name").long("node-name")
    1483         6108 :                     .help("Name to assign to the imported timeline"))
    1484         6108 :                 .arg(Arg::new("base-tarfile")
    1485         6108 :                     .long("base-tarfile")
    1486         6108 :                     .value_parser(value_parser!(PathBuf))
    1487         6108 :                     .help("Basebackup tarfile to import")
    1488         6108 :                 )
    1489         6108 :                 .arg(Arg::new("base-lsn").long("base-lsn")
    1490         6108 :                     .help("Lsn the basebackup starts at"))
    1491         6108 :                 .arg(Arg::new("wal-tarfile")
    1492         6108 :                     .long("wal-tarfile")
    1493         6108 :                     .value_parser(value_parser!(PathBuf))
    1494         6108 :                     .help("Wal to add after base")
    1495         6108 :                 )
    1496         6108 :                 .arg(Arg::new("end-lsn").long("end-lsn")
    1497         6108 :                     .help("Lsn the basebackup ends at"))
    1498         6108 :                 .arg(pg_version_arg.clone())
    1499         6108 :             )
    1500         6108 :         ).subcommand(
    1501         6108 :             Command::new("tenant")
    1502         6108 :             .arg_required_else_help(true)
    1503         6108 :             .about("Manage tenants")
    1504         6108 :             .subcommand(Command::new("list"))
    1505         6108 :             .subcommand(Command::new("create")
    1506         6108 :                 .arg(tenant_id_arg.clone())
    1507         6108 :                 .arg(timeline_id_arg.clone().help("Use a specific timeline id when creating a tenant and its initial timeline"))
    1508         6108 :                 .arg(Arg::new("config").short('c').num_args(1).action(ArgAction::Append).required(false))
    1509         6108 :                 .arg(pg_version_arg.clone())
    1510         6108 :                 .arg(Arg::new("set-default").long("set-default").action(ArgAction::SetTrue).required(false)
    1511         6108 :                     .help("Use this tenant in future CLI commands where tenant_id is needed, but not specified"))
    1512         6108 :                 .arg(Arg::new("shard-count").value_parser(value_parser!(u8)).long("shard-count").action(ArgAction::Set).help("Number of shards in the new tenant (default 1)"))
    1513         6108 :                 .arg(Arg::new("shard-stripe-size").value_parser(value_parser!(u32)).long("shard-stripe-size").action(ArgAction::Set).help("Sharding stripe size in pages"))
    1514         6108 :                 )
    1515         6108 :             .subcommand(Command::new("set-default").arg(tenant_id_arg.clone().required(true))
    1516         6108 :                 .about("Set a particular tenant as default in future CLI commands where tenant_id is needed, but not specified"))
    1517         6108 :             .subcommand(Command::new("config")
    1518         6108 :                 .arg(tenant_id_arg.clone())
    1519         6108 :                 .arg(Arg::new("config").short('c').num_args(1).action(ArgAction::Append).required(false)))
    1520         6108 :             .subcommand(Command::new("migrate")
    1521         6108 :                 .about("Migrate a tenant from one pageserver to another")
    1522         6108 :                 .arg(tenant_id_arg.clone())
    1523         6108 :                 .arg(pageserver_id_arg.clone()))
    1524         6108 :             .subcommand(Command::new("status")
    1525         6108 :                 .about("Human readable summary of the tenant's shards and attachment locations")
    1526         6108 :                 .arg(tenant_id_arg.clone()))
    1527         6108 :         )
    1528         6108 :         .subcommand(
    1529         6108 :             Command::new("pageserver")
    1530         6108 :                 .arg_required_else_help(true)
    1531         6108 :                 .about("Manage pageserver")
    1532         6108 :                 .arg(pageserver_id_arg)
    1533         6108 :                 .subcommand(Command::new("status"))
    1534         6108 :                 .subcommand(Command::new("start")
    1535         6108 :                     .about("Start local pageserver")
    1536         6108 :                     .arg(pageserver_config_args.clone()).arg(Arg::new("register")
    1537         6108 :                     .long("register")
    1538         6108 :                     .default_value("true").required(false)
    1539         6108 :                     .value_parser(value_parser!(bool))
    1540         6108 :                     .value_name("register"))
    1541         6108 :                 )
    1542         6108 :                 .subcommand(Command::new("stop")
    1543         6108 :                     .about("Stop local pageserver")
    1544         6108 :                     .arg(stop_mode_arg.clone())
    1545         6108 :                 )
    1546         6108 :                 .subcommand(Command::new("restart")
    1547         6108 :                     .about("Restart local pageserver")
    1548         6108 :                     .arg(pageserver_config_args.clone())
    1549         6108 :                 )
    1550         6108 :                 .subcommand(Command::new("set-state")
    1551         6108 :                     .arg(Arg::new("availability").value_parser(value_parser!(NodeAvailability)).long("availability").action(ArgAction::Set).help("Availability state: offline,active"))
    1552         6108 :                     .arg(Arg::new("scheduling").value_parser(value_parser!(NodeSchedulingPolicy)).long("scheduling").action(ArgAction::Set).help("Scheduling state: draining,pause,filling,active"))
    1553         6108 :                     .about("Set scheduling or availability state of pageserver node")
    1554         6108 :                     .arg(pageserver_config_args.clone())
    1555         6108 :                 )
    1556         6108 :         )
    1557         6108 :         .subcommand(
    1558         6108 :             Command::new("attachment_service")
    1559         6108 :                 .arg_required_else_help(true)
    1560         6108 :                 .about("Manage attachment_service")
    1561         6108 :                 .subcommand(Command::new("start").about("Start local pageserver").arg(pageserver_config_args.clone()))
    1562         6108 :                 .subcommand(Command::new("stop").about("Stop local pageserver")
    1563         6108 :                             .arg(stop_mode_arg.clone()))
    1564         6108 :         )
    1565         6108 :         .subcommand(
    1566         6108 :             Command::new("safekeeper")
    1567         6108 :                 .arg_required_else_help(true)
    1568         6108 :                 .about("Manage safekeepers")
    1569         6108 :                 .subcommand(Command::new("start")
    1570         6108 :                             .about("Start local safekeeper")
    1571         6108 :                             .arg(safekeeper_id_arg.clone())
    1572         6108 :                             .arg(safekeeper_extra_opt_arg.clone())
    1573         6108 :                 )
    1574         6108 :                 .subcommand(Command::new("stop")
    1575         6108 :                             .about("Stop local safekeeper")
    1576         6108 :                             .arg(safekeeper_id_arg.clone())
    1577         6108 :                             .arg(stop_mode_arg.clone())
    1578         6108 :                 )
    1579         6108 :                 .subcommand(Command::new("restart")
    1580         6108 :                             .about("Restart local safekeeper")
    1581         6108 :                             .arg(safekeeper_id_arg)
    1582         6108 :                             .arg(stop_mode_arg.clone())
    1583         6108 :                             .arg(safekeeper_extra_opt_arg)
    1584         6108 :                 )
    1585         6108 :         )
    1586         6108 :         .subcommand(
    1587         6108 :             Command::new("endpoint")
    1588         6108 :                 .arg_required_else_help(true)
    1589         6108 :                 .about("Manage postgres instances")
    1590         6108 :                 .subcommand(Command::new("list").arg(tenant_id_arg.clone()))
    1591         6108 :                 .subcommand(Command::new("create")
    1592         6108 :                     .about("Create a compute endpoint")
    1593         6108 :                     .arg(endpoint_id_arg.clone())
    1594         6108 :                     .arg(branch_name_arg.clone())
    1595         6108 :                     .arg(tenant_id_arg.clone())
    1596         6108 :                     .arg(lsn_arg.clone())
    1597         6108 :                     .arg(pg_port_arg.clone())
    1598         6108 :                     .arg(http_port_arg.clone())
    1599         6108 :                     .arg(endpoint_pageserver_id_arg.clone())
    1600         6108 :                     .arg(
    1601         6108 :                         Arg::new("config-only")
    1602         6108 :                             .help("Don't do basebackup, create endpoint directory with only config files")
    1603         6108 :                             .long("config-only")
    1604         6108 :                             .required(false))
    1605         6108 :                     .arg(pg_version_arg.clone())
    1606         6108 :                     .arg(hot_standby_arg.clone())
    1607         6108 :                 )
    1608         6108 :                 .subcommand(Command::new("start")
    1609         6108 :                     .about("Start postgres.\n If the endpoint doesn't exist yet, it is created.")
    1610         6108 :                     .arg(endpoint_id_arg.clone())
    1611         6108 :                     .arg(endpoint_pageserver_id_arg.clone())
    1612         6108 :                     .arg(safekeepers_arg)
    1613         6108 :                     .arg(remote_ext_config_args)
    1614         6108 :                 )
    1615         6108 :                 .subcommand(Command::new("reconfigure")
    1616         6108 :                             .about("Reconfigure the endpoint")
    1617         6108 :                             .arg(endpoint_pageserver_id_arg)
    1618         6108 :                             .arg(endpoint_id_arg.clone())
    1619         6108 :                             .arg(tenant_id_arg.clone())
    1620         6108 :                 )
    1621         6108 :                 .subcommand(
    1622         6108 :                     Command::new("stop")
    1623         6108 :                     .arg(endpoint_id_arg)
    1624         6108 :                     .arg(
    1625         6108 :                         Arg::new("destroy")
    1626         6108 :                             .help("Also delete data directory (now optional, should be default in future)")
    1627         6108 :                             .long("destroy")
    1628         6108 :                             .action(ArgAction::SetTrue)
    1629         6108 :                             .required(false)
    1630         6108 :                         )
    1631         6108 :                 )
    1632         6108 : 
    1633         6108 :         )
    1634         6108 :         .subcommand(
    1635         6108 :             Command::new("mappings")
    1636         6108 :                 .arg_required_else_help(true)
    1637         6108 :                 .about("Manage neon_local branch name mappings")
    1638         6108 :                 .subcommand(
    1639         6108 :                     Command::new("map")
    1640         6108 :                         .about("Create new mapping which cannot exist already")
    1641         6108 :                         .arg(branch_name_arg.clone())
    1642         6108 :                         .arg(tenant_id_arg.clone())
    1643         6108 :                         .arg(timeline_id_arg.clone())
    1644         6108 :                 )
    1645         6108 :         )
    1646         6108 :         // Obsolete old name for 'endpoint'. We now just print an error if it's used.
    1647         6108 :         .subcommand(
    1648         6108 :             Command::new("pg")
    1649         6108 :                 .hide(true)
    1650         6108 :                 .arg(Arg::new("ignore-rest").allow_hyphen_values(true).num_args(0..).required(false))
    1651         6108 :                 .trailing_var_arg(true)
    1652         6108 :         )
    1653         6108 :         .subcommand(
    1654         6108 :             Command::new("start")
    1655         6108 :                 .about("Start page server and safekeepers")
    1656         6108 :                 .arg(pageserver_config_args)
    1657         6108 :         )
    1658         6108 :         .subcommand(
    1659         6108 :             Command::new("stop")
    1660         6108 :                 .about("Stop page server and safekeepers")
    1661         6108 :                 .arg(stop_mode_arg)
    1662         6108 :         )
    1663         6108 : }
    1664              : 
    1665            2 : #[test]
    1666            2 : fn verify_cli() {
    1667            2 :     cli().debug_assert();
    1668            2 : }
        

Generated by: LCOV version 2.1-beta