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)