Refactor 'zenith' CLI subcommand handling

Also fixes 'zenith safekeeper restart -m immediate'. The stop-mode was
previously ignored.

zenith/src/main.rs

@@ -203,96 +203,38 @@ fn main() -> Result<()> {
         )
         .get_matches();
 
-    // Create config file
-    if let ("init", Some(init_match)) = matches.subcommand() {
-        let toml_file: String = if let Some(config_path) = init_match.value_of("config") {
-            // load and parse the file
-            std::fs::read_to_string(std::path::Path::new(config_path))
-                .with_context(|| format!("Could not read configuration file \"{}\"", config_path))?
-        } else {
-            // Built-in default config
-            default_conf()
+    let (sub_name, sub_args) = matches.subcommand();
+    let sub_args = sub_args.expect("no subcommand");
+
+    // Check for 'zenith init' command first.
+    let subcmd_result = if sub_name == "init" {
+        handle_init(sub_args)
+    } else {
+        // all other commands need an existing config
+        let env = match LocalEnv::load_config() {
+            Ok(conf) => conf,
+            Err(e) => {
+                eprintln!("Error loading config: {}", e);
+                exit(1);
+            }
         };
 
-        let mut env = LocalEnv::create_config(&toml_file)
-            .with_context(|| "Failed to create zenith configuration")?;
-        env.init()
-            .with_context(|| "Failed to initialize zenith repository")?;
+        match sub_name {
+            "tenant" => handle_tenant(sub_args, &env),
+            "branch" => handle_branch(sub_args, &env),
+            "start" => handle_start_all(sub_args, &env),
+            "stop" => handle_stop_all(sub_args, &env),
+            "pageserver" => handle_pageserver(sub_args, &env),
+            "pg" => handle_pg(sub_args, &env),
+            "safekeeper" => handle_safekeeper(sub_args, &env),
+            _ => bail!("unexpected subcommand {}", sub_name),
+        }
+    };
+    if let Err(e) = subcmd_result {
+        eprintln!("command failed: {}", e);
+        exit(1);
     }
 
-    // all other commands would need config
-    let env = match LocalEnv::load_config() {
-        Ok(conf) => conf,
-        Err(e) => {
-            eprintln!("Error loading config: {}", e);
-            exit(1);
-        }
-    };
-
-    match matches.subcommand() {
-        ("init", Some(_sub_m)) => {
-            // The options were handled above already
-            let pageserver = PageServerNode::from_env(&env);
-            if let Err(e) = pageserver.init(
-                // default_tenantid was generated by the `env.init()` call above
-                Some(&env.default_tenantid.unwrap().to_string()),
-            ) {
-                eprintln!("pageserver init failed: {}", e);
-                exit(1);
-            }
-        }
-        ("tenant", Some(args)) => {
-            if let Err(e) = handle_tenant(args, &env) {
-                eprintln!("tenant command failed: {}", e);
-                exit(1);
-            }
-        }
-
-        ("branch", Some(sub_args)) => {
-            if let Err(e) = handle_branch(sub_args, &env) {
-                eprintln!("branch command failed: {}", e);
-                exit(1);
-            }
-        }
-
-        ("start", Some(sub_match)) => {
-            if let Err(e) = handle_start_all(sub_match, &env) {
-                eprintln!("start command failed: {}", e);
-                exit(1);
-            }
-        }
-
-        ("stop", Some(sub_match)) => {
-            if let Err(e) = handle_stop_all(sub_match, &env) {
-                eprintln!("stop command failed: {}", e);
-                exit(1);
-            }
-        }
-
-        ("pageserver", Some(sub_match)) => {
-            if let Err(e) = handle_pageserver(sub_match, &env) {
-                eprintln!("pageserver command failed: {}", e);
-                exit(1);
-            }
-        }
-
-        ("pg", Some(pg_match)) => {
-            if let Err(e) = handle_pg(pg_match, &env) {
-                eprintln!("pg operation failed: {:?}", e);
-                exit(1);
-            }
-        }
-
-        ("safekeeper", Some(sub_match)) => {
-            if let Err(e) = handle_safekeeper(sub_match, &env) {
-                eprintln!("safekeeper command failed: {}", e);
-                exit(1);
-            }
-        }
-
-        _ => {}
-    };
-
     Ok(())
 }
 
@@ -438,6 +380,35 @@ fn get_tenantid(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<ZTe
     }
 }
 
