From 841c25e1bdc74db2f264d11b6b0b35c9bf0ff347 Mon Sep 17 00:00:00 2001
From: yuvipanda <yuvipanda@gmail.com>
Date: Sat, 11 Aug 2018 01:19:51 -0700
Subject: [PATCH] Add plugin hook to modify config.yaml post install

---
 tljh/hooks.py     | 13 +++++++++++++
 tljh/installer.py | 45 ++++++++++++++++++++++++++++++++++++---------
 2 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/tljh/hooks.py b/tljh/hooks.py
index f1a673e..8cd6d56 100644
--- a/tljh/hooks.py
+++ b/tljh/hooks.py
@@ -30,4 +30,17 @@ def tljh_extra_apt_packages():
 
     These will be installed before additional pip or conda packages.
     """
+    pass
+
+
+@hookspec
+def tljh_config_post_install(config):
+    """
+    Modify on-disk tljh-config after installation.
+
+    config is a dict-like object that should be modified
+    in-place. The contents of the on-disk config.yaml will
+    be the serialized contents of config, so try to not
+    overwrite anything the user might have explicitly set.
+    """
     pass
\ No newline at end of file
diff --git a/tljh/installer.py b/tljh/installer.py
index c384338..301e83c 100644
--- a/tljh/installer.py
+++ b/tljh/installer.py
@@ -13,7 +13,7 @@ import pluggy
 from ruamel.yaml import YAML
 
 from tljh import conda, systemd, traefik, user, apt, hooks
-from tljh.config import INSTALL_PREFIX, HUB_ENV_PREFIX, USER_ENV_PREFIX, STATE_DIR
+from tljh.config import INSTALL_PREFIX, HUB_ENV_PREFIX, USER_ENV_PREFIX, STATE_DIR, CONFIG_FILE
 
 HERE = os.path.abspath(os.path.dirname(__file__))
 
@@ -307,11 +307,11 @@ def ensure_symlinks(prefix):
             return
     os.symlink(tljh_config_src, tljh_config_dest)
 
-def run_plugin_actions(plugins):
-    """
-    Run installer hooks defined in plugins
-    """
 
+def setup_plugins(plugins):
+    """
+    Install plugins & setup a pluginmanager
+    """
     # Install plugins
     if plugins:
         conda.ensure_pip_packages(HUB_ENV_PREFIX, plugins)
@@ -321,8 +321,15 @@ def run_plugin_actions(plugins):
     pm.add_hookspecs(hooks)
     pm.load_setuptools_entrypoints('tljh')
 
+    return pm
+
+def run_plugin_actions(plugin_manager, plugins):
+    """
+    Run installer hooks defined in plugins
+    """
+    hook = plugin_manager.hook
     # Install apt packages
-    apt_packages = list(set(itertools.chain(*pm.hook.tljh_extra_apt_packages())))
+    apt_packages = list(set(itertools.chain(*hook.tljh_extra_apt_packages())))
     if apt_packages:
         logger.info('Installing {} apt packages collected from plugins: {}'.format(
             len(apt_packages), ' '.join(apt_packages)
@@ -330,7 +337,7 @@ def run_plugin_actions(plugins):
         apt.install_packages(apt_packages)
 
     # Install conda packages
-    conda_packages = list(set(itertools.chain(*pm.hook.tljh_extra_user_conda_packages())))
+    conda_packages = list(set(itertools.chain(*hook.tljh_extra_user_conda_packages())))
     if conda_packages:
         logger.info('Installing {} conda packages collected from plugins: {}'.format(
             len(conda_packages), ' '.join(conda_packages)
@@ -338,7 +345,7 @@ def run_plugin_actions(plugins):
         conda.ensure_conda_packages(USER_ENV_PREFIX, conda_packages)
 
     # Install pip packages
-    pip_packages = list(set(itertools.chain(*pm.hook.tljh_extra_user_pip_packages())))
+    pip_packages = list(set(itertools.chain(*hook.tljh_extra_user_pip_packages())))
     if pip_packages:
         logger.info('Installing {} pip packages collected from plugins: {}'.format(
             len(pip_packages), ' '.join(pip_packages)
@@ -346,6 +353,23 @@ def run_plugin_actions(plugins):
         conda.ensure_pip_packages(USER_ENV_PREFIX, pip_packages)
 
 
+def ensure_config_yaml(plugin_manager):
+    """
+    Ensure we have a config.yaml present
+    """
+    if os.path.exists(CONFIG_FILE):
+        with open(CONFIG_FILE, 'r') as f:
+            config = rt_yaml.load(f)
+    else:
+        config = {}
+
+    hook = plugin_manager.hook
+    hook.tljh_config_post_install(config=config)
+
+    with open(CONFIG_FILE, 'w+') as f:
+        rt_yaml.dump(config, f)
+
+
 def main():
     argparser = argparse.ArgumentParser()
     argparser.add_argument(
@@ -365,6 +389,8 @@ def main():
 
     args = argparser.parse_args()
 
+    pm = setup_plugins(args.plugin)
+
     ensure_admins(args.admin)
 
     ensure_usergroups()
@@ -374,12 +400,13 @@ def main():
     ensure_node()
     ensure_jupyterhub_package(HUB_ENV_PREFIX)
     ensure_chp_package(HUB_ENV_PREFIX)
+    ensure_config_yaml(pm)
     ensure_jupyterhub_service(HUB_ENV_PREFIX)
     ensure_jupyterhub_running()
     ensure_symlinks(HUB_ENV_PREFIX)
 
     # Run installer plugins last
-    run_plugin_actions(args.plugin)
+    run_plugin_actions(pm, args.plugin)
 
     logger.info("Done!")