def pool_import(args):
"""Imports pools from the file"""
api_client = get_current_api_client()
log = LoggingMixin().log
if os.path.exists(args.file):
pools = pool_import_helper(args.file)
else:
print("Missing pools file.")
pools = api_client.get_pools()
log.info(_tabulate_pools(pools=pools, tablefmt=args.output))
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 task_run(args, dag=None):
"""Runs a single task instance"""
if dag:
args.dag_id = dag.dag_id
log = LoggingMixin().log
# Load custom airflow config
if args.cfg_path:
with open(args.cfg_path, 'r') as conf_file:
conf_dict = json.load(conf_file)
if os.path.exists(args.cfg_path):
os.remove(args.cfg_path)
conf.read_dict(conf_dict, source=args.cfg_path)
settings.configure_vars()
# IMPORTANT, have to use the NullPool, otherwise, each "run" command may leave
# behind multiple open sleeping connections while heartbeating, which could
# easily exceed the database connection limit when
# processing hundreds of simultaneous tasks.
settings.configure_orm(disable_connection_pool=True)
if not args.pickle and not dag:
dag = get_dag(args)
elif not dag:
with db.create_session() as session:
log.info('Loading pickle id %s', args.pickle)
dag_pickle = session.query(DagPickle).filter(
DagPickle.id == args.pickle).first()
if not dag_pickle:
raise AirflowException("Who hid the pickle!? [missing pickle]")
dag = dag_pickle.pickle
task = dag.get_task(task_id=args.task_id)
ti = TaskInstance(task, args.execution_date)
ti.refresh_from_db()
ti.init_run_context(raw=args.raw)
hostname = get_hostname()
log.info("Running %s on host %s", ti, hostname)
if args.interactive:
_run(args, dag, ti)
else:
with redirect_stdout(ti.log, logging.INFO), redirect_stderr(
ti.log, logging.WARN):
_run(args, dag, ti)
logging.shutdown()
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 get_default_executor(cls) -> BaseExecutor:
"""Creates a new instance of the configured executor if none exists and returns it"""
if cls._default_executor is not None:
return cls._default_executor
from airflow.configuration import conf
executor_name = conf.get('core', 'EXECUTOR')
cls._default_executor = ExecutorLoader._get_executor(executor_name)
from airflow import LoggingMixin
log = LoggingMixin().log
log.info("Using executor %s", executor_name)
return cls._default_executor
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 create_app(config=None, testing=False):
log = LoggingMixin().log
app = Flask(__name__)
app.wsgi_app = ProxyFix(app.wsgi_app)
if configuration.conf.get('webserver', 'SECRET_KEY') == "temporary_key":
log.info("SECRET_KEY for Flask App is not specified. Using a random one.")
app.secret_key = os.urandom(16)
else:
app.secret_key = configuration.conf.get('webserver', 'SECRET_KEY')
app.config['LOGIN_DISABLED'] = not configuration.conf.getboolean(
'webserver', 'AUTHENTICATE')
csrf.init_app(app)
app.config['TESTING'] = testing
airflow.load_login()
airflow.login.login_manager.init_app(app)
from airflow import api
api.load_auth()
api.api_auth.init_app(app)
cache = Cache(
app=app, config={'CACHE_TYPE': 'filesystem', 'CACHE_DIR': '/tmp'})
app.register_blueprint(routes)
configure_logging()
with app.app_context():
from airflow.www import views
admin = Admin(
app, name='Airflow',
static_url_path='/admin',
index_view=views.HomeView(endpoint='', url='/admin', name="DAGs"),
template_mode='bootstrap3',
)
av = admin.add_view
vs = views
av(vs.Airflow(name='DAGs', category='DAGs'))
if not conf.getboolean('core', 'secure_mode'):
av(vs.QueryView(name='Ad Hoc Query', category="Data Profiling"))
av(vs.ChartModelView(
models.Chart, Session, name="Charts", category="Data Profiling"))
av(vs.KnownEventView(
models.KnownEvent,
Session, name="Known Events", category="Data Profiling"))
av(vs.SlaMissModelView(
models.SlaMiss,
Session, name="SLA Misses", category="Browse"))
av(vs.TaskInstanceModelView(models.TaskInstance,
Session, name="Task Instances", category="Browse"))
av(vs.LogModelView(
models.Log, Session, name="Logs", category="Browse"))
av(vs.JobModelView(
jobs.BaseJob, Session, name="Jobs", category="Browse"))
av(vs.PoolModelView(
models.Pool, Session, name="Pools", category="Admin"))
av(vs.ConfigurationView(
name='Configuration', category="Admin"))
av(vs.UserModelView(
models.User, Session, name="Users", category="Admin"))
av(vs.ConnectionModelView(
models.Connection, Session, name="Connections", category="Admin"))
av(vs.VariableView(
models.Variable, Session, name="Variables", category="Admin"))
av(vs.XComView(
models.XCom, Session, name="XComs", category="Admin"))
admin.add_link(base.MenuLink(
category='Docs', name='Documentation',
url='https://airflow.incubator.apache.org/'))
admin.add_link(
base.MenuLink(category='Docs',
name='Github',
url='https://github.com/apache/incubator-airflow'))
av(vs.VersionView(name='Version', category="About"))
av(vs.DagRunModelView(
models.DagRun, Session, name="DAG Runs", category="Browse"))
av(vs.DagModelView(models.DagModel, Session, name=None))
# Hack to not add this view to the menu
admin._menu = admin._menu[:-1]
def integrate_plugins():
"""Integrate plugins to the context"""
from airflow.plugins_manager import (
admin_views, flask_blueprints, menu_links)
for v in admin_views:
log.debug('Adding view %s', v.name)
admin.add_view(v)
for bp in flask_blueprints:
log.debug('Adding blueprint %s', bp.name)
app.register_blueprint(bp)
for ml in sorted(menu_links, key=lambda x: x.name):
log.debug('Adding menu link %s', ml.name)
admin.add_link(ml)
integrate_plugins()
import airflow.www.api.experimental.endpoints as e
# required for testing purposes otherwise the module retains
# a link to the default_auth
if app.config['TESTING']:
if six.PY2:
reload(e)
else:
import importlib
importlib.reload(e)
app.register_blueprint(e.api_experimental, url_prefix='/api/experimental')
@app.context_processor
def jinja_globals():
return {
'hostname': get_hostname(),
'navbar_color': configuration.get('webserver', 'NAVBAR_COLOR'),
}
@app.teardown_appcontext
def shutdown_session(exception=None):
settings.Session.remove()
return app
def clean_before_retry(context):
logger = LoggingMixin().log
logger.info("executing default retry handler")
if 'retry_command' in context['params']:
bash_command = context['params']['retry_command']
if 'retry_java_args_method' in context['params']:
bash_command = bash_command + ' ' + context['params'][
'retry_java_args_method'](context)
logger.info("tmp dir root location: \n" + gettempdir())
task_instance_key_str = context['task_instance_key_str']
with TemporaryDirectory(prefix='airflowtmp') as tmp_dir:
with NamedTemporaryFile(dir=tmp_dir,
prefix=("retry_%s" %
task_instance_key_str)) as f:
f.write(bash_command)
f.flush()
fname = f.name
script_location = tmp_dir + "/" + fname
logger.info("Temporary script "
"location :{0}".format(script_location))
logger.info("Running retry command: " + bash_command)
sp = Popen(['bash', fname],
stdout=PIPE,
stderr=STDOUT,
cwd=tmp_dir,
preexec_fn=os.setsid)
logger.info("Retry command output:")
line = ''
for line in iter(sp.stdout.readline, b''):
line = line.decode("UTF-8").strip()
logger.info(line)
sp.wait()
logger.info("Retry command exited with "
"return code {0}".format(sp.returncode))
if sp.returncode:
raise AirflowException("Retry bash command failed")
def pool_export(args):
"""Exports all of the pools to the file"""
log = LoggingMixin().log
pools = pool_export_helper(args.file)
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_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))
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 create_app(config=None, testing=False):
log = LoggingMixin().log
app = Flask(__name__)
app.wsgi_app = ProxyFix(app.wsgi_app)
if configuration.conf.get('webserver', 'SECRET_KEY') == "temporary_key":
log.info(
"SECRET_KEY for Flask App is not specified. Using a random one.")
app.secret_key = os.urandom(16)
else:
app.secret_key = configuration.conf.get('webserver', 'SECRET_KEY')
app.config['LOGIN_DISABLED'] = not configuration.conf.getboolean(
'webserver', 'AUTHENTICATE')
csrf.init_app(app)
app.config['TESTING'] = testing
airflow.load_login()
airflow.login.login_manager.init_app(app)
from airflow import api
api.load_auth()
api.api_auth.init_app(app)
cache = Cache(app=app,
config={
'CACHE_TYPE': 'filesystem',
'CACHE_DIR': '/tmp'
})
app.register_blueprint(routes)
configure_logging()
with app.app_context():
from airflow.www import views
admin = Admin(
app,
name='Airflow',
static_url_path='/admin',
index_view=views.HomeView(endpoint='', url='/admin', name="DAGs"),
template_mode='bootstrap3',
)
av = admin.add_view
vs = views
av(vs.Airflow(name='DAGs', category='DAGs'))
if not conf.getboolean('core', 'secure_mode'):
av(vs.QueryView(name='Ad Hoc Query', category="Data Profiling"))
av(
vs.ChartModelView(models.Chart,
Session,
name="Charts",
category="Data Profiling"))
av(
vs.KnownEventView(models.KnownEvent,
Session,
name="Known Events",
category="Data Profiling"))
av(
vs.SlaMissModelView(models.SlaMiss,
Session,
name="SLA Misses",
category="Browse"))
av(
vs.TaskInstanceModelView(models.TaskInstance,
Session,
name="Task Instances",
category="Browse"))
av(vs.LogModelView(models.Log, Session, name="Logs",
category="Browse"))
av(
vs.JobModelView(jobs.BaseJob,
Session,
name="Jobs",
category="Browse"))
av(
vs.PoolModelView(models.Pool,
Session,
name="Pools",
category="Admin"))
av(vs.ConfigurationView(name='Configuration', category="Admin"))
av(
vs.UserModelView(models.User,
Session,
name="Users",
category="Admin"))
av(
vs.ConnectionModelView(models.Connection,
Session,
name="Connections",
category="Admin"))
av(
vs.VariableView(models.Variable,
Session,
name="Variables",
category="Admin"))
av(vs.XComView(models.XCom, Session, name="XComs", category="Admin"))
admin.add_link(
base.MenuLink(category='Docs',
name='Documentation',
url='https://airflow.incubator.apache.org/'))
admin.add_link(
base.MenuLink(category='Docs',
name='Github',
url='https://github.com/apache/incubator-airflow'))
av(vs.VersionView(name='Version', category="About"))
av(
vs.DagRunModelView(models.DagRun,
Session,
name="DAG Runs",
category="Browse"))
av(vs.DagModelView(models.DagModel, Session, name=None))
# Hack to not add this view to the menu
admin._menu = admin._menu[:-1]
def integrate_plugins():
"""Integrate plugins to the context"""
from airflow.plugins_manager import (admin_views, flask_blueprints,
menu_links)
for v in admin_views:
log.debug('Adding view %s', v.name)
admin.add_view(v)
for bp in flask_blueprints:
log.debug('Adding blueprint %s', bp.name)
app.register_blueprint(bp)
for ml in sorted(menu_links, key=lambda x: x.name):
log.debug('Adding menu link %s', ml.name)
admin.add_link(ml)
integrate_plugins()
import airflow.www.api.experimental.endpoints as e
# required for testing purposes otherwise the module retains
# a link to the default_auth
if app.config['TESTING']:
if six.PY2:
reload(e)
else:
import importlib
importlib.reload(e)
app.register_blueprint(e.api_experimental,
url_prefix='/api/experimental')
@app.context_processor
def jinja_globals():
return {
'hostname': get_hostname(),
'navbar_color': configuration.get('webserver', 'NAVBAR_COLOR'),
}
@app.teardown_appcontext
def shutdown_session(exception=None):
settings.Session.remove()
return app