From 17419b8a628dc9453bfcdb0a9086024da0db0b6c Mon Sep 17 00:00:00 2001
From: Arthur Petukhovsky <petuhovskiy@yandex.ru>
Date: Tue, 25 Jan 2022 14:28:55 +0000
Subject: [PATCH] Add arg to override config in zenith_cli

---
 control_plane/src/local_env.rs          | 26 ++++++++++++--------
 pyproject.toml                          |  1 +
 test_runner/fixtures/zenith_fixtures.py | 32 +++++++++++++++++++++++--
 zenith/src/main.rs                      | 14 +++++++++--
 4 files changed, 59 insertions(+), 14 deletions(-)

diff --git a/control_plane/src/local_env.rs b/control_plane/src/local_env.rs
index b80e137cb9..e76f858fd0 100644
--- a/control_plane/src/local_env.rs
+++ b/control_plane/src/local_env.rs
@@ -189,20 +189,26 @@ impl LocalEnv {
     }
 
     /// Locate and load config
-    pub fn load_config() -> anyhow::Result<Self> {
+    pub fn load_config(override_config: Option<&str>) -> anyhow::Result<Self> {
         let repopath = base_path();
 
-        if !repopath.exists() {
-            bail!(
-                "Zenith config is not found in {}. You need to run 'zenith init' first",
-                repopath.to_str().unwrap()
-            );
-        }
-
-        // TODO: check that it looks like a zenith repository
+        let config_path = if let Some(override_config) = override_config {
+            PathBuf::from(override_config)
+        } else {
+            if !repopath.exists() {
+                bail!(
+                    "Zenith config is not found in {}. You need to run 'zenith init' first",
+                    repopath.to_str().unwrap()
+                );
+            }
+    
+            // TODO: check that it looks like a zenith repository
+    
+            repopath.join("config")
+        };
 
         // load and parse file
-        let config = fs::read_to_string(repopath.join("config"))?;
+        let config = fs::read_to_string(config_path)?;
         let mut env: LocalEnv = toml::from_str(config.as_str())?;
 
         env.base_data_dir = repopath;
diff --git a/pyproject.toml b/pyproject.toml
index 394c645104..f34f3c2ae3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,6 +21,7 @@ types-psycopg2 = "^2.9.6"
 boto3 = "^1.20.40"
 boto3-stubs = "^1.20.40"
 moto = {version = "^3.0.0", extras = ["server"]}
+toml = "^0.10.2"
 
 [tool.poetry.dev-dependencies]
 yapf = "==0.31.0"
diff --git a/test_runner/fixtures/zenith_fixtures.py b/test_runner/fixtures/zenith_fixtures.py
index c80eb35568..4cf24c2160 100644
--- a/test_runner/fixtures/zenith_fixtures.py
+++ b/test_runner/fixtures/zenith_fixtures.py
@@ -19,6 +19,7 @@ import subprocess
 import time
 import filecmp
 import tempfile
+import toml
 
 from contextlib import closing
 from pathlib import Path
@@ -372,7 +373,6 @@ class MockS3Server:
     def kill(self):
         self.subprocess.kill()
 
-
 class ZenithEnvBuilder:
     """
     Builder object to create a Zenith runtime environment
@@ -502,6 +502,9 @@ class ZenithEnv:
         self.port_distributor = config.port_distributor
         self.s3_mock_server = config.s3_mock_server
 
+        # If specified, this config will be passed for all zenith_cli calls
+        self.override_config: Optional[Dict] = None
+
         self.postgres = PostgresFactory(self)
 
         self.safekeepers: List[Safekeeper] = []
@@ -599,9 +602,22 @@ sync = false # Disable fsyncs to make the tests go faster
 
         assert type(arguments) == list
 
+        if self.override_config is not None:
+            tmp_config = tempfile.NamedTemporaryFile(mode='w+')
+            tmp_config.write(toml.dumps(self.override_config))
+            tmp_config.flush()
+            log.info('Using overriden config to run next command!')
+            log.info(f'Config: {toml.dumps(self.override_config)}')
+        else:
+            tmp_config = None
+
         bin_zenith = os.path.join(str(zenith_binpath), 'zenith')
 
-        args = [bin_zenith] + arguments
+        args = [bin_zenith]
+        if tmp_config is not None:
+            args += ['--override-config', tmp_config.name]
+        args += arguments
+
         log.info('Running command "{}"'.format(' '.join(args)))
         log.info(f'Running in "{self.repo_dir}"')
 
@@ -637,9 +653,21 @@ sync = false # Disable fsyncs to make the tests go faster
             log.info(msg)
 
             raise Exception(msg) from exc
+        finally:
+            if tmp_config is not None:
+                tmp_config.close()
 
         return res
 
+    def read_toml_config(self) -> Dict:
+        """
+        Read the config file from the repo directory.
+
+        Returns a dictionary of the config file.
+        """
+        with open(os.path.join(str(self.repo_dir), 'config')) as f:
+            return toml.load(f)
+
     @cached_property
     def auth_keys(self) -> AuthKeys:
         pub = (Path(self.repo_dir) / 'auth_public_key.pem').read_bytes()
diff --git a/zenith/src/main.rs b/zenith/src/main.rs
index 976cbfac0c..76a2684e6f 100644
--- a/zenith/src/main.rs
+++ b/zenith/src/main.rs
@@ -112,6 +112,13 @@ fn main() -> Result<()> {
     let matches = App::new("Zenith CLI")
         .setting(AppSettings::ArgRequiredElseHelp)
         .version(GIT_VERSION)
+        .arg(
+            Arg::with_name("override-config")
+                .long("override-config")
+                .takes_value(true)
+                .help("Override the default configuration file")
+                .required(false)
+        )
         .subcommand(
             SubCommand::with_name("init")
                 .about("Initialize a new Zenith repository")
@@ -220,8 +227,11 @@ fn main() -> Result<()> {
     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() {
+        // all other commands need an existing config, which is found in base directory,
+        // or can be overridden with --override-config
+
+        let override_config = matches.value_of("override-config");
+        let env = match LocalEnv::load_config(override_config) {
             Ok(conf) => conf,
             Err(e) => {
                 eprintln!("Error loading config: {}", e);