Make decorators pluggable

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()

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.<name>'",
+        "items": {
+            "name": {
+                "type": "string"
+            },
+            "path": {
+                "type": "string"
+            }
+        }
     }
   },
   "required": [

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."""

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 <https://docs.docker.com/install/>`__
+
+   task-decorators:
+       docker: airflow.providers.docker.operators.docker.docker_decorator
+
+
 Multiple outputs inference
 --------------------------
 Tasks can also infer multiple outputs by using dict python typing.

docs/spelling_wordlist.txt

@@ -139,6 +139,7 @@ Firehose
 Firestore
 Flink
 FluentD
+FooDecoratedOperator
 Formaturas
 Fundera
 GCS