+fn handle_init(init_match: &ArgMatches) -> Result<()> {
+    // Create config file
+    let toml_file: String = if let Some(config_path) = init_match.value_of("config") {
+        // load and parse the file
+        std::fs::read_to_string(std::path::Path::new(config_path))
+            .with_context(|| format!("Could not read configuration file \"{}\"", config_path))?
+    } else {
+        // Built-in default config
+        default_conf()
+    };
+
+    let mut env = LocalEnv::create_config(&toml_file)
+        .with_context(|| "Failed to create zenith configuration")?;
+    env.init()
+        .with_context(|| "Failed to initialize zenith repository")?;
+
+    // Call 'pageserver init'.
+    let pageserver = PageServerNode::from_env(&env);
+    if let Err(e) = pageserver.init(
+        // default_tenantid was generated by the `env.init()` call above
+        Some(&env.default_tenantid.unwrap().to_string()),
+    ) {
+        eprintln!("pageserver init failed: {}", e);
+        exit(1);
+    }
+
+    Ok(())
+}
+
 fn handle_tenant(tenant_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
     let pageserver = PageServerNode::from_env(env);
     match tenant_match.subcommand() {
@@ -486,10 +457,11 @@ fn handle_branch(branch_match: &ArgMatches, env: &local_env::LocalEnv) -> Result
 fn handle_pg(pg_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
     let mut cplane = ComputeControlPlane::load(env.clone())?;
 
-    match pg_match.subcommand() {
-        ("list", Some(list_match)) => {
-            let tenantid = get_tenantid(list_match, env)?;
+    // All subcommands take an optional --tenantid option
+    let tenantid = get_tenantid(pg_match, env)?;
 
+    match pg_match.subcommand() {
+        ("list", Some(_list_match)) => {
             let branch_infos = get_branch_infos(env, &tenantid).unwrap_or_else(|e| {
                 eprintln!("Failed to load branch info: {}", e);
                 HashMap::new()
@@ -520,7 +492,6 @@ fn handle_pg(pg_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
             }
         }
         ("create", Some(create_match)) => {
-            let tenantid = get_tenantid(create_match, env)?;
             let node_name = create_match.value_of("node").unwrap_or("main");
             let timeline_name = create_match.value_of("timeline").unwrap_or(node_name);
 
@@ -531,7 +502,6 @@ fn handle_pg(pg_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
             cplane.new_node(tenantid, node_name, timeline_name, port)?;
         }
         ("start", Some(start_match)) => {
-            let tenantid = get_tenantid(start_match, env)?;
             let node_name = start_match.value_of("node").unwrap_or("main");
             let timeline_name = start_match.value_of("timeline");
 
@@ -572,7 +542,6 @@ fn handle_pg(pg_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
             }
         }
         ("stop", Some(stop_match)) => {
-            let tenantid = get_tenantid(stop_match, env)?;
             let node_name = stop_match.value_of("node").unwrap_or("main");
             let destroy = stop_match.is_present("destroy");
 
@@ -636,26 +605,23 @@ fn get_safekeeper(env: &local_env::LocalEnv, name: &str) -> Result<SafekeeperNod
 }
 
 fn handle_safekeeper(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
-    match sub_match.subcommand() {
-        ("start", Some(sub_match)) => {
-            let node_name = sub_match
-                .value_of("node")
-                .unwrap_or(DEFAULT_SAFEKEEPER_NAME);
-            let safekeeper = get_safekeeper(env, node_name)?;
+    let (sub_name, sub_args) = sub_match.subcommand();
+    let sub_args = sub_args.expect("no safekeeper subcommand");
 
+    // All the commands take an optional safekeeper name argument
+    let node_name = sub_args.value_of("node").unwrap_or(DEFAULT_SAFEKEEPER_NAME);
+    let safekeeper = get_safekeeper(env, node_name)?;
+
+    match sub_name {
+        "start" => {
             if let Err(e) = safekeeper.start() {
                 eprintln!("safekeeper start failed: {}", e);
                 exit(1);
             }
         }
 
-        ("stop", Some(sub_match)) => {
-            let node_name = sub_match
-                .value_of("node")
-                .unwrap_or(DEFAULT_SAFEKEEPER_NAME);
-            let immediate = sub_match.value_of("stop-mode") == Some("immediate");
-
-            let safekeeper = get_safekeeper(env, node_name)?;
+        "stop" => {
+            let immediate = sub_args.value_of("stop-mode") == Some("immediate");
 
             if let Err(e) = safekeeper.stop(immediate) {
                 eprintln!("safekeeper stop failed: {}", e);
@@ -663,15 +629,10 @@ fn handle_safekeeper(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> Resul
             }
         }
 
-        ("restart", Some(sub_match)) => {
-            let node_name = sub_match
-                .value_of("node")
-                .unwrap_or(DEFAULT_SAFEKEEPER_NAME);
+        "restart" => {
+            let immediate = sub_args.value_of("stop-mode") == Some("immediate");
 
-            let safekeeper = get_safekeeper(env, node_name)?;
-
-            //TODO what shutdown strategy should we use here?
-            if let Err(e) = safekeeper.stop(false) {
+            if let Err(e) = safekeeper.stop(immediate) {
                 eprintln!("safekeeper stop failed: {}", e);
                 exit(1);
             }
@@ -682,7 +643,9 @@ fn handle_safekeeper(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> Resul
             }
         }
 
-        _ => {}
+        _ => {
+            bail!("Unexpected safekeeper subcommand '{}'", sub_name)
+        }
     }
     Ok(())
 }