def init_context(self): if not self._databand_context: self._databand_context = DatabandContext(name="luigi") # Must enter context before creating dbnd task! dc = self._enter_cm( DatabandContext.context(_context=self._databand_context) )
def run_context(self): # type: (DatabandRun) -> Iterator[DatabandRun] from dbnd._core.context.databand_context import DatabandContext # noqa: F811 with DatabandContext.context(_context=self.context): with DatabandRun.context(_context=self) as dr: yield dr # type: DatabandRun
def set_active_run_context(run): # type: (DatabandRun) -> Iterator[DatabandRun] from dbnd._core.context.databand_context import DatabandContext # noqa: F811 with DatabandContext.context(_context=run.context): with DatabandRun.context(_context=run) as dr: yield dr
def __init__(self, dag): self.dag = dag self.dbnd_airflow_name = {} config_store = self.get_and_process_dbnd_dag_config() with dbnd_config( config_values=config_store, source="airflow" ) as current_config: self.dbnd_context = DatabandContext(name="airflow__%s" % self.dag.dag_id) with DatabandContext.context(_context=self.dbnd_context): # we need databand context to update config first self.dbnd_config_layer = current_config.config_layer
def build_airflow_operator(self, task_cls, call_args, call_kwargs): if try_get_databand_context() is self.dbnd_context: # we are already in the context of build return self._build_airflow_operator( task_cls=task_cls, call_args=call_args, call_kwargs=call_kwargs ) # we are coming from external world with dbnd_config.config_layer_context( self.dbnd_config_layer ) as c, DatabandContext.context(_context=self.dbnd_context) as dc: return self._build_airflow_operator( task_cls=task_cls, call_args=call_args, call_kwargs=call_kwargs )
def execute(self, context): logger.debug("Running dbnd dbnd_task from airflow operator %s", self.task_id) dag = context["dag"] execution_date = context["execution_date"] dag_id = dag.dag_id run_uid = get_job_run_uid(dag_id=dag_id, execution_date=execution_date) # Airflow has updated all relevant fields in Operator definition with XCom values # now we can create a real dbnd dbnd_task with real references to dbnd_task new_kwargs = {} for p_name in self.dbnd_task_params_fields: new_kwargs[p_name] = getattr(self, p_name, None) # this is the real input value after if p_name in self.dbnd_xcom_inputs: new_kwargs[p_name] = target(new_kwargs[p_name]) new_kwargs["_dbnd_disable_airflow_inplace"] = True dag_ctrl = self.get_dbnd_dag_ctrl() with DatabandContext.context(_context=dag_ctrl.dbnd_context) as dc: logger.debug("Running %s with kwargs=%s ", self.task_id, new_kwargs) dbnd_task = dc.task_instance_cache.get_task_by_id(self.dbnd_task_id) # rebuild task with new values coming from xcom->operator with dbnd_task.ctrl.task_context(phase=TaskContextPhase.BUILD): dbnd_task = dbnd_task.clone( output_params_to_clone=self.dbnd_overridden_output_params, **new_kwargs ) logger.debug("Creating inplace databand run for driver dump") dag_task = Task(task_name=dag.dag_id, task_target_date=execution_date) dag_task.set_upstream(dbnd_task) # create databand run with new_databand_run( context=dc, task_or_task_name=dag_task, run_uid=run_uid, existing_run=False, job_name=dag.dag_id, ) as dr: # type: DatabandRun dr._init_without_run() # dr.driver_task_run.set_task_run_state(state=TaskRunState.RUNNING) # "make dag run" # dr.root_task_run.set_task_run_state(state=TaskRunState.RUNNING) dbnd_task_run = dr.get_task_run_by_id(dbnd_task.task_id) needs_databand_run_save = dbnd_task._conf__require_run_dump_file if needs_databand_run_save: dr.save_run() logger.info( dbnd_task.ctrl.banner( "Running task '%s'." % dbnd_task.task_name, color="cyan" ) ) # should be replaced with tr._execute call dbnd_task_run.runner.execute( airflow_context=context, handle_sigterm=False ) logger.debug("Finished to run %s", self) result = { output_name: convert_to_safe_types(getattr(dbnd_task, output_name)) for output_name in self.dbnd_xcom_outputs } return result