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 __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 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 try_get_databand_context():
# type: () -> Optional[DatabandContext]
from dbnd._core.context.databand_context import DatabandContext # noqa: F811
if not DatabandContext.has_instance():
return None
return get_databand_context()
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 dbnd_context():
context = try_get_databand_context()
if not context:
# we are running without Databand Context
# let create one inplace
from dbnd._core.context.databand_context import DatabandContext # noqa: F811
context = DatabandContext.try_instance(name="inplace_run")
return context
def dbnd_context():
"""Get or create databand context."""
context = try_get_databand_context()
if not context:
# we are running without Databand Context
# let create one inplace
from dbnd._core.context.databand_context import (
DatabandContext as _DatabandContext,
)
context = _DatabandContext.try_instance(name="inplace_run")
return context
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
def get_databand_context():
# type: () -> DatabandContext
from dbnd._core.context.databand_context import DatabandContext # noqa: F811
return DatabandContext.try_instance()
def get_databand_context():
# type: () -> DatabandContext
"""Databand Context get or create instance."""
from dbnd._core.context.databand_context import DatabandContext as _DatabandContext
return _DatabandContext.try_instance()
