def pool_set(args): """Creates new pool with a given name and slots""" api_client = get_current_api_client() pools = [api_client.create_pool(name=args.pool, slots=args.slots, description=args.description)] print(_tabulate_pools(pools=pools, tablefmt=args.output))
def test_should_create_client(self, mock_client): result = get_current_api_client() mock_client.assert_called_once_with( api_base_url='http://localhost:1234', auth='CLIENT_AUTH', session=None ) self.assertEqual(mock_client.return_value, result)
def pool_import_helper(filepath): """Helps import pools from the json file""" api_client = get_current_api_client() with open(filepath, 'r') as poolfile: data = poolfile.read() 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
def pool_set(args): """Creates new pool with a given name and slots""" api_client = get_current_api_client() api_client.create_pool(name=args.pool, slots=args.slots, description=args.description) print("Pool created")
def pool_get(args): """Displays pool info by a given name""" api_client = get_current_api_client() try: pools = [api_client.get_pool(name=args.pool)] _show_pools(pools=pools, output=args.output) except PoolNotFound: raise SystemExit(f"Pool {args.pool} does not exist")
def pool_delete(args): """Deletes pool by a given name""" api_client = get_current_api_client() try: api_client.delete_pool(name=args.pool) print("Pool deleted") except PoolNotFound: sys.exit(f"Pool {args.pool} does not exist")
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() print(_tabulate_pools(pools=pools, tablefmt=args.output))
def test_should_create_google_open_id_client(self, mock_create_client_session, mock_client): result = get_current_api_client() mock_client.assert_called_once_with( api_base_url='http://localhost:1234', auth=None, session=mock_create_client_session.return_value ) self.assertEqual(mock_client.return_value, result)
def dag_trigger(args): """Creates a dag run for the specified dag""" api_client = get_current_api_client() try: message = api_client.trigger_dag( dag_id=args.dag_id, run_id=args.run_id, conf=args.conf, execution_date=args.exec_date ) print(message) except OSError as err: raise AirflowException(err)
def pool_export_helper(filepath): """Helps export all of the pools to the json file""" api_client = get_current_api_client() 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)) return pools
def pool_set(args): """Creates new pool with a given name and slots""" api_client = get_current_api_client() log = LoggingMixin().log pools = [ api_client.create_pool(name=args.pool, slots=args.slots, description=args.description) ] log.info(_tabulate_pools(pools=pools, tablefmt=args.output))
def dag_delete(args): """Deletes all DB records related to the specified dag""" api_client = get_current_api_client() if (args.yes or input( "This will drop all existing records related to the specified DAG. Proceed? (y/n)" ).upper() == "Y"): try: message = api_client.delete_dag(dag_id=args.dag_id) print(message) except OSError as err: raise AirflowException(err) else: print("Cancelled")
def dag_trigger(args): """ Creates a dag run for the specified dag """ api_client = get_current_api_client() log = LoggingMixin().log try: message = api_client.trigger_dag(dag_id=args.dag_id, run_id=args.run_id, conf=args.conf, execution_date=args.exec_date) except OSError as err: log.error(err) raise AirflowException(err) log.info(message)
def dag_delete(args): """ Deletes all DB records related to the specified dag """ api_client = get_current_api_client() log = LoggingMixin().log if args.yes or input( "This will drop all existing records related to the specified DAG. " "Proceed? (y/n)").upper() == "Y": try: message = api_client.delete_dag(dag_id=args.dag_id) except OSError as err: log.error(err) raise AirflowException(err) log.info(message) else: print("Bail.")
def pool_import_helper(filepath): """Helps import pools from the json file""" api_client = get_current_api_client() with open(filepath) as poolfile: data = poolfile.read() try: # pylint: disable=too-many-nested-blocks pools_json = json.loads(data) except JSONDecodeError as e: raise SystemExit("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) return pools, failed
def pool_list(args): """Displays info of all the pools""" api_client = get_current_api_client() log = LoggingMixin().log pools = api_client.get_pools() log.info(_tabulate_pools(pools=pools, tablefmt=args.output))
def pool_delete(args): """Deletes pool by a given name""" api_client = get_current_api_client() log = LoggingMixin().log pools = [api_client.delete_pool(name=args.pool)] log.info(_tabulate_pools(pools=pools, tablefmt=args.output))
def pool_delete(args): """Deletes pool by a given name""" api_client = get_current_api_client() pools = [api_client.delete_pool(name=args.pool)] print(_tabulate_pools(pools=pools, tablefmt=args.output))
def pool_list(args): """Displays info of all the pools""" api_client = get_current_api_client() pools = api_client.get_pools() _show_pools(pools=pools, output=args.output)
def pool_get(args): """Displays pool info by a given name""" api_client = get_current_api_client() pools = [api_client.get_pool(name=args.pool)] print(_tabulate_pools(pools=pools, tablefmt=args.output))
def pool_list(args): """Displays info of all the pools""" api_client = get_current_api_client() pools = api_client.get_pools() print(_tabulate_pools(pools=pools, tablefmt=args.output))
def pool_get(args): """Displays pool info by a given name""" api_client = get_current_api_client() pools = [api_client.get_pool(name=args.pool)] _show_pools(pools=pools, output=args.output)
def pool_get(args): """Displays pool info by a given name""" api_client = get_current_api_client() log = LoggingMixin().log pools = [api_client.get_pool(name=args.pool)] log.info(_tabulate_pools(pools=pools, tablefmt=args.output))