[airflow] branch master updated: Return nonzero exit codes on pool import errors. (#12095)

[jm...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jmcarp pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/master by this push:
     new 397d912  Return nonzero exit codes on pool import errors. (#12095)
397d912 is described below

commit 397d9128b6abb201e4eafd4a1dd7ece9db5661c5
Author: Joshua Carp <jm...@gmail.com>
AuthorDate: Sat Nov 21 19:33:06 2020 -0500

    Return nonzero exit codes on pool import errors. (#12095)
    
    The pool import command returns an exit code of zero in a few different
    error cases. This causes problems for scripts that invoke the command,
    since commands that actually failed will appear to have worked. This
    patch returns a nonzero code if the pool file doesn't exist, if the file
    isn't valid json, or if any of the pools in the file is invalid.
---
 airflow/cli/commands/pool_command.py    | 40 +++++++++++++--------------------
 tests/cli/commands/test_pool_command.py | 19 ++++++++++++++++
 2 files changed, 35 insertions(+), 24 deletions(-)

diff --git a/airflow/cli/commands/pool_command.py b/airflow/cli/commands/pool_command.py
index 4984618..641b977 100644
--- a/airflow/cli/commands/pool_command.py
+++ b/airflow/cli/commands/pool_command.py
@@ -18,6 +18,7 @@
 """Pools sub-commands"""
 import json
 import os
+import sys
 from json import JSONDecodeError
 
 from tabulate import tabulate
@@ -63,13 +64,12 @@ def pool_delete(args):
 @cli_utils.action_logging
 def pool_import(args):
     """Imports pools from the file"""
-    api_client = get_current_api_client()
-    if os.path.exists(args.file):
-        pools = pool_import_helper(args.file)
-    else:
-        print("Missing pools file.")
-        pools = api_client.get_pools()
+    if not os.path.exists(args.file):
+        sys.exit("Missing pools file.")
+    pools, failed = pool_import_helper(args.file)
     print(_tabulate_pools(pools=pools, tablefmt=args.output))
+    if len(failed) > 0:
+        sys.exit("Failed to update pool(s): {}".format(", ".join(failed)))
 
 
 def pool_export(args):
@@ -87,24 +87,16 @@ def pool_import_helper(filepath):
     try:  # pylint: disable=too-many-nested-blocks
         pools_json = json.loads(data)
     except JSONDecodeError as e:
-        print("Please check the validity of the json file: " + str(e))
-    else:
-        try:
-            pools = []
-            counter = 0
-            for k, v in pools_json.items():
-                if isinstance(v, dict) and len(v) == 2:
-                    pools.append(
-                        api_client.create_pool(name=k, slots=v["slots"], description=v["description"])
-                    )
-                    counter += 1
-                else:
-                    pass
-        except Exception:  # pylint: disable=broad-except
-            pass
-        finally:
-            print("{} of {} pool(s) successfully updated.".format(counter, len(pools_json)))
-            return pools  # pylint: disable=lost-exception
+        sys.exit("Invalid json file: " + str(e))
+    pools = []
+    failed = []
+    for k, v in pools_json.items():
+        if isinstance(v, dict) and len(v) == 2:
+            pools.append(api_client.create_pool(name=k, slots=v["slots"], description=v["description"]))
+        else:
+            failed.append(k)
+    print(f"{len(pools)} of {len(pools_json)} pool(s) successfully updated.")
+    return pools, failed
 
 
 def pool_export_helper(filepath):
diff --git a/tests/cli/commands/test_pool_command.py b/tests/cli/commands/test_pool_command.py
index be95d09..5d2e875 100644
--- a/tests/cli/commands/test_pool_command.py
+++ b/tests/cli/commands/test_pool_command.py
@@ -78,6 +78,25 @@ class TestCliPools(unittest.TestCase):
         pool_command.pool_delete(self.parser.parse_args(['pools', 'delete', 'foo']))
         self.assertEqual(self.session.query(Pool).count(), 1)
 
+    def test_pool_import_nonexistent(self):
+        with self.assertRaises(SystemExit):
+            pool_command.pool_import(self.parser.parse_args(['pools', 'import', 'nonexistent.json']))
+
+    def test_pool_import_invalid_json(self):
+        with open('pools_import_invalid.json', mode='w') as file:
+            file.write("not valid json")
+
+        with self.assertRaises(SystemExit):
+            pool_command.pool_import(self.parser.parse_args(['pools', 'import', 'pools_import_invalid.json']))
+
+    def test_pool_import_invalid_pools(self):
+        pool_config_input = {"foo": {"description": "foo_test"}}
+        with open('pools_import_invalid.json', mode='w') as file:
+            json.dump(pool_config_input, file)
+
+        with self.assertRaises(SystemExit):
+            pool_command.pool_import(self.parser.parse_args(['pools', 'import', 'pools_import_invalid.json']))
+
     def test_pool_import_export(self):
         # Create two pools first
         pool_config_input = {