You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ma...@apache.org on 2018/04/02 22:16:36 UTC
incubator-airflow git commit: [AIRFLOW-2253] Add Airflow CLI
instrumentation
Repository: incubator-airflow
Updated Branches:
refs/heads/master a6487d619 -> 9dba430b6
[AIRFLOW-2253] Add Airflow CLI instrumentation
Closes #3159 from jinhyukchang/cli-instrumentation
Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/9dba430b
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/9dba430b
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/9dba430b
Branch: refs/heads/master
Commit: 9dba430b683361fc0ed7f50de6daa03c971a476b
Parents: a6487d6
Author: Jin Hyuk Chang <ji...@lyft.com>
Authored: Mon Apr 2 15:16:25 2018 -0700
Committer: Maxime Beauchemin <ma...@gmail.com>
Committed: Mon Apr 2 15:16:25 2018 -0700
----------------------------------------------------------------------
airflow/bin/cli.py | 30 ++++++++
airflow/settings.py | 11 ++-
airflow/utils/cli.py | 114 +++++++++++++++++++++++++++++++
airflow/utils/cli_action_loggers.py | 101 +++++++++++++++++++++++++++
tests/utils/test_cli_util.py | 97 ++++++++++++++++++++++++++
5 files changed, 352 insertions(+), 1 deletion(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/9dba430b/airflow/bin/cli.py
----------------------------------------------------------------------
diff --git a/airflow/bin/cli.py b/airflow/bin/cli.py
index 24f5d31..eb96e77 100755
--- a/airflow/bin/cli.py
+++ b/airflow/bin/cli.py
@@ -54,6 +54,7 @@ from airflow.models import (DagModel, DagBag, TaskInstance,
Connection, DAG)
from airflow.ti_deps.dep_context import (DepContext, SCHEDULER_DEPS)
+from airflow.utils import cli as cli_utils
from airflow.utils import db as db_utils
from airflow.utils.net import get_hostname
from airflow.utils.log.logging_mixin import (LoggingMixin, redirect_stderr,
@@ -149,6 +150,7 @@ def get_dags(args):
return matched_dags
+@cli_utils.action_logging
def backfill(args, dag=None):
logging.basicConfig(
level=settings.LOGGING_LEVEL,
@@ -190,6 +192,7 @@ def backfill(args, dag=None):
delay_on_limit_secs=args.delay_on_limit)
+@cli_utils.action_logging
def trigger_dag(args):
"""
Creates a dag run for the specified dag
@@ -208,6 +211,7 @@ def trigger_dag(args):
log.info(message)
+@cli_utils.action_logging
def delete_dag(args):
"""
Deletes all DB records related to the specified dag
@@ -228,6 +232,7 @@ def delete_dag(args):
print("Bail.")
+@cli_utils.action_logging
def pool(args):
log = LoggingMixin().log
@@ -252,6 +257,7 @@ def pool(args):
log.info(_tabulate(pools=pools))
+@cli_utils.action_logging
def variables(args):
if args.get:
try:
@@ -328,10 +334,12 @@ def export_helper(filepath):
print("{} variables successfully exported to {}".format(len(var_dict), filepath))
+@cli_utils.action_logging
def pause(args, dag=None):
set_is_paused(True, args, dag)
+@cli_utils.action_logging
def unpause(args, dag=None):
set_is_paused(False, args, dag)
@@ -401,6 +409,7 @@ def _run(args, dag, ti):
executor.end()
+@cli_utils.action_logging
def run(args, dag=None):
# Disable connection pooling to reduce the # of connections on the DB
# while it's waiting for the task to finish.
@@ -453,6 +462,8 @@ def run(args, dag=None):
_run(args, dag, ti)
logging.shutdown()
+
+@cli_utils.action_logging
def task_failed_deps(args):
"""
Returns the unmet dependencies for a task instance from the perspective of the
@@ -479,6 +490,7 @@ def task_failed_deps(args):
print("Task instance dependencies are all met.")
+@cli_utils.action_logging
def task_state(args):
"""
Returns the state of a TaskInstance at the command line.
@@ -492,6 +504,7 @@ def task_state(args):
print(ti.current_state())
+@cli_utils.action_logging
def dag_state(args):
"""
Returns the state of a DagRun at the command line.
@@ -504,6 +517,7 @@ def dag_state(args):
print(dr[0].state if len(dr) > 0 else None)
+@cli_utils.action_logging
def list_dags(args):
dagbag = DagBag(process_subdir(args.subdir))
s = textwrap.dedent("""\n
@@ -518,6 +532,7 @@ def list_dags(args):
print(dagbag.dagbag_report())
+@cli_utils.action_logging
def list_tasks(args, dag=None):
dag = dag or get_dag(args)
if args.tree:
@@ -527,6 +542,7 @@ def list_tasks(args, dag=None):
print("\n".join(sorted(tasks)))
+@cli_utils.action_logging
def test(args, dag=None):
dag = dag or get_dag(args)
@@ -543,6 +559,7 @@ def test(args, dag=None):
ti.run(ignore_task_deps=True, ignore_ti_state=True, test_mode=True)
+@cli_utils.action_logging
def render(args):
dag = get_dag(args)
task = dag.get_task(task_id=args.task_id)
@@ -557,6 +574,7 @@ def render(args):
""".format(attr, getattr(task, attr))))
+@cli_utils.action_logging
def clear(args):
logging.basicConfig(
level=settings.LOGGING_LEVEL,
@@ -713,6 +731,7 @@ def restart_workers(gunicorn_master_proc, num_workers_expected, master_timeout):
sys.exit(1)
+@cli_utils.action_logging
def webserver(args):
print(settings.HEADER)
@@ -838,6 +857,7 @@ def webserver(args):
monitor_gunicorn(gunicorn_master_proc)
+@cli_utils.action_logging
def scheduler(args):
print(settings.HEADER)
job = jobs.SchedulerJob(
@@ -871,6 +891,7 @@ def scheduler(args):
job.run()
+@cli_utils.action_logging
def serve_logs(args):
print("Starting flask")
import flask
@@ -891,6 +912,7 @@ def serve_logs(args):
host='0.0.0.0', port=WORKER_LOG_SERVER_PORT)
+@cli_utils.action_logging
def worker(args):
env = os.environ.copy()
env['AIRFLOW_HOME'] = settings.AIRFLOW_HOME
@@ -937,12 +959,14 @@ def worker(args):
sp.kill()
+@cli_utils.action_logging
def initdb(args): # noqa
print("DB: " + repr(settings.engine.url))
db_utils.initdb(settings.RBAC)
print("Done.")
+@cli_utils.action_logging
def resetdb(args):
print("DB: " + repr(settings.engine.url))
if args.yes or input(
@@ -953,6 +977,7 @@ def resetdb(args):
print("Bail.")
+@cli_utils.action_logging
def upgradedb(args): # noqa
print("DB: " + repr(settings.engine.url))
db_utils.upgradedb()
@@ -970,6 +995,7 @@ def upgradedb(args): # noqa
session.commit()
+@cli_utils.action_logging
def version(args): # noqa
print(settings.HEADER + " v" + airflow.__version__)
@@ -978,6 +1004,7 @@ alternative_conn_specs = ['conn_type', 'conn_host',
'conn_login', 'conn_password', 'conn_schema', 'conn_port']
+@cli_utils.action_logging
def connections(args):
if args.list:
# Check that no other flags were passed to the command
@@ -1098,6 +1125,7 @@ def connections(args):
return
+@cli_utils.action_logging
def flower(args):
broka = conf.get('celery', 'BROKER_URL')
address = '--address={}'.format(args.hostname)
@@ -1139,6 +1167,7 @@ def flower(args):
broka, address, port, api, flower_conf, url_prefix])
+@cli_utils.action_logging
def kerberos(args): # noqa
print(settings.HEADER)
import airflow.security.kerberos
@@ -1167,6 +1196,7 @@ def kerberos(args): # noqa
airflow.security.kerberos.run()
+@cli_utils.action_logging
def create_user(args):
fields = {
'role': args.role,
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/9dba430b/airflow/settings.py
----------------------------------------------------------------------
diff --git a/airflow/settings.py b/airflow/settings.py
index 06748ce..1abafa4 100644
--- a/airflow/settings.py
+++ b/airflow/settings.py
@@ -188,9 +188,17 @@ def configure_adapters():
pass
+def configure_action_logging():
+ """
+ Any additional configuration (register callback) for airflow.utils.action_loggers
+ module
+ :return: None
+ """
+ pass
+
+
try:
from airflow_local_settings import *
-
log.info("Loaded airflow_local_settings.")
except:
pass
@@ -199,6 +207,7 @@ configure_logging()
configure_vars()
configure_adapters()
configure_orm()
+configure_action_logging()
# Ensure we close DB connections at scheduler and gunicon worker terminations
atexit.register(dispose_orm)
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/9dba430b/airflow/utils/cli.py
----------------------------------------------------------------------
diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py
new file mode 100644
index 0000000..5f71bca
--- /dev/null
+++ b/airflow/utils/cli.py
@@ -0,0 +1,114 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""
+Utilities module for cli
+"""
+from __future__ import absolute_import
+
+import functools
+import getpass
+import json
+import socket
+import sys
+from argparse import Namespace
+from datetime import datetime
+
+import airflow.models
+from airflow.utils import cli_action_loggers
+
+
+def action_logging(f):
+ """
+ Decorates function to execute function at the same time submitting action_logging
+ but in CLI context. It will call action logger callbacks twice,
+ one for pre-execution and the other one for post-execution.
+
+ Action logger will be called with below keyword parameters:
+ sub_command : name of sub-command
+ start_datetime : start datetime instance by utc
+ end_datetime : end datetime instance by utc
+ full_command : full command line arguments
+ user : current user
+ log : airflow.models.Log ORM instance
+ dag_id : dag id (optional)
+ task_id : task_id (optional)
+ execution_date : execution date (optional)
+ error : exception instance if there's an exception
+
+ :param f: function instance
+ :return: wrapped function
+ """
+ @functools.wraps(f)
+ def wrapper(*args, **kwargs):
+ """
+ An wrapper for cli functions. It assumes to have Namespace instance
+ at 1st positional argument
+ :param args: Positional argument. It assumes to have Namespace instance
+ at 1st positional argument
+ :param kwargs: A passthrough keyword argument
+ """
+ assert args
+ assert isinstance(args[0], Namespace), \
+ "1st positional argument should be argparse.Namespace instance, " \
+ "but {}".format(args[0])
+ metrics = _build_metrics(f.__name__, args[0])
+ cli_action_loggers.on_pre_execution(**metrics)
+ try:
+ return f(*args, **kwargs)
+ except Exception as e:
+ metrics['error'] = e
+ raise e
+ finally:
+ metrics['end_datetime'] = datetime.utcnow()
+ cli_action_loggers.on_post_execution(**metrics)
+
+ return wrapper
+
+
+def _build_metrics(func_name, namespace):
+ """
+ Builds metrics dict from function args
+ It assumes that function arguments is from airflow.bin.cli module's function
+ and has Namespace instance where it optionally contains "dag_id", "task_id",
+ and "execution_date".
+
+ :param func_name: name of function
+ :param namespace: Namespace instance from argparse
+ :return: dict with metrics
+ """
+
+ metrics = {'sub_command': func_name}
+ metrics['start_datetime'] = datetime.utcnow()
+ metrics['full_command'] = '{}'.format(list(sys.argv))
+ metrics['user'] = getpass.getuser()
+
+ assert isinstance(namespace, Namespace)
+ tmp_dic = vars(namespace)
+ metrics['dag_id'] = tmp_dic.get('dag_id')
+ metrics['task_id'] = tmp_dic.get('task_id')
+ metrics['execution_date'] = tmp_dic.get('execution_date')
+ metrics['host_name'] = socket.gethostname()
+
+ extra = json.dumps(dict((k, metrics[k]) for k in ('host_name', 'full_command')))
+ log = airflow.models.Log(
+ event='cli_{}'.format(func_name),
+ task_instance=None,
+ owner=metrics['user'],
+ extra=extra,
+ task_id=metrics.get('task_id'),
+ dag_id=metrics.get('dag_id'),
+ execution_date=metrics.get('execution_date'))
+ metrics['log'] = log
+ return metrics
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/9dba430b/airflow/utils/cli_action_loggers.py
----------------------------------------------------------------------
diff --git a/airflow/utils/cli_action_loggers.py b/airflow/utils/cli_action_loggers.py
new file mode 100644
index 0000000..1bc6865
--- /dev/null
+++ b/airflow/utils/cli_action_loggers.py
@@ -0,0 +1,101 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""
+An Action Logger module. Singleton pattern has been applied into this module
+so that registered callbacks can be used all through the same python process.
+"""
+from __future__ import absolute_import
+
+import logging
+
+import airflow.settings
+
+
+def register_pre_exec_callback(action_logger):
+ """
+ Registers more action_logger function callback for pre-execution.
+ This function callback is expected to be called with keyword args.
+ For more about the arguments that is being passed to the callback,
+ refer to airflow.utils.cli.action_logging()
+ :param action_logger: An action logger function
+ :return: None
+ """
+ logging.debug("Adding {} to pre execution callback".format(action_logger))
+ __pre_exec_callbacks.append(action_logger)
+
+
+def register_post_exec_callback(action_logger):
+ """
+ Registers more action_logger function callback for post-execution.
+ This function callback is expected to be called with keyword args.
+ For more about the arguments that is being passed to the callback,
+ refer to airflow.utils.cli.action_logging()
+ :param action_logger: An action logger function
+ :return: None
+ """
+ logging.debug("Adding {} to post execution callback".format(action_logger))
+ __post_exec_callbacks.append(action_logger)
+
+
+def on_pre_execution(**kwargs):
+ """
+ Calls callbacks before execution.
+ Note that any exception from callback will be logged but won't be propagated.
+ :param kwargs:
+ :return: None
+ """
+ logging.debug("Calling callbacks: {}".format(__pre_exec_callbacks))
+ for cb in __pre_exec_callbacks:
+ try:
+ cb(**kwargs)
+ except Exception:
+ logging.exception('Failed on pre-execution callback using {}'.format(cb))
+
+
+def on_post_execution(**kwargs):
+ """
+ Calls callbacks after execution.
+ As it's being called after execution, it can capture status of execution,
+ duration, etc. Note that any exception from callback will be logged but
+ won't be propagated.
+ :param kwargs:
+ :return: None
+ """
+ logging.debug("Calling callbacks: {}".format(__post_exec_callbacks))
+ for cb in __post_exec_callbacks:
+ try:
+ cb(**kwargs)
+ except Exception:
+ logging.exception('Failed on post-execution callback using {}'.format(cb))
+
+
+def default_action_log(log, **_):
+ """
+ A default action logger callback that behave same as www.utils.action_logging
+ which uses global session and pushes log ORM object.
+ :param log: An log ORM instance
+ :param **_: other keyword arguments that is not being used by this function
+ :return: None
+ """
+ session = airflow.settings.Session()
+ session.add(log)
+ session.commit()
+
+
+__pre_exec_callbacks = []
+__post_exec_callbacks = []
+
+# By default, register default action log into pre-execution callback
+register_pre_exec_callback(default_action_log)
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/9dba430b/tests/utils/test_cli_util.py
----------------------------------------------------------------------
diff --git a/tests/utils/test_cli_util.py b/tests/utils/test_cli_util.py
new file mode 100644
index 0000000..c4a76cc
--- /dev/null
+++ b/tests/utils/test_cli_util.py
@@ -0,0 +1,97 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import unittest
+from argparse import Namespace
+from contextlib import contextmanager
+from datetime import datetime
+
+from airflow.utils import cli, cli_action_loggers
+
+
+class CliUtilTest(unittest.TestCase):
+
+ def test_metrics_build(self):
+ func_name = 'test'
+ exec_date = datetime.utcnow()
+ ns = Namespace(dag_id='foo', task_id='bar',
+ subcommand='test', execution_date=exec_date)
+ metrics = cli._build_metrics(func_name, ns)
+
+ expected = {'user': os.environ.get('USER'),
+ 'sub_command': 'test',
+ 'dag_id': 'foo',
+ 'task_id': 'bar',
+ 'execution_date': exec_date}
+ for k, v in expected.items():
+ self.assertEquals(v, metrics.get(k))
+
+ self.assertTrue(metrics.get('start_datetime') <= datetime.utcnow())
+ self.assertTrue(metrics.get('full_command'))
+
+ log_dao = metrics.get('log')
+ self.assertTrue(log_dao)
+ self.assertEquals(log_dao.dag_id, metrics.get('dag_id'))
+ self.assertEquals(log_dao.task_id, metrics.get('task_id'))
+ self.assertEquals(log_dao.execution_date, metrics.get('execution_date'))
+ self.assertEquals(log_dao.owner, metrics.get('user'))
+
+ def test_fail_function(self):
+ """
+ Actual function is failing and fail needs to be propagated.
+ :return:
+ """
+ with self.assertRaises(NotImplementedError):
+ fail_func(Namespace())
+
+ def test_success_function(self):
+ """
+ Test success function but with failing callback.
+ In this case, failure should not propagate.
+ :return:
+ """
+ with fail_action_logger_callback():
+ success_func(Namespace())
+
+
+@contextmanager
+def fail_action_logger_callback():
+ """
+ Adding failing callback and revert it back when closed.
+ :return:
+ """
+ tmp = cli_action_loggers.__pre_exec_callbacks[:]
+
+ def fail_callback(**_):
+ raise NotImplementedError
+
+ cli_action_loggers.register_pre_exec_callback(fail_callback)
+ yield
+ cli_action_loggers.__pre_exec_callbacks = tmp
+
+
+@cli.action_logging
+def fail_func(_):
+ raise NotImplementedError
+
+
+@cli.action_logging
+def success_func(_):
+ pass
+
+
+if __name__ == '__main__':
+ unittest.main()