[GitHub] r39132 closed pull request #3730: [AIRFLOW-2882] Add import and export for pool cli using JSON

[GitBox <gi...@apache.org>]

r39132 closed pull request #3730: [AIRFLOW-2882] Add import and export for pool cli using JSON
URL: https://github.com/apache/incubator-airflow/pull/3730
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/airflow/bin/cli.py b/airflow/bin/cli.py
index e22427cf40..4ff1ae3679 100644
--- a/airflow/bin/cli.py
+++ b/airflow/bin/cli.py
@@ -267,6 +267,7 @@ def _tabulate(pools):
                                  tablefmt="fancy_grid")
 
     try:
+        imp = getattr(args, 'import')
         if args.get is not None:
             pools = [api_client.get_pool(name=args.get)]
         elif args.set:
@@ -275,6 +276,14 @@ def _tabulate(pools):
                                             description=args.set[2])]
         elif args.delete:
             pools = [api_client.delete_pool(name=args.delete)]
+        elif imp:
+            if os.path.exists(imp):
+                pools = pool_import_helper(imp)
+            else:
+                print("Missing pools file.")
+                pools = api_client.get_pools()
+        elif args.export:
+            pools = pool_export_helper(args.export)
         else:
             pools = api_client.get_pools()
     except (AirflowException, IOError) as err:
@@ -283,6 +292,43 @@ def _tabulate(pools):
         log.info(_tabulate(pools=pools))
 
 
+def pool_import_helper(filepath):
+    with open(filepath, 'r') as poolfile:
+        pl = poolfile.read()
+    try:
+        d = json.loads(pl)
+    except Exception as e:
+        print("Please check the validity of the json file: " + str(e))
+    else:
+        try:
+            pools = []
+            n = 0
+            for k, v in d.items():
+                if isinstance(v, dict) and len(v) == 2:
+                    pools.append(api_client.create_pool(name=k,
+                                                        slots=v["slots"],
+                                                        description=v["description"]))
+                    n += 1
+                else:
+                    pass
+        except Exception:
+            pass
+        finally:
+            print("{} of {} pool(s) successfully updated.".format(n, len(d)))
+            return pools
+
+
+def pool_export_helper(filepath):
+    pool_dict = {}
+    pools = api_client.get_pools()
+    for pool in pools:
+        pool_dict[pool[0]] = {"slots": pool[1], "description": pool[2]}
+    with open(filepath, 'w') as poolfile:
+        poolfile.write(json.dumps(pool_dict, sort_keys=True, indent=4))
+    print("{} pools successfully exported to {}".format(len(pool_dict), filepath))
+    return pools
+
+
 @cli_utils.action_logging
 def variables(args):
     if args.get:
@@ -1551,6 +1597,14 @@ class CLIFactory(object):
             ("-x", "--delete"),
             metavar="NAME",
             help="Delete a pool"),
+        'pool_import': Arg(
+            ("-i", "--import"),
+            metavar="FILEPATH",
+            help="Import pool from JSON file"),
+        'pool_export': Arg(
+            ("-e", "--export"),
+            metavar="FILEPATH",
+            help="Export pool to JSON file"),
         # variables
         'set': Arg(
             ("-s", "--set"),
@@ -1870,7 +1924,7 @@ class CLIFactory(object):
         }, {
             'func': pool,
             'help': "CRUD operations on pools",
-            "args": ('pool_set', 'pool_get', 'pool_delete'),
+            "args": ('pool_set', 'pool_get', 'pool_delete', 'pool_import', 'pool_export'),
         }, {
             'func': variables,
             'help': "CRUD operations on variables",
diff --git a/tests/core.py b/tests/core.py
index f8b8691912..b59ae7048d 100644
--- a/tests/core.py
+++ b/tests/core.py
@@ -1512,6 +1512,42 @@ def test_pool_no_args(self):
         except Exception as e:
             self.fail("The 'pool' command raised unexpectedly: %s" % e)
 
+    def test_pool_import_export(self):
+        # Create two pools first
+        pool_config_input = {
+            "foo": {
+                "description": "foo_test",
+                "slots": 1
+            },
+            "baz": {
+                "description": "baz_test",
+                "slots": 2
+            }
+        }
+        with open('pools_import.json', mode='w') as f:
+            json.dump(pool_config_input, f)
+
+        # Import json
+        try:
+            cli.pool(self.parser.parse_args(['pool', '-i', 'pools_import.json']))
+        except Exception as e:
+            self.fail("The 'pool -i pools_import.json' failed: %s" % e)
+
+        # Export json
+        try:
+            cli.pool(self.parser.parse_args(['pool', '-e', 'pools_export.json']))
+        except Exception as e:
+            self.fail("The 'pool -e pools_export.json' failed: %s" % e)
+
+        with open('pools_export.json', mode='r') as f:
+            pool_config_output = json.load(f)
+            self.assertEqual(
+                pool_config_input,
+                pool_config_output,
+                "Input and output pool files are not same")
+        os.remove('pools_import.json')
+        os.remove('pools_export.json')
+
     def test_variables(self):
         # Checks if all subcommands are properly received
         cli.variables(self.parser.parse_args([


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services