def main(args=None):
# from dbnd._core.log.config import configure_basic_logging
# configure_basic_logging(None)
dbnd_system_bootstrap()
# LET'S PATCH AIRFLOW FIRST
from dbnd_airflow.bootstrap import dbnd_airflow_bootstrap
dbnd_airflow_bootstrap()
from airflow.bin.cli import CLIFactory
from airflow.configuration import conf
from dbnd_airflow.plugins.setup_plugins import (
setup_scheduled_dags,
setup_versioned_dags,
)
# ORIGINAL CODE from airflow/bin/airflow
if conf.get("core", "security") == "kerberos":
os.environ["KRB5CCNAME"] = conf.get("kerberos", "ccache")
os.environ["KRB5_KTNAME"] = conf.get("kerberos", "keytab")
import argcomplete
from dbnd_airflow.scheduler.zombies import find_and_kill_dagrun_zombies
CLIFactory.subparsers_dict[find_and_kill_dagrun_zombies.__name__] = {
"func": find_and_kill_dagrun_zombies,
"help": "Clean up BackfillJob zombie tasks",
"args": tuple(),
}
parser = CLIFactory.get_parser()
argcomplete.autocomplete(parser)
args = parser.parse_args(args=args)
func_name = args.func.__name__
# DBND PATCH:
if dbnd_config.getboolean("airflow",
"auto_add_scheduled_dags") and func_name in [
"scheduler",
"webserver",
]:
setup_scheduled_dags()
if dbnd_config.getboolean(
"airflow",
"auto_add_versioned_dags") and func_name in ["webserver"]:
setup_versioned_dags()
args.func(args)
if func_name in ["resetdb", "initdb"]:
pool_name = dbnd_config.get("airflow", "dbnd_pool")
if pool_name == "dbnd_pool":
create_airflow_pool(pool_name)
def dbnd_operator__execute(dbnd_operator, context):
from dbnd._core.current import try_get_databand_run
from dbnd._core.run.databand_run import DatabandRun
from targets import target
run = try_get_databand_run()
if not run:
# we are not inside dbnd run, probably we are running from native airflow
# let's try to load it:
try:
executor_config = dbnd_operator.executor_config
logger.info("context: %s", context)
logger.info("task.executor_config: %s",
dbnd_operator.executor_config)
logger.info("ti.executor_config: %s",
context["ti"].executor_config)
driver_dump = executor_config["DatabandExecutor"].get(
"dbnd_driver_dump")
print(
"Running dbnd task %s from %s" %
(dbnd_operator.dbnd_task_id, driver_dump),
file=sys.__stderr__,
)
if executor_config["DatabandExecutor"].get(
"remove_airflow_std_redirect", False):
sys.stdout = sys.__stdout__
sys.stderr = sys.__stderr__
dbnd_bootstrap()
dbnd_airflow_bootstrap()
run = DatabandRun.load_run(dump_file=target(driver_dump),
disable_tracking_api=False)
except Exception as e:
print(
"Failed to load dbnd task in native airflow execution! Exception: %s"
% (e, ),
file=sys.__stderr__,
)
dump_trace()
raise
with run.run_context() as dr:
task_run = run.get_task_run_by_id(dbnd_operator.dbnd_task_id)
ret_value = task_run.runner.execute(airflow_context=context)
else:
task_run = run.get_task_run_by_id(dbnd_operator.dbnd_task_id)
ret_value = task_run.runner.execute(airflow_context=context)
return ret_value
def patched_create_app(*args, **kwargs):
from dbnd._core.configuration.dbnd_config import config
logger.info("Setting SQL from databand configuration.")
config.load_system_configs()
from dbnd_airflow.bootstrap import dbnd_airflow_bootstrap
dbnd_airflow_bootstrap()
res = create_app_func(*args, **kwargs)
try:
use_databand_airflow_dagbag()
except Exception:
logger.info("Failed to apply dbnd versioned dagbag")
return res
def __init__(self, run, task_executor_type, host_engine, target_engine,
task_runs):
super(AirflowTaskExecutor, self).__init__(
run=run,
task_executor_type=task_executor_type,
host_engine=host_engine,
target_engine=target_engine,
task_runs=task_runs,
)
from dbnd_airflow.bootstrap import dbnd_airflow_bootstrap
dbnd_airflow_bootstrap()
self.airflow_config = AirflowConfig()
self.airflow_task_executor = self._get_airflow_executor()
logger.info("Using airflow executor: %s" %
self.airflow_task_executor.__class__.__name__)
def __init__(self, run, task_executor_type, host_engine, target_engine, task_runs):
super(AirflowTaskExecutor, self).__init__(
run=run,
task_executor_type=task_executor_type,
host_engine=host_engine,
target_engine=target_engine,
task_runs=task_runs,
)
# we want to use reference to AIRFLOW HOME as it can be changed in runtime
from dbnd_airflow.bootstrap import dbnd_airflow_bootstrap
dbnd_airflow_bootstrap()
self.airflow_config = AirflowConfig()
self.airflow_task_executor = self._get_airflow_executor()
self._validate_airflow_db()
def main(args=None):
# from dbnd._core.log.config import configure_basic_logging
# configure_basic_logging(None)
dbnd_system_bootstrap()
# LET'S PATCH AIRFLOW FIRST
from dbnd_airflow.bootstrap import dbnd_airflow_bootstrap
dbnd_airflow_bootstrap()
from airflow.bin.cli import CLIFactory
from airflow.configuration import conf
from dbnd_airflow.plugins.setup_plugins import (
setup_scheduled_dags,
setup_versioned_dags,
)
# ORIGINAL CODE from airflow/bin/airflow
if conf.get("core", "security") == "kerberos":
os.environ["KRB5CCNAME"] = conf.get("kerberos", "ccache")
os.environ["KRB5_KTNAME"] = conf.get("kerberos", "keytab")
import argcomplete
parser = CLIFactory.get_parser()
argcomplete.autocomplete(parser)
args = parser.parse_args(args=args)
func_name = args.func.__name__
# DBND PATCH:
if dbnd_config.getboolean("airflow",
"auto_add_scheduled_dags") and func_name in [
"scheduler",
"webserver",
]:
setup_scheduled_dags()
if dbnd_config.getboolean(
"airflow",
"auto_add_versioned_dags") and func_name in ["webserver"]:
setup_versioned_dags()
args.func(args)
