From c3f270f11c4aa60314a125fdd1ebbed65b8a643b Mon Sep 17 00:00:00 2001
From: Joseph Hamman <jhamman@ucar.edu>
Date: Thu, 12 Mar 2020 23:21:17 -0700
Subject: [PATCH] add atomic and cleanup args to hubploy cli

---
 hubploy/__main__.py | 10 +++++++++-
 hubploy/helm.py     | 16 +++++++++++++---
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/hubploy/__main__.py b/hubploy/__main__.py
index 0ded564..6a2186f 100644
--- a/hubploy/__main__.py
+++ b/hubploy/__main__.py
@@ -67,6 +67,14 @@ def main():
         '--force',
         action='store_true'
     )
+    deploy_parser.add_argument(
+        '--atomic',
+        action='store_true'
+    )
+    deploy_parser.add_argument(
+        '--cleanup-on-fail',
+        action='store_true'
+    )
 
     args = argparser.parse_args()
 
@@ -91,7 +99,7 @@ def main():
 
     elif args.command == 'deploy':
         auth.cluster_auth(args.deployment)
-        helm.deploy(args.deployment, args.chart, args.environment, args.namespace, args.set, args.version, args.timeout, args.force)
+        helm.deploy(args.deployment, args.chart, args.environment, args.namespace, args.set, args.version, args.timeout, args.force, args.atomic, args.cleanup_on_fail)
 
 if __name__ == '__main__':
     main()
\ No newline at end of file
diff --git a/hubploy/helm.py b/hubploy/helm.py
index 85430b9..3b8e657 100644
--- a/hubploy/helm.py
+++ b/hubploy/helm.py
@@ -37,7 +37,9 @@ def helm_upgrade(
     config_overrides,
     version,
     timeout,
-    force
+    force,
+    atomic,
+    cleanup_on_fail
 ):
     # Clear charts and do a helm dep up before installing
     # Clearing charts is important so we don't deploy charts that
@@ -82,6 +84,10 @@ def helm_upgrade(
         cmd += ['--timeout', timeout]
     if force:
         cmd += ['--force']
+    if atomic:
+        cmd += ['--atomic']
+    if cleanup_on_fail:
+        cmd += ['--cleanup-on-fail']
     cmd += itertools.chain(*[['-f', cf] for cf in config_files])
     cmd += itertools.chain(*[['--set', v] for v in config_overrides])
     subprocess.check_call(cmd)
@@ -95,7 +101,9 @@ def deploy(
     helm_config_overrides=None,
     version=None,
     timeout=None,
-    force=False
+    force=False,
+    atomic=False,
+    cleanup_on_fail=False
 ):
     """
     Deploy a JupyterHub.
@@ -148,5 +156,7 @@ def deploy(
         helm_config_overrides,
         version,
         timeout,
-        force
+        force,
+        atomic,
+        cleanup_on_fail
     )