def _task_decorator(*decorator_args, **decorator_kwargs):
task_type = decorator_kwargs.pop(
"_task_type") # type: Type[_DecoratedTask]
task_default_result = decorator_kwargs.pop(
"_task_default_result") # ParameterFactory
task_defaults = decorator_kwargs.pop("defaults", None)
if not is_databand_enabled():
# simple `@task` decorator, no options were (probably) given.
if len(decorator_args) == 1 and callable(decorator_args[0]):
return decorator_args[0]
return __passthrough_decorator
def decorated(item):
try:
func_spec = build_task_decorator_spec(
item=item,
decorator_kwargs=decorator_kwargs,
default_result=task_default_result,
)
class DynamicFuncAsTask(task_type):
_conf__decorator_spec = func_spec
_callable_item = None
__doc__ = item.__doc__
__module__ = item.__module__
defaults = task_defaults
# we can't create class dynamically because of python2/3
# __name__ is not overridable
task_cls = DynamicFuncAsTask
except Exception as ex:
logger.error(
"Failed to create task %s: %s\n%s\n",
item.__name__,
str(ex),
user_side_code(context=5),
exc_info=show_exc_info(ex),
)
raise
if func_spec.is_class:
callable_item = six.add_metaclass(_DecoratedUserClassMeta)(item)
else:
callable_item = _decorated_user_func(task_cls=task_cls)
task_cls._callable_item = callable_item
callable_item.func = item
callable_item.task_cls = task_cls
callable_item.task = task_cls
callable_item.t = task_cls
return task_cls._callable_item
# simple `@task` decorator, no options were (probably) given.
if len(decorator_args) == 1 and callable(decorator_args[0]):
return decorated(decorator_args[0])
return decorated
def build_task_decorator(*decorator_args, **decorator_kwargs):
# this code creates a new decorator that can be applied on any User Code
if not is_databand_enabled():
# simple `@task` decorator, no options were (probably) given.
if len(decorator_args) == 1 and callable(decorator_args[0]):
return decorator_args[0]
return _do_nothing_decorator
def class_or_func_decorator(class_or_func):
# this code will run during compile time, when we apply dbnd decorator (for example: @task)
task_decorator = TaskDecorator(class_or_func,
decorator_kwargs=decorator_kwargs)
tp = task_decorator.task_passport
# we need to manually register the task here, since in regular flow
# this happens in TaskMetaclass, but it's not invoked here due to lazy
# evaluation task_cls
r = get_task_registry()
r.register_task_cls_factory(
task_cls_factory=task_decorator.get_task_cls,
full_task_family=tp.full_task_family,
task_family=tp.task_family,
)
if task_decorator.is_class:
# we will change metaclass for UserClass so we will process all UserClass calls
#
# @task
# class UserClass():
# pass
# so the the moment user call UserClass(), -> _DecoratedUserClassMeta.__call__ will be called
dbnd_decorated_class = six.add_metaclass(
_UserClassWithTaskDecoratorMetaclass)(class_or_func)
dbnd_decorated_class.task_decorator = task_decorator
task_decorator.class_or_func = dbnd_decorated_class
return dbnd_decorated_class
else:
# @task
# def user_func():
# pass
# we will return our wrapper, that will be called during a runtime,
# when user calls his own code.
return build_dbnd_decorated_func(task_decorator)
# simple `@task` decorator in opposite to @task(...), no options were (probably) given.
if len(decorator_args) == 1 and callable(decorator_args[0]):
return class_or_func_decorator(decorator_args[0])
return class_or_func_decorator
def _call_handler(cls, call_user_code, call_args, call_kwargs):
"""
-= Use "Step into My Code"" to get back from Databand code! =-
decorated object call/creation ( my_func(), MyDecoratedTask()
"""
force_invoke = call_kwargs.pop("__force_invoke", False)
if force_invoke or not is_databand_enabled():
# 1. Databand is not enabled
# 2. we have this call coming from Task.run / Task.band direct invocation
return call_user_code(*call_args, **call_kwargs)
func_call = FuncCall(
task_cls=cls,
call_args=call_args,
call_kwargs=call_kwargs,
call_user_code=call_user_code,
)
if is_in_airflow_dag_build_context(
): # we are in Airflow DAG building mode
return build_task_at_airflow_dag_context(task_cls=cls,
call_args=call_args,
call_kwargs=call_kwargs)
airflow_task_context = try_get_airflow_context()
if airflow_task_context:
return track_airflow_dag_run_operator_run(
func_call=func_call, airflow_task_context=airflow_task_context)
current = try_get_current_task()
if not current and is_inplace_run():
from dbnd._core.inplace_run.inplace_run_manager import dbnd_run_start
task_run = dbnd_run_start()
if task_run:
current = task_run.task
if not current: # direct call to the function
return func_call.invoke()
######
# DBND HANDLING OF CALL
# now we can make some decisions what we do with the call
# it's not coming from _invoke_func
# but from user code ... some_func() or SomeTask()
phase = current_phase()
if phase is TaskContextPhase.BUILD:
# we are in the @pipeline context, we are building execution plan
t = cls(*call_args, **call_kwargs)
# we are in inline debug mode -> we are going to execute the task
# we are in the band
# and want to return result of the object
if t.task_definition.single_result_output:
return t.result
# we have multiple outputs ( result, another output.. )
# -> just return task object
return t
if phase is TaskContextPhase.RUN:
# we are in the run function!
if (current.settings.dynamic_task.enabled
and current.task_supports_dynamic_tasks):
# isinstance() check required to prevent infinite recursion when @task is on
# class and not on func (example: see test_task_decorated_class.py)
# and the current task supports inline calls
# that's extra mechanism in addition to __force_invoke
# on pickle/unpickle isinstance fails to run.
return create_and_run_dynamic_task_safe(func_call=func_call)
# we can not call it in"databand" way, fallback to normal execution
return func_call.invoke()
def _task_decorator(*decorator_args, **decorator_kwargs):
if not is_databand_enabled():
# simple `@task` decorator, no options were (probably) given.
if len(decorator_args) == 1 and callable(decorator_args[0]):
return decorator_args[0]
return _passthrough_decorator
task_type = decorator_kwargs.pop(
"_task_type") # type: Type[_DecoratedTask]
task_default_result = decorator_kwargs.pop(
"_task_default_result") # ParameterFactory
task_defaults = decorator_kwargs.pop("defaults", None)
def decorated(class_or_func):
try:
func_spec = build_task_decorator_spec(
class_or_func=class_or_func,
decorator_kwargs=decorator_kwargs,
default_result=task_default_result,
)
except Exception as ex:
logger.error(
"Failed to create task %s: %s\n%s\n",
class_or_func.__name__,
str(ex),
user_side_code(context=5),
exc_info=show_exc_info(ex),
)
raise
fp = TaskClsBuilder(func_spec, task_type, task_defaults)
if func_spec.is_class:
wrapper = six.add_metaclass(_DecoratedUserClassMeta)(class_or_func)
fp._callable_item = wrapper
else:
@functools.wraps(class_or_func)
def wrapper(*args, **kwargs):
if in_tracking_mode():
with fp.tracking_context(args,
kwargs) as track_result_callback:
return track_result_callback(fp.func(*args, **kwargs))
return _call_handler(
fp.get_task_cls(),
call_user_code=fp.func,
call_args=args,
call_kwargs=kwargs,
)
wrapper.dbnd_run = fp.dbnd_run
wrapper.__is_dbnd_task__ = True
wrapper.func = class_or_func
# we're using CallableLazyObjectProxy to have lazy evaluation for creating task_cls
# this is only orchestration scenarios
task_cls = CallableLazyObjectProxy(fp.get_task_cls)
wrapper.task_cls = task_cls
wrapper.task = task_cls
wrapper.t = task_cls
# we need lazy task_definition here, for example for dbnd_task_as_bash_operator
wrapper.task_definition = CallableLazyObjectProxy(
fp.get_task_definition)
# we need to manually register the task here, since in regular flow
# this happens in TaskMetaclass, but it's not invoked here due to lazy
# evaluation using CallableLazyObjectProxy
tp = TaskPassport.from_func_spec(func_spec, decorator_kwargs)
# TODO: we can use CallableLazyObjectProxy object (task_cls) instead of task_cls_factory
r = get_task_registry()
r.register_task_cls_factory(
task_cls_factory=fp.get_task_cls,
full_task_family=tp.full_task_family,
task_family=tp.task_family,
)
return wrapper
# simple `@task` decorator, no options were (probably) given.
if len(decorator_args) == 1 and callable(decorator_args[0]):
return decorated(decorator_args[0])
return decorated