diff --git a/airflow/decorators/__init__.py b/airflow/decorators/__init__.py index 1250f321de291..53691725a2e4c 100644 --- a/airflow/decorators/__init__.py +++ b/airflow/decorators/__init__.py @@ -15,15 +15,23 @@ # specific language governing permissions and limitations # under the License. -from typing import Callable, Dict, Iterable, List, Optional, Union +from typing import Callable, Optional from airflow.decorators.python import python_task from airflow.decorators.python_virtualenv import _virtualenv_task from airflow.decorators.task_group import task_group # noqa +from airflow.exceptions import AirflowException from airflow.models.dag import dag # noqa +from airflow.providers_manager import ProvidersManager class _TaskDecorator: + def __init__(self): + self.store = {"python": python_task, "virtualenv": _virtualenv_task} + decorator = ProvidersManager().taskflow_decorators + for decorator_name, decorator_class in decorator.items(): + self.store[decorator_name] = decorator_class + def __call__( self, python_callable: Optional[Callable] = None, multiple_outputs: Optional[bool] = None, **kwargs ): @@ -41,103 +49,10 @@ def __call__( """ return self.python(python_callable=python_callable, multiple_outputs=multiple_outputs, **kwargs) - @staticmethod - def python(python_callable: Optional[Callable] = None, multiple_outputs: Optional[bool] = None, **kwargs): - """ - Python operator decorator. Wraps a function into an Airflow operator. - Accepts kwargs for operator kwarg. This decorator can be reused in a single DAG. - - :param python_callable: Function to decorate - :type python_callable: Optional[Callable] - :param multiple_outputs: if set, function return value will be - unrolled to multiple XCom values. List/Tuples will unroll to xcom values - with index as key. Dict will unroll to xcom values with keys as XCom keys. - Defaults to False. - :type multiple_outputs: bool - """ - return python_task(python_callable=python_callable, multiple_outputs=multiple_outputs, **kwargs) - - @staticmethod - def virtualenv( - python_callable: Optional[Callable] = None, - multiple_outputs: Optional[bool] = None, - requirements: Optional[Iterable[str]] = None, - python_version: Optional[Union[str, int, float]] = None, - use_dill: bool = False, - system_site_packages: bool = True, - string_args: Optional[Iterable[str]] = None, - templates_dict: Optional[Dict] = None, - templates_exts: Optional[List[str]] = None, - **kwargs, - ): - """ - Allows one to run a function in a virtualenv that is - created and destroyed automatically (with certain caveats). - - The function must be defined using def, and not be - part of a class. All imports must happen inside the function - and no variables outside of the scope may be referenced. A global scope - variable named virtualenv_string_args will be available (populated by - string_args). In addition, one can pass stuff through op_args and op_kwargs, and one - can use a return value. - Note that if your virtualenv runs in a different Python major version than Airflow, - you cannot use return values, op_args, op_kwargs, or use any macros that are being provided to - Airflow through plugins. You can use string_args though. - - .. seealso:: - For more information on how to use this operator, take a look at the guide: - :ref:`howto/operator:PythonVirtualenvOperator` - - :param python_callable: A python function with no references to outside variables, - defined with def, which will be run in a virtualenv - :type python_callable: function - :param multiple_outputs: if set, function return value will be - unrolled to multiple XCom values. List/Tuples will unroll to xcom values - with index as key. Dict will unroll to xcom values with keys as XCom keys. - Defaults to False. - :type multiple_outputs: bool - :param requirements: A list of requirements as specified in a pip install command - :type requirements: list[str] - :param python_version: The Python version to run the virtualenv with. Note that - both 2 and 2.7 are acceptable forms. - :type python_version: Optional[Union[str, int, float]] - :param use_dill: Whether to use dill to serialize - the args and result (pickle is default). This allow more complex types - but requires you to include dill in your requirements. - :type use_dill: bool - :param system_site_packages: Whether to include - system_site_packages in your virtualenv. - See virtualenv documentation for more information. - :type system_site_packages: bool - :param op_args: A list of positional arguments to pass to python_callable. - :type op_args: list - :param op_kwargs: A dict of keyword arguments to pass to python_callable. - :type op_kwargs: dict - :param string_args: Strings that are present in the global var virtualenv_string_args, - available to python_callable at runtime as a list[str]. Note that args are split - by newline. - :type string_args: list[str] - :param templates_dict: a dictionary where the values are templates that - will get templated by the Airflow engine sometime between - ``__init__`` and ``execute`` takes place and are made available - in your callable's context after the template has been applied - :type templates_dict: dict of str - :param templates_exts: a list of file extensions to resolve while - processing templated fields, for examples ``['.sql', '.hql']`` - :type templates_exts: list[str] - """ - return _virtualenv_task( - python_callable=python_callable, - multiple_outputs=multiple_outputs, - requirements=requirements, - python_version=python_version, - use_dill=use_dill, - system_site_packages=system_site_packages, - string_args=string_args, - templates_dict=templates_dict, - templates_exts=templates_exts, - **kwargs, - ) + def __getattr__(self, name): + if self.store.get(name, None): + return self.store[name] + raise AirflowException("Decorator %s not found", name) task = _TaskDecorator() diff --git a/airflow/provider_info.schema.json b/airflow/provider_info.schema.json index 656adbd0db16d..31b2f463b4db6 100644 --- a/airflow/provider_info.schema.json +++ b/airflow/provider_info.schema.json @@ -27,6 +27,18 @@ "items": { "type": "string" } + }, + "task-decorators": { + "type": "array", + "description": "Apply custom decorators to the TaskFlow API. Can be accessed by users via '@task.'", + "items": { + "name": { + "type": "string" + }, + "path": { + "type": "string" + } + } } }, "required": [ diff --git a/airflow/providers_manager.py b/airflow/providers_manager.py index 17bcf15ce583c..fc421d5181fe7 100644 --- a/airflow/providers_manager.py +++ b/airflow/providers_manager.py @@ -85,6 +85,13 @@ class ConnectionFormWidgetInfo(NamedTuple): field: Field +class TaskflowDecoratorInfo(NamedTuple): + """Taskflow decorator information""" + + name: str + decorator: object + + class ProvidersManager(LoggingMixin): """ Manages all provider packages. This is a Singleton class. The first time it is @@ -106,6 +113,8 @@ def __init__(self): self._provider_dict: Dict[str, ProviderInfo] = {} # Keeps dict of hooks keyed by connection type self._hooks_dict: Dict[str, HookInfo] = {} + + self._taskflow_decorator_dict: Dict[str, object] = {} # Keeps methods that should be used to add custom widgets tuple of keyed by name of the extra field self._connection_form_widgets: Dict[str, ConnectionFormWidgetInfo] = {} # Customizations for javascript fields are kept here @@ -118,6 +127,7 @@ def __init__(self): self._providers_list_initialized = False self._providers_hooks_initialized = False self._providers_extra_links_initialized = False + self._providers_taskflow_decorator_initialized = False def initialize_providers_list(self): """Lazy initialization of providers list.""" @@ -158,6 +168,20 @@ def initialize_providers_hooks(self): ) self._providers_hooks_initialized = True + def initialize_providers_taskflow_decorator(self): + """Lazy initialization of providers hooks.""" + if self._providers_taskflow_decorator_initialized: + return + self.initialize_providers_list() + start_time = perf_counter() + self.log.debug("Initializing Taskflow Decorators") + self._discover_taskflow_decorators() + self.log.debug( + "Initialization of Providers Manager Taskflow Decorators took %.2f seconds", + perf_counter() - start_time, + ) + self._providers_taskflow_decorator_initialized = True + def initialize_providers_extra_links(self): """Lazy initialization of providers extra links.""" if self._providers_extra_links_initialized: @@ -270,6 +294,13 @@ def _discover_hooks(self) -> None: for hook_class_name in hook_class_names: self._add_hook(hook_class_name, provider_package) + def _discover_taskflow_decorators(self) -> None: + for name, provider in self._provider_dict.items(): + taskflow_decorators = provider[1].get("task-decorators") + if taskflow_decorators: + for taskflow_decorator in taskflow_decorators: + self._add_taskflow_decorator(taskflow_decorator["name"], taskflow_decorator["path"], name) + @staticmethod def _get_attr(obj: Any, attr_name: str): """Retrieves attributes of an object, or warns if not found""" @@ -278,6 +309,51 @@ def _get_attr(obj: Any, attr_name: str): return None return getattr(obj, attr_name) + def _add_taskflow_decorator( + self, decorator_name, decorator_class_name: str, provider_package: str + ) -> None: + if provider_package.startswith("apache-airflow"): + provider_path = provider_package[len("apache-") :].replace("-", ".") + if not decorator_class_name.startswith(provider_path): + log.warning( + "Sanity check failed when importing '%s' from '%s' package. It should start with '%s'", + decorator_class_name, + provider_package, + provider_path, + ) + return + if decorator_name in self._taskflow_decorator_dict: + log.warning( + "The hook_class '%s' has been already registered.", + decorator_class_name, + ) + return + try: + module, class_name = decorator_class_name.rsplit('.', maxsplit=1) + decorator_class = getattr(importlib.import_module(module), class_name) + self._taskflow_decorator_dict[decorator_name] = decorator_class + # Do not use attr here. We want to check only direct class fields not those + # inherited from parent hook. This way we add form fields only once for the whole + # hierarchy and we add it only from the parent hook that provides those! + except ImportError as e: + # When there is an ImportError we turn it into debug warnings as this is + # an expected case when only some providers are installed + log.debug( + "Exception when importing '%s' from '%s' package: %s", + decorator_name, + provider_package, + e, + ) + return + except Exception as e: + log.warning( + "Exception when importing '%s' from '%s' package: %s", + decorator_name, + provider_package, + e, + ) + return + def _add_hook(self, hook_class_name: str, provider_package: str) -> None: """ Adds hook class name to list of hooks @@ -440,6 +516,11 @@ def hooks(self) -> Dict[str, HookInfo]: self.initialize_providers_hooks() return self._hooks_dict + @property + def taskflow_decorators(self) -> Dict[str, TaskflowDecoratorInfo]: + self.initialize_providers_taskflow_decorator() + return self._taskflow_decorator_dict + @property def extra_links_class_names(self) -> Set[str]: """Returns set of extra link class names.""" diff --git a/docs/apache-airflow/tutorial_taskflow_api.rst b/docs/apache-airflow/tutorial_taskflow_api.rst index 03da9e01b69d3..222de2e5baa73 100644 --- a/docs/apache-airflow/tutorial_taskflow_api.rst +++ b/docs/apache-airflow/tutorial_taskflow_api.rst @@ -181,6 +181,59 @@ and even a different python version to run your function. This option should allow for far greater flexibility for users who wish to keep their workflows more simple and pythonic. +Creating Custom TaskFlow Decorators +----------------------------------- + +As of Airflow 2.2, users can now integrate custom decorators into their provider packages and have those decorators +appear natively as part of the ``@task.____`` design. + +For an example. Let's say you were trying to create a "foo" decorator. To create ``@task.foo``, follow the following +steps: + +1. Create a ``FooDecoratedOperator`` + +In this case, we are assuming that you have a ``FooOperator`` that takes a python function as an argument. +By creating a ``FooDecoratedOperator`` that inherits from ``FooOperator`` and +``airflow.decorators.base.DecoratedOperator``, Airflow will supply much of the needed functionality required to treat +your new class as a taskflow native class. + +2. Create a ``foo_task`` function + +Once you have your decorated class, create a function that takes arguments ``python_callable``\, ``multiple_outputs``\, +and ``kwargs``\. This function will use the ``airflow.decorators.base.task_decorator_factory`` function to convert +the new ``FooDecoratedOperator`` into a TaskFlow function decorator! + +.. code-block:: python + + def foo_task( + python_callable: Optional[Callable] = None, + multiple_outputs: Optional[bool] = None, + **kwargs + ): + return task_decorator_factory( + python_callable=python_callable, + multiple_outputs=multiple_outputs, + decorated_operator_class=FooDecoratedOperator, + **kwargs, + ) + +3. Register your new decorator in the provider.yaml of your provider + +Finally, add a key-value of ``decorator-name``:``path-to-function`` to your provider.yaml. When Airflow starts, the +``ProviderManager`` class will automatically import this value and ``task.decorator-name`` will work as a new +decorator! + +.. code-block:: yaml + + package-name: apache-airflow-providers-docker + name: Docker + description: | + `Docker `__ + + task-decorators: + docker: airflow.providers.docker.operators.docker.docker_decorator + + Multiple outputs inference -------------------------- Tasks can also infer multiple outputs by using dict python typing. diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 47c82baf461bb..59715ca74cfd6 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -139,6 +139,7 @@ Firehose Firestore Flink FluentD +FooDecoratedOperator Formaturas Fundera GCS