def conf_vars(overrides):
original = {}
original_env_vars = {}
for (section, key), value in overrides.items():
env = conf._env_var_name(section, key)
if env in os.environ:
original_env_vars[env] = os.environ.pop(env)
if conf.has_option(section, key):
original[(section, key)] = conf.get(section, key)
else:
original[(section, key)] = None
if value is not None:
conf.set(section, key, value)
else:
conf.remove_option(section, key)
settings.configure_vars()
try:
yield
finally:
for (section, key), value in original.items():
if value is not None:
conf.set(section, key, value)
else:
conf.remove_option(section, key)
for env, value in original_env_vars.items():
os.environ[env] = value
settings.configure_vars()
def conf_vars(overrides):
original = {}
original_env_vars = {}
reconfigure_vars = False
for (section, key), value in overrides.items():
env = conf._env_var_name(section, key)
if env in os.environ:
original_env_vars[env] = os.environ.pop(env)
if conf.has_option(section, key):
original[(section, key)] = conf.get(section, key)
else:
original[(section, key)] = None
if value is not None:
conf.set(section, key, value)
else:
conf.remove_option(section, key)
if section == 'core' and key.lower().endswith('_folder'):
reconfigure_vars = True
if reconfigure_vars:
settings.configure_vars()
yield
for (section, key), value in original.items():
if value is not None:
conf.set(section, key, value)
else:
conf.remove_option(section, key)
for env, value in original_env_vars.items():
os.environ[env] = value
if reconfigure_vars:
settings.configure_vars()
def reinit_airflow_sql_conn():
from airflow.settings import configure_orm, configure_vars
from dbnd._core.configuration.dbnd_config import config as dbnd_config
configure_vars()
# The webservers import this file from models.py with the default settings.
configure_orm()
# add query handler before every execute
# this will print query, code line and stack trace
if dbnd_config.getboolean("log", "sqlalchemy_trace"):
from airflow import settings as airflow_settings
from sqlalchemy import event
from dbnd_airflow.db_utils import trace_sqlalchemy_query
event.listen(airflow_settings.engine, "before_cursor_execute",
trace_sqlalchemy_query)
# this will print query execution time
from airflow import settings as airflow_settings
from sqlalchemy import event
from dbnd_airflow.db_utils import (
profile_after_cursor_execute,
profile_before_cursor_execute,
)
event.listen(airflow_settings.engine, "before_cursor_execute",
profile_before_cursor_execute)
event.listen(airflow_settings.engine, "after_cursor_execute",
profile_after_cursor_execute)
def set_airflow_db(sql_alchemy_conn: str, fernet_key: str):
with set_env(
AIRFLOW__CORE__SQL_ALCHEMY_CONN=sql_alchemy_conn,
AIRFLOW__CORE__FERNET_KEY=fernet_key,
):
settings.configure_vars()
settings.configure_orm()
assert repr(settings.engine.url) == sql_alchemy_conn
yield
settings.configure_vars()
settings.configure_orm()
def mock_airflow_db() -> ContextManager[AirflowDb]:
with tempfile.TemporaryDirectory() as temp_dir:
test_db_path = os.path.join(temp_dir, 'airflow.db')
sql_alchemy_conn = f'sqlite:///{test_db_path}'
with set_env(AIRFLOW__CORE__SQL_ALCHEMY_CONN=sql_alchemy_conn):
settings.configure_vars()
settings.configure_orm()
assert repr(settings.engine.url) == sql_alchemy_conn
initdb()
yield AirflowDb(sql_alchemy_conn=sql_alchemy_conn)
settings.configure_vars()
settings.configure_orm()
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 task_run(args, dag=None):
"""Runs a single task instance"""
# 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 dag and args.pickle:
raise AirflowException(
"You cannot use the --pickle option when using DAG.cli() method.")
elif args.pickle:
print(f'Loading pickle id: {args.pickle}')
dag = get_dag_by_pickle(args.pickle)
elif not dag:
dag = get_dag(args.subdir, args.dag_id)
else:
# Use DAG from parameter
pass
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()
print(f"Running {ti} on host {hostname}")
if args.interactive:
_run_task_by_selected_method(args, dag, ti)
else:
with redirect_stdout(StreamLogWriter(ti.log, logging.INFO)), \
redirect_stderr(StreamLogWriter(ti.log, logging.WARN)):
_run_task_by_selected_method(args, dag, ti)
logging.shutdown()
def set_airflow_db(sql_alchemy_conn: Optional[str],
fernet_key: Optional[str]) -> ContextManager[AirflowDb]:
env = {}
if sql_alchemy_conn is not None:
env['AIRFLOW__CORE__SQL_ALCHEMY_CONN'] = sql_alchemy_conn
if fernet_key is not None:
env['AIRFLOW__CORE__FERNET_KEY'] = fernet_key
with set_env(**env):
settings.configure_vars()
settings.configure_orm()
if sql_alchemy_conn is not None:
assert str(
settings.engine.url
) == sql_alchemy_conn, f'{settings.engine.url} != {sql_alchemy_conn}'
yield AirflowDb(sql_alchemy_conn or settings.SQL_ALCHEMY_CONN)
settings.configure_vars()
settings.configure_orm()
def task_run(args, dag=None):
"""Runs a single task instance"""
# 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 dag and args.pickle:
raise AirflowException("You cannot use the --pickle option when using DAG.cli() method.")
elif args.pickle:
print(f'Loading pickle id: {args.pickle}')
dag = get_dag_by_pickle(args.pickle)
elif not dag:
dag = get_dag(args.subdir, args.dag_id)
else:
# Use DAG from parameter
pass
task = dag.get_task(task_id=args.task_id)
ti = TaskInstance(task, args.execution_date)
ti.init_run_context(raw=args.raw)
hostname = get_hostname()
print(f"Running {ti} on host {hostname}")
if args.interactive:
_run_task_by_selected_method(args, dag, ti)
else:
if settings.DONOT_MODIFY_HANDLERS:
with redirect_stdout(StreamLogWriter(ti.log, logging.INFO)), \
redirect_stderr(StreamLogWriter(ti.log, logging.WARN)):
_run_task_by_selected_method(args, dag, ti)
else:
# Get all the Handlers from 'airflow.task' logger
# Add these handlers to the root logger so that we can get logs from
# any custom loggers defined in the DAG
airflow_logger_handlers = logging.getLogger('airflow.task').handlers
root_logger = logging.getLogger()
root_logger_handlers = root_logger.handlers
# Remove all handlers from Root Logger to avoid duplicate logs
for handler in root_logger_handlers:
root_logger.removeHandler(handler)
for handler in airflow_logger_handlers:
root_logger.addHandler(handler)
root_logger.setLevel(logging.getLogger('airflow.task').level)
with redirect_stdout(StreamLogWriter(ti.log, logging.INFO)), \
redirect_stderr(StreamLogWriter(ti.log, logging.WARN)):
_run_task_by_selected_method(args, dag, ti)
# We need to restore the handlers to the loggers as celery worker process
# can call this command multiple times,
# so if we don't reset this then logs from next task would go to the wrong place
for handler in airflow_logger_handlers:
root_logger.removeHandler(handler)
for handler in root_logger_handlers:
root_logger.addHandler(handler)
logging.shutdown()
def task_run(args, dag=None):
"""Run a single task instance.
Note that there must be at least one DagRun for this to start,
i.e. it must have been scheduled and/or triggered previously.
Alternatively, if you just need to run it for testing then use
"airflow tasks test ..." command instead.
"""
# Load custom airflow config
if args.local and args.raw:
raise AirflowException(
"Option --raw and --local are mutually exclusive. "
"Please remove one option to execute the command.")
if args.raw:
unsupported_options = [
o for o in RAW_TASK_UNSUPPORTED_OPTION if getattr(args, o)
]
if unsupported_options:
unsupported_raw_task_flags = ', '.join(
f'--{o}' for o in RAW_TASK_UNSUPPORTED_OPTION)
unsupported_flags = ', '.join(f'--{o}'
for o in unsupported_options)
raise AirflowException(
"Option --raw does not work with some of the other options on this command. "
"You can't use --raw option and the following options: "
f"{unsupported_raw_task_flags}. "
f"You provided the option {unsupported_flags}. "
"Delete it to execute the command.")
if dag and args.pickle:
raise AirflowException(
"You cannot use the --pickle option when using DAG.cli() method.")
if args.cfg_path:
with open(args.cfg_path) 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()
settings.MASK_SECRETS_IN_LOGS = True
# 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 args.pickle:
print(f'Loading pickle id: {args.pickle}')
dag = get_dag_by_pickle(args.pickle)
elif not dag:
dag = get_dag(args.subdir, args.dag_id)
else:
# Use DAG from parameter
pass
task = dag.get_task(task_id=args.task_id)
ti = _get_ti(task, args.execution_date_or_run_id)
ti.init_run_context(raw=args.raw)
hostname = get_hostname()
print(f"Running {ti} on host {hostname}")
if args.interactive:
_run_task_by_selected_method(args, dag, ti)
else:
with _capture_task_logs(ti):
_run_task_by_selected_method(args, dag, ti)
