Store branch name mappings in separate branches.toml file

control_plane/src/bin/neon_local.rs

@@ -90,8 +90,11 @@ fn main() -> Result<()> {
         handle_init(sub_args).map(Some)
     } else {
         // all other commands need an existing config
-        let mut env =
-            LocalEnv::load_config(&local_env::base_path()).context("Error loading config")?;
+        let base_path = local_env::base_path();
+        let branch_mappings_path = local_env::branch_mappings_path();
+
+        let mut env = LocalEnv::load_config(&base_path, Some(&branch_mappings_path))
+            .context("Error loading config")?;
         let original_env = env.clone();
 
         let rt = tokio::runtime::Builder::new_current_thread()
@@ -264,7 +267,7 @@ async fn get_timeline_infos(
 fn get_tenant_id(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> anyhow::Result<TenantId> {
     if let Some(tenant_id_from_arguments) = parse_tenant_id(sub_match).transpose() {
         tenant_id_from_arguments
-    } else if let Some(default_id) = env.default_tenant_id {
+    } else if let Some(default_id) = env.branch_mappings.default_tenant_id {
         Ok(default_id)
     } else {
         anyhow::bail!("No tenant id. Use --tenant-id, or set a default tenant");
@@ -278,7 +281,7 @@ fn get_tenant_shard_id(
 ) -> anyhow::Result<TenantShardId> {
     if let Some(tenant_id_from_arguments) = parse_tenant_shard_id(sub_match).transpose() {
         tenant_id_from_arguments
-    } else if let Some(default_id) = env.default_tenant_id {
+    } else if let Some(default_id) = env.branch_mappings.default_tenant_id {
         Ok(TenantShardId::unsharded(default_id))
     } else {
         anyhow::bail!("No tenant shard id. Use --tenant-id, or set a default tenant");
@@ -360,7 +363,6 @@ fn handle_init(init_match: &ArgMatches) -> anyhow::Result<LocalEnv> {
                 .collect(),
             pg_distrib_dir: None,
             neon_distrib_dir: None,
-            default_tenant_id: TenantId::from_array(std::array::from_fn(|_| 0)),
             storage_controller: None,
             control_plane_compute_hook_api: None,
         }
@@ -368,8 +370,12 @@ fn handle_init(init_match: &ArgMatches) -> anyhow::Result<LocalEnv> {
 
     LocalEnv::init(init_conf, force)
         .context("materialize initial neon_local environment on disk")?;
-    Ok(LocalEnv::load_config(&local_env::base_path())
-        .expect("freshly written config should be loadable"))
+    let base_path = local_env::base_path();
+    let branch_mappings_path = local_env::branch_mappings_path();
+    let env = LocalEnv::load_config(&base_path, Some(&branch_mappings_path))
+        .expect("freshly written config should be loadable");
+
+    Ok(env)
 }
 
 /// The default pageserver is the one where CLI tenant/timeline operations are sent by default.
@@ -525,14 +531,14 @@ async fn handle_tenant(
 
             if create_match.get_flag("set-default") {
                 println!("Setting tenant {tenant_id} as a default one");
-                env.default_tenant_id = Some(tenant_id);
+                env.branch_mappings.default_tenant_id = Some(tenant_id);
             }
         }
         Some(("set-default", set_default_match)) => {
             let tenant_id =
                 parse_tenant_id(set_default_match)?.context("No tenant id specified")?;
             println!("Setting tenant {tenant_id} as a default one");
-            env.default_tenant_id = Some(tenant_id);
+            env.branch_mappings.default_tenant_id = Some(tenant_id);
         }
         Some(("config", create_match)) => {
             let tenant_id = get_tenant_id(create_match, env)?;
@@ -693,6 +699,7 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re
         Some(ep_subcommand_data) => ep_subcommand_data,
         None => bail!("no endpoint subcommand provided"),
     };
+
     let mut cplane = ComputeControlPlane::load(env.clone())?;
 
     match sub_name {
@@ -1360,6 +1367,7 @@ async fn handle_stop_all(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> R
 
 async fn try_stop_all(env: &local_env::LocalEnv, immediate: bool) {
     // Stop all endpoints
+    // NOTE: This only knows about endpoints in the default endpoints dir
     match ComputeControlPlane::load(env.clone()) {
         Ok(cplane) => {
             for (_k, node) in cplane.endpoints {
@@ -1420,6 +1428,18 @@ fn cli() -> Command {
         .default_value("10s")
         .required(false);
 
+    let branch_mappings_arg = Arg::new("branch-mappings")
+        .long("branch-mappings")
+        .help("File holding all branch names. Default is <repo dir>/branches.toml")
+        .value_parser(value_parser!(PathBuf))
+        .required(false);
+
+    let endpoints_dir_arg = Arg::new("endpoints-dir")
+        .long("endpoints-dir")
+        .help("Path to directory holding all endpoints. Default is <repo dir>/endpoints")
+        .value_parser(value_parser!(PathBuf))
+        .required(false);
+
     let branch_name_arg = Arg::new("branch-name")
         .long("branch-name")
         .help("Name of the branch to be created or used as an alias for other services")
@@ -1560,6 +1580,8 @@ fn cli() -> Command {
     Command::new("Neon CLI")
         .arg_required_else_help(true)
         .version(GIT_VERSION)
+        .arg(branch_mappings_arg)
+        .arg(endpoints_dir_arg)
         .subcommand(
             Command::new("init")
                 .about("Initialize a new Neon repository, preparing configs for services to start with")
@@ -1571,7 +1593,6 @@ fn cli() -> Command {
                         .value_parser(value_parser!(PathBuf))
                         .value_name("config")
                 )
-                .arg(pg_version_arg.clone())
                 .arg(force_arg)
         )
         .subcommand(

control_plane/src/local_env.rs

@@ -46,6 +46,11 @@ pub struct LocalEnv {
     // must be an absolute path. If the env var is not set, $PWD/.neon is used.
     pub base_data_dir: PathBuf,
 
+    // Similarly, path to branch mappings file. Not stored in the config file but
+    // read from the NEON_BRANCH_MAPPINGS env variable. "None" means no mappings
+    // are loaded, and they cannot be saved either.
+    pub branch_name_mappings_path: Option<PathBuf>,
+
     // Path to postgres distribution. It's expected that "bin", "include",
     // "lib", "share" from postgres distribution are there. If at some point
     // in time we will be able to run against vanilla postgres we may split that
@@ -55,10 +60,6 @@ pub struct LocalEnv {
     // Path to pageserver binary.
     pub neon_distrib_dir: PathBuf,
 
-    // Default tenant ID to use with the 'neon_local' command line utility, when
-    // --tenant_id is not explicitly specified.
-    pub default_tenant_id: Option<TenantId>,
-
     // used to issue tokens during e.g pg start
     pub private_key_path: PathBuf,
 
@@ -82,11 +83,7 @@ pub struct LocalEnv {
     // storage controller's configuration.
     pub control_plane_compute_hook_api: Option<Url>,
 
-    /// Keep human-readable aliases in memory (and persist them to config), to hide ZId hex strings from the user.
-    // A `HashMap<String, HashMap<TenantId, TimelineId>>` would be more appropriate here,
-    // but deserialization into a generic toml object as `toml::Value::try_from` fails with an error.
-    // https://toml.io/en/v1.0.0 does not contain a concept of "a table inside another table".
-    pub branch_name_mappings: HashMap<String, Vec<(TenantId, TimelineId)>>,
+    pub branch_mappings: BranchMappings,
 }
 
 /// On-disk state stored in `.neon/config`.
@@ -95,7 +92,6 @@ pub struct LocalEnv {
 pub struct OnDiskConfig {
     pub pg_distrib_dir: PathBuf,
     pub neon_distrib_dir: PathBuf,
-    pub default_tenant_id: Option<TenantId>,
     pub private_key_path: PathBuf,
     pub broker: NeonBroker,
     pub storage_controller: NeonStorageControllerConf,
@@ -107,7 +103,20 @@ pub struct OnDiskConfig {
     pub safekeepers: Vec<SafekeeperConf>,
     pub control_plane_api: Option<Url>,
     pub control_plane_compute_hook_api: Option<Url>,
-    branch_name_mappings: HashMap<String, Vec<(TenantId, TimelineId)>>,
+}
+
+#[derive(PartialEq, Eq, Clone, Debug, Default, Serialize, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct BranchMappings {
+    // Default tenant ID to use with the 'neon_local' command line utility, when
+    // --tenant_id is not explicitly specified. This comes from the branches.
+    pub default_tenant_id: Option<TenantId>,
+
+    /// Keep human-readable aliases in memory (and persist them to config XXX), to hide ZId hex strings from the user.
+    // A `HashMap<String, HashMap<TenantId, TimelineId>>` would be more appropriate here,
+    // but deserialization into a generic toml object as `toml::Value::try_from` fails with an error.
+    // https://toml.io/en/v1.0.0 does not contain a concept of "a table inside another table".
+    pub mappings: HashMap<String, Vec<(TenantId, TimelineId)>>,
 }
 
 fn fail_if_pageservers_field_specified<'de, D>(_: D) -> Result<Vec<PageServerConf>, D::Error>
@@ -128,7 +137,6 @@ pub struct NeonLocalInitConf {
     pub pg_distrib_dir: Option<PathBuf>,
     // TODO: do we need this? Seems unused
     pub neon_distrib_dir: Option<PathBuf>,
-    pub default_tenant_id: TenantId,
     pub broker: NeonBroker,
     pub storage_controller: Option<NeonStorageControllerConf>,
     pub pageservers: Vec<NeonLocalInitPageserverConf>,
@@ -443,7 +451,8 @@ impl LocalEnv {
         timeline_id: TimelineId,
     ) -> anyhow::Result<()> {
         let existing_values = self
-            .branch_name_mappings
+            .branch_mappings
+            .mappings
             .entry(branch_name.clone())
             .or_default();
 
@@ -468,7 +477,8 @@ impl LocalEnv {
         branch_name: &str,
         tenant_id: TenantId,
     ) -> Option<TimelineId> {
-        self.branch_name_mappings
+        self.branch_mappings
+            .mappings
             .get(branch_name)?
             .iter()
             .find(|(mapped_tenant_id, _)| mapped_tenant_id == &tenant_id)
@@ -477,7 +487,8 @@ impl LocalEnv {
     }
 
     pub fn timeline_name_mappings(&self) -> HashMap<TenantTimelineId, String> {
-        self.branch_name_mappings
+        self.branch_mappings
+            .mappings
             .iter()
             .flat_map(|(name, tenant_timelines)| {
                 tenant_timelines.iter().map(|&(tenant_id, timeline_id)| {
@@ -488,7 +499,10 @@ impl LocalEnv {
     }
 
     ///  Construct `Self` from on-disk state.
-    pub fn load_config(repopath: &Path) -> anyhow::Result<Self> {
+    pub fn load_config(
+        repopath: &Path,
+        branch_name_mappings_path: Option<&Path>,
+    ) -> anyhow::Result<Self> {
         if !repopath.exists() {
             bail!(
                 "Neon config is not found in {}. You need to run 'neon_local init' first",
@@ -498,37 +512,43 @@ impl LocalEnv {
 
         // TODO: check that it looks like a neon repository
 
-        // load and parse file
+        // load and parse config file
         let config_file_contents = fs::read_to_string(repopath.join("config"))?;
         let on_disk_config: OnDiskConfig = toml::from_str(config_file_contents.as_str())?;
-        let mut env = {
-            let OnDiskConfig {
-                pg_distrib_dir,
-                neon_distrib_dir,
-                default_tenant_id,
-                private_key_path,
-                broker,
-                storage_controller,
-                pageservers,
-                safekeepers,
-                control_plane_api,
-                control_plane_compute_hook_api,
-                branch_name_mappings,
-            } = on_disk_config;
-            LocalEnv {
-                base_data_dir: repopath.to_owned(),
-                pg_distrib_dir,
-                neon_distrib_dir,
-                default_tenant_id,
-                private_key_path,
-                broker,
-                storage_controller,
-                pageservers,
-                safekeepers,
-                control_plane_api,
-                control_plane_compute_hook_api,
-                branch_name_mappings,
-            }
+        let OnDiskConfig {
+            pg_distrib_dir,
+            neon_distrib_dir,
+            private_key_path,
+            broker,
+            storage_controller,
+            pageservers,
+            safekeepers,
+            control_plane_api,
+            control_plane_compute_hook_api,
+        } = on_disk_config;
+
+        // load and parse "branches.toml" file
+        let branch_mappings = if let Some(path) = branch_name_mappings_path {
+            let contents = fs::read_to_string(path)
+                .context(format!("load branch mappings file {}", path.display()))?;
+            toml::from_str::<BranchMappings>(contents.as_str())?
+        } else {
+            BranchMappings::default()
+        };
+
+        let mut env = LocalEnv {
+            base_data_dir: repopath.to_owned(),
+            branch_name_mappings_path: branch_name_mappings_path.map(|p| p.to_owned()),
+            pg_distrib_dir,
+            neon_distrib_dir,
+            private_key_path,
+            broker,
+            storage_controller,
+            pageservers,
+            safekeepers,
+            control_plane_api,
+            control_plane_compute_hook_api,
+            branch_mappings,
         };
 
         // The source of truth for pageserver configuration is the pageserver.toml.
@@ -618,7 +638,6 @@ impl LocalEnv {
             &OnDiskConfig {
                 pg_distrib_dir: self.pg_distrib_dir.clone(),
                 neon_distrib_dir: self.neon_distrib_dir.clone(),
-                default_tenant_id: self.default_tenant_id,
                 private_key_path: self.private_key_path.clone(),
                 broker: self.broker.clone(),
                 storage_controller: self.storage_controller.clone(),
@@ -626,9 +645,20 @@ impl LocalEnv {
                 safekeepers: self.safekeepers.clone(),
                 control_plane_api: self.control_plane_api.clone(),
                 control_plane_compute_hook_api: self.control_plane_compute_hook_api.clone(),
-                branch_name_mappings: self.branch_name_mappings.clone(),
             },
-        )
+        )?;
+
+        if let Some(path) = &self.branch_name_mappings_path {
+            Self::persist_branches_impl(path, &self.branch_mappings)?;
+        } else {
+            if !self.branch_mappings.mappings.is_empty() {
+                tracing::warn!("command created a branch mapping, but it was not saved because no mappings file was configured")
+            } else if self.branch_mappings.default_tenant_id.is_some() {
+                tracing::warn!("command created a tenant default, but it was not saved because no mappings file was configured")
+            }
+        }
+
+        Ok(())
     }
 
     pub fn persist_config_impl(base_path: &Path, config: &OnDiskConfig) -> anyhow::Result<()> {
@@ -642,6 +672,19 @@ impl LocalEnv {
         })
     }
 
+    pub fn persist_branches_impl(
+        branch_name_mappings_path: &Path,
+        branch_mappings: &BranchMappings,
+    ) -> anyhow::Result<()> {
+        let content = &toml::to_string_pretty(branch_mappings)?;
+        fs::write(branch_name_mappings_path, content).with_context(|| {
+            format!(
+                "Failed to write branch information into path '{}'",
+                branch_name_mappings_path.display()
+            )
+        })
+    }
+
     // this function is used only for testing purposes in CLI e g generate tokens during init
     pub fn generate_auth_token(&self, claims: &Claims) -> anyhow::Result<String> {
         let private_key_path = self.get_private_key_path();
@@ -702,7 +745,6 @@ impl LocalEnv {
         let NeonLocalInitConf {
             pg_distrib_dir,
             neon_distrib_dir,
-            default_tenant_id,
             broker,
             storage_controller,
             pageservers,
@@ -746,9 +788,9 @@ impl LocalEnv {
         // TODO: refactor to avoid this, LocalEnv should only be constructed from on-disk state
         let env = LocalEnv {
             base_data_dir: base_path.clone(),
+            branch_name_mappings_path: Some(base_path.join("branches.toml")),
             pg_distrib_dir,
             neon_distrib_dir,
-            default_tenant_id: Some(default_tenant_id),
             private_key_path,
             broker,
             storage_controller: storage_controller.unwrap_or_default(),
@@ -756,10 +798,10 @@ impl LocalEnv {
             safekeepers,
             control_plane_api: control_plane_api.unwrap_or_default(),
             control_plane_compute_hook_api: control_plane_compute_hook_api.unwrap_or_default(),
-            branch_name_mappings: Default::default(),
+            branch_mappings: Default::default(),
         };
 
-        // create endpoints dir
+        // create the default endpoints dir
         fs::create_dir_all(env.endpoints_path())?;
 
         // create safekeeper dirs
@@ -806,6 +848,24 @@ pub fn base_path() -> PathBuf {
     path
 }
 
+pub fn branch_mappings_path() -> PathBuf {
+    let path = match std::env::var_os("NEON_BRANCH_MAPPINGS") {
+        Some(val) => {
+            let path = PathBuf::from(val);
+
+            // a relative path is relative to repo dir
+            if !path.is_absolute() {
+                base_path().join(path)
+            } else {
+                path
+            }
+        }
+        None => base_path().join("branches.toml"),
+    };
+    assert!(path.is_absolute());
+    path
+}
+
 /// Generate a public/private key pair for JWT authentication
 fn generate_auth_keys(private_key_path: &Path, public_key_path: &Path) -> anyhow::Result<()> {
     // Generate the key pair

storage_controller/src/compute_hook.rs

@@ -292,7 +292,7 @@ impl ComputeHook {
             );
             return Ok(());
         };
-        let env = match LocalEnv::load_config(repo_dir) {
+        let env = match LocalEnv::load_config(repo_dir, None) {
             Ok(e) => e,
             Err(e) => {
                 tracing::warn!("Couldn't load neon_local config, skipping compute update ({e})");

test_runner/fixtures/neon_fixtures.py

@@ -654,6 +654,10 @@ class NeonEnvBuilder:
         with snapshot_config_toml.open("r") as f:
             snapshot_config = toml.load(f)
 
+        snapshot_branches_toml = repo_dir / "branches.toml"
+        with snapshot_branches_toml.open("r") as f:
+            snapshot_branch_mappings = toml.load(f)
+
         self.initial_tenant = TenantId(snapshot_config["default_tenant_id"])
         self.initial_timeline = TimelineId(
             dict(snapshot_config["branch_name_mappings"][DEFAULT_BRANCH_NAME])[
@@ -729,9 +733,6 @@ class NeonEnvBuilder:
         with (self.repo_dir / "config").open("r") as f:
             config = toml.load(f)
 
-        config["default_tenant_id"] = snapshot_config["default_tenant_id"]
-        config["branch_name_mappings"] = snapshot_config["branch_name_mappings"]
-
         # Update the config with new neon + postgres path in case of compat test
         config["pg_distrib_dir"] = str(self.pg_distrib_dir)
         config["neon_distrib_dir"] = str(self.neon_binpath)
@@ -739,6 +740,9 @@ class NeonEnvBuilder:
         with (self.repo_dir / "config").open("w") as f:
             toml.dump(config, f)
 
+        with (self.repo_dir / "branches.toml").open("w") as f:
+            toml.dump(snapshot_branch_mappings, f)
+
         return self.env
 
     def overlay_mount(self, ident: str, srcdir: Path, dstdir: Path):
@@ -1118,7 +1122,6 @@ class NeonEnv:
 
         # Create the neon_local's `NeonLocalInitConf`
         cfg: Dict[str, Any] = {
-            "default_tenant_id": str(self.initial_tenant),
             "broker": {
                 "listen_addr": self.broker.listen_addr(),
             },