Merge pull request PrefectHQ#195 from PrefectHQ/user-config

jlowin · web-flow · commit 1772d4e6b311 · 2018-09-06T12:49:52.000-04:00
Add support for user configuration files
diff --git a/src/prefect/config.toml b/src/prefect/config.toml
@@ -1,3 +1,9 @@
+[general]
+
+# the location of the user's config file
+user_config_path = "$HOME/.prefect/config.toml"
+
+
 [server]
 # the Prefect Server address
 api_server = "http://127.0.0.1:4200"
diff --git a/src/prefect/configuration.py b/src/prefect/configuration.py
@@ -6,11 +6,9 @@
 
 import toml
 
-import prefect
 from prefect.utilities import collections
 
 DEFAULT_CONFIG = os.path.join(os.path.dirname(__file__), "config.toml")
-USER_CONFIG = "~/.prefect/config.toml"
 ENV_VAR_PREFIX = "PREFECT"
 INTERPOLATION_REGEX = re.compile(r"\${(.[^${}]*)}")
 
@@ -74,7 +72,7 @@ def validate_config(config: Config) -> None:
 # Load configuration ----------------------------------------------------------
 
 
-def load_config_file(path: str, env_var_prefix: str = ENV_VAR_PREFIX) -> Config:
+def load_config_file(path: str, env_var_prefix: str = None) -> Config:
     """
     Loads a configuration file from a path, optionally merging it into an existing
     configuration.
@@ -90,19 +88,22 @@ def load_config_file(path: str, env_var_prefix: str = ENV_VAR_PREFIX) -> Config:
     # --------------------- Interpolate env vars -----------------------
     # check if any env var sets a configuration value with the format:
     # [ENV_VAR_PREFIX]__[Section]__[Optional Sub-Sections...]__[Key] = Value
-    for env_var in os.environ:
-        if env_var.startswith(env_var_prefix + "__"):
+    if env_var_prefix:
+        for env_var in os.environ:
+            if env_var.startswith(env_var_prefix + "__"):
 
-            # strip the prefix off the env var
-            env_var_option = env_var[len(env_var_prefix + "__") :]
+                # strip the prefix off the env var
+                env_var_option = env_var[len(env_var_prefix + "__") :]
 
-            # make sure the resulting env var has at least one delimitied section and key
-            if "__" not in env_var:
-                continue
+                # make sure the resulting env var has at least one delimitied section and key
+                if "__" not in env_var:
+                    continue
 
-            # place the env var in the flat config as a compound key
-            config_option = collections.CompoundKey(env_var_option.lower().split("__"))
-            flat_config[config_option] = interpolate_env_var(os.getenv(env_var))
+                # place the env var in the flat config as a compound key
+                config_option = collections.CompoundKey(
+                    env_var_option.lower().split("__")
+                )
+                flat_config[config_option] = interpolate_env_var(os.getenv(env_var))
 
     # interpolate any env vars referenced
     for k, v in list(flat_config.items()):
@@ -153,28 +154,39 @@ def load_config_file(path: str, env_var_prefix: str = ENV_VAR_PREFIX) -> Config:
 
 
 def load_configuration(
-    default_config_path: str, user_config_path: str, env_var_prefix: str = None
+    config_path: str, env_var_prefix: str = None, merge_into_config: Config = None
 ) -> Config:
+    """
+    Given a `config_path` with a toml configuration file, returns a Config object.
+
+    Args:
+        - config_path (str): the path to the toml configuration file
+        - env_var_prefix (str): if provided, environment variables starting with this prefix
+            will be added as configuration settings.
+        - merge_into_config (Config): if provided, the configuration loaded from
+            `config_path` will be merged into a copy of this configuration file. The merged
+            Config is returned.
+    """
 
     # load default config
-    config = load_config_file(default_config_path, env_var_prefix=env_var_prefix or "")
+    config = load_config_file(config_path, env_var_prefix=env_var_prefix or "")
 
-    # if user config exists, load and merge it with default config
-    if os.path.isfile(user_config_path):
-        user_config = load_config_file(
-            user_config_path, env_var_prefix=env_var_prefix or ""
-        )
-        config = collections.merge_dicts(config, user_config)
+    if merge_into_config is not None:
+        config = collections.merge_dicts(merge_into_config, config)
 
     validate_config(config)
 
     return config
 
 
-config = load_configuration(
-    default_config_path=DEFAULT_CONFIG,
-    user_config_path=USER_CONFIG,
-    env_var_prefix=ENV_VAR_PREFIX,
-)
+config = load_configuration(config_path=DEFAULT_CONFIG, env_var_prefix=ENV_VAR_PREFIX)
+
+# if user config exists, load and merge it with default config
+if os.path.isfile(config.get("general", {}).get("user_config_path", "")):
+    config = load_configuration(
+        config_path=config.general.user_config_path,
+        env_var_prefix=ENV_VAR_PREFIX,
+        merge_into_config=config,
+    )
 
 configure_logging(logger_name="Prefect")
diff --git a/tests/test_configuration.py b/tests/test_configuration.py
@@ -78,22 +78,50 @@ def test_env_var_interpolation(config):
 
 
 def test_env_var_overrides(test_config_file_path):
-    os.environ["PREFECT__ENV_VARS__TEST"] = "OVERRIDE!"
-    assert config(test_config_file_path).env_vars.test == "OVERRIDE!"
+    try:
+        ev = "PREFECT__ENV_VARS__TEST"
+        os.environ[ev] = "OVERRIDE!"
+        config = configuration.load_config_file(
+            test_config_file_path, env_var_prefix="PREFECT"
+        )
+        assert config.env_vars.test == "OVERRIDE!"
+    finally:
+        del os.environ[ev]
 
 
-def test_load_user_config_and_update_default(test_config_file_path):
+def test_merge_configurations(test_config_file_path):
+
+    default_config = configuration.config
+
+    assert default_config.logging.format != "log-format"
+    assert default_config.flows.default_version == "1"
+
     config = configuration.load_configuration(
-        default_config_path=configuration.DEFAULT_CONFIG,
-        user_config_path=test_config_file_path,
+        config_path=test_config_file_path, merge_into_config=default_config
     )
-    assert "logging" in config
 
-    # this comes from default
-    assert config.logging.level == "INFO"
-    # this comes from user
     assert config.logging.format == "log-format"
-
-    # this comes from user
-    assert "general" in config
-    assert config.general.nested.x == 1
+    assert config.flows.default_version == "1"
+    assert config.interpolation.value == 1
+
+
+def test_load_user_config(test_config_file_path):
+
+    with tempfile.NamedTemporaryFile() as user_config:
+        user_config.write(
+            b"""
+            [general]
+            x = 2
+
+            [user]
+            foo = "bar"
+            """
+        )
+        user_config.seek(0)
+
+        test_config = configuration.load_configuration(test_config_file_path)
+        config = configuration.load_configuration(
+            user_config.name, merge_into_config=test_config
+        )
+        assert config.general.x == 2
+        assert config.user.foo == "bar"
