From 23be5021f8d0dabf8b84b075b2d21385b6b06bcc Mon Sep 17 00:00:00 2001
From: Patrick Insinger <patrickinsinger@gmail.com>
Date: Wed, 12 May 2021 14:17:52 -0400
Subject: [PATCH] Remote CLI command

---
 control_plane/src/lib.rs       |  1 +
 control_plane/src/local_env.rs |  6 +++-
 control_plane/src/remotes.rs   | 25 ++++++++++++++++
 zenith/src/main.rs             | 55 +++++++++++++++++++++++++++++++---
 4 files changed, 82 insertions(+), 5 deletions(-)
 create mode 100644 control_plane/src/remotes.rs

diff --git a/control_plane/src/lib.rs b/control_plane/src/lib.rs
index 7ce1cffb9c..cd6636aada 100644
--- a/control_plane/src/lib.rs
+++ b/control_plane/src/lib.rs
@@ -12,6 +12,7 @@ use std::path::Path;
 
 pub mod compute;
 pub mod local_env;
+pub mod remotes;
 pub mod storage;
 
 /// Read a PID file
diff --git a/control_plane/src/local_env.rs b/control_plane/src/local_env.rs
index 2ee377e0fb..c2ad533a80 100644
--- a/control_plane/src/local_env.rs
+++ b/control_plane/src/local_env.rs
@@ -6,11 +6,13 @@
 //
 use anyhow::{anyhow, Result};
 use serde::{Deserialize, Serialize};
-use std::env;
+use std::{collections::BTreeMap, env};
 use std::fs;
 use std::path::PathBuf;
 use url::Url;
 
+use crate::remotes;
+
 //
 // This data structures represent deserialized zenith CLI config
 //
@@ -128,6 +130,8 @@ pub fn init(remote_pageserver: Option<&str>) -> Result<()> {
 
     let toml = toml::to_string(&conf)?;
     fs::write(conf.base_data_dir.join("config"), toml)?;
+    // initialize remotes file
+    remotes::save_remotes(&conf, &BTreeMap::default())?;
 
     Ok(())
 }
diff --git a/control_plane/src/remotes.rs b/control_plane/src/remotes.rs
new file mode 100644
index 0000000000..e1d2238357
--- /dev/null
+++ b/control_plane/src/remotes.rs
@@ -0,0 +1,25 @@
+//! Utility module for managing the remote pageserver file
+//! Currently a TOML file is used to map remote names to URIs
+
+use anyhow::Result;
+use std::collections::BTreeMap;
+use std::{fs, path::PathBuf};
+
+use crate::local_env::LocalEnv;
+
+pub type Remotes = BTreeMap<String, String>;
+
+fn remotes_path(local_env: &LocalEnv) -> PathBuf {
+    local_env.base_data_dir.join("remotes")
+}
+
+pub fn load_remotes(local_env: &LocalEnv) -> Result<Remotes> {
+    let remotes_str = fs::read_to_string(remotes_path(local_env))?;
+    Ok(toml::from_str(&remotes_str)?)
+}
+
+pub fn save_remotes(local_env: &LocalEnv, remotes: &Remotes) -> Result<()> {
+    let remotes_str = toml::to_string_pretty(remotes)?;
+    fs::write(remotes_path(local_env), remotes_str)?;
+    Ok(())
+}
diff --git a/zenith/src/main.rs b/zenith/src/main.rs
index a847ad888e..02bc36c3a5 100644
--- a/zenith/src/main.rs
+++ b/zenith/src/main.rs
@@ -1,12 +1,13 @@
+use std::{collections::btree_map::Entry, str::FromStr};
 use anyhow::Result;
-use anyhow::{anyhow, Context};
+use anyhow::{anyhow, Context, bail};
 use clap::{App, Arg, ArgMatches, SubCommand};
 use std::collections::HashMap;
 use std::process::exit;
-
-use control_plane::compute::ComputeControlPlane;
-use control_plane::local_env;
 use control_plane::storage::PageServerNode;
+use control_plane::{compute::ComputeControlPlane, local_env};
+use control_plane::{local_env::LocalEnv, remotes};
+
 use pageserver::{branches::BranchInfo, ZTimelineId};
 use zenith_utils::lsn::Lsn;
 
@@ -61,6 +62,20 @@ fn main() -> Result<()> {
                 .subcommand(SubCommand::with_name("stop").arg(name_arg.clone()))
                 .subcommand(SubCommand::with_name("destroy").arg(name_arg.clone())),
         )
+        .subcommand(
+            SubCommand::with_name("remote")
+                .about("Manage remote pagerservers")
+                .subcommand(
+                    SubCommand::with_name("add")
+                        .about("Add a new remote pageserver")
+                        .arg(Arg::with_name("name").required(true))
+                        .arg(
+                            Arg::with_name("url")
+                                .help("PostgreSQL connection URI")
+                                .required(true),
+                        ),
+                ),
+        )
         .get_matches();
 
     // Create config file
@@ -146,6 +161,13 @@ fn main() -> Result<()> {
                 exit(1);
             }
         }
+
+        ("remote", Some(remote_match)) => {
+            if let Err(e) = handle_remote(remote_match, &env) {
+                eprintln!("remote operation failed: {}", e);
+                exit(1);
+            }
+        }
         _ => {}
     };
 
@@ -233,3 +255,28 @@ fn handle_pg(pg_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<()> {
 
     Ok(())
 }
+
+fn handle_remote(remote_match: &ArgMatches, local_env: &LocalEnv) -> Result<()> {
+    let mut remotes = remotes::load_remotes(local_env)?;
+    match remote_match.subcommand() {
+        ("add", Some(args)) => {
+            let name = args.value_of("name").unwrap();
+            let url = args.value_of("url").unwrap();
+
+            // validate the URL
+            postgres::Config::from_str(url)?;
+
+            match remotes.entry(name.to_string()) {
+                Entry::Vacant(vacant) => {
+                    vacant.insert(url.to_string());
+                }
+                Entry::Occupied(_) => bail!("origin '{}' already exists", name),
+            }
+
+            remotes::save_remotes(local_env, &remotes)?;
+        }
+        _ => bail!("unknown command"),
+    }
+
+    Ok(())
+}