From fb495df94ca788df5f7b8236b415fc6eecabdb15 Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Tue, 27 Jul 2021 11:09:47 -0700 Subject: [PATCH 01/13] Make decorators pluggable This PR will allow users to add custom "@task.____" decorators by adding to their setup.cfg files. This will make decorators seem native to airflow while living in provider packages. --- airflow/decorator_implementations/__init__.py | 16 +++ .../python.py | 0 .../python_virtualenv.py | 0 .../task_group.py | 0 airflow/decorators/__init__.py | 114 ++---------------- setup.cfg | 5 + 6 files changed, 34 insertions(+), 101 deletions(-) create mode 100644 airflow/decorator_implementations/__init__.py rename airflow/{decorators => decorator_implementations}/python.py (100%) rename airflow/{decorators => decorator_implementations}/python_virtualenv.py (100%) rename airflow/{decorators => decorator_implementations}/task_group.py (100%) diff --git a/airflow/decorator_implementations/__init__.py b/airflow/decorator_implementations/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow/decorator_implementations/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow/decorators/python.py b/airflow/decorator_implementations/python.py similarity index 100% rename from airflow/decorators/python.py rename to airflow/decorator_implementations/python.py diff --git a/airflow/decorators/python_virtualenv.py b/airflow/decorator_implementations/python_virtualenv.py similarity index 100% rename from airflow/decorators/python_virtualenv.py rename to airflow/decorator_implementations/python_virtualenv.py diff --git a/airflow/decorators/task_group.py b/airflow/decorator_implementations/task_group.py similarity index 100% rename from airflow/decorators/task_group.py rename to airflow/decorator_implementations/task_group.py diff --git a/airflow/decorators/__init__.py b/airflow/decorators/__init__.py index 1250f321de291..07f6e445d63b3 100644 --- a/airflow/decorators/__init__.py +++ b/airflow/decorators/__init__.py @@ -15,15 +15,17 @@ # specific language governing permissions and limitations # under the License. -from typing import Callable, Dict, Iterable, List, Optional, Union +from typing import Callable, Optional + +import importlib_metadata as metadata -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.models.dag import dag # noqa class _TaskDecorator: + def __init__(self): + self.store = {} + def __call__( self, python_callable: Optional[Callable] = None, multiple_outputs: Optional[bool] = None, **kwargs ): @@ -41,103 +43,13 @@ 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] + connections = [e for e in metadata.entry_points()['task_decorator_connections'] if e.name == name] + mod = connections[0].load() + self.store[name] = mod + return mod task = _TaskDecorator() diff --git a/setup.cfg b/setup.cfg index d3c5f574c0b7d..a173e363600e0 100644 --- a/setup.cfg +++ b/setup.cfg @@ -180,6 +180,11 @@ airflow.serialization=*.json console_scripts= airflow=airflow.__main__:main +task_decorator_connections= + python=airflow.decorator_implementations.python:python_task + virtualenv=airflow.decorator_implementations.python_virtualenv:_virtualenv_task + task_group=airflow.decorator_implementations.task_group:task_group + [bdist_wheel] python-tag=py3 From 92bf68bc6c1024da1f9919b3454aebc6add0398a Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Tue, 27 Jul 2021 18:18:16 -0700 Subject: [PATCH 02/13] change entrypoint name --- airflow/decorators/__init__.py | 2 +- setup.cfg | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/decorators/__init__.py b/airflow/decorators/__init__.py index 07f6e445d63b3..c0073defa5088 100644 --- a/airflow/decorators/__init__.py +++ b/airflow/decorators/__init__.py @@ -46,7 +46,7 @@ def __call__( def __getattr__(self, name): if self.store.get(name, None): return self.store[name] - connections = [e for e in metadata.entry_points()['task_decorator_connections'] if e.name == name] + connections = [e for e in metadata.entry_points()['airflow.task_decorators'] if e.name == name] mod = connections[0].load() self.store[name] = mod return mod diff --git a/setup.cfg b/setup.cfg index a173e363600e0..6267747267e30 100644 --- a/setup.cfg +++ b/setup.cfg @@ -180,7 +180,7 @@ airflow.serialization=*.json console_scripts= airflow=airflow.__main__:main -task_decorator_connections= +airflow.task_decorators= python=airflow.decorator_implementations.python:python_task virtualenv=airflow.decorator_implementations.python_virtualenv:_virtualenv_task task_group=airflow.decorator_implementations.task_group:task_group From 8861fe5dcd7087df44634a8914d8c69253f95e29 Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Wed, 28 Jul 2021 10:05:18 -0700 Subject: [PATCH 03/13] store python and virtualenv as defaults --- airflow/decorator_implementations/__init__.py | 16 ---------------- airflow/decorators/__init__.py | 4 +++- .../python.py | 0 .../python_virtualenv.py | 0 .../task_group.py | 0 setup.cfg | 5 ----- 6 files changed, 3 insertions(+), 22 deletions(-) delete mode 100644 airflow/decorator_implementations/__init__.py rename airflow/{decorator_implementations => decorators}/python.py (100%) rename airflow/{decorator_implementations => decorators}/python_virtualenv.py (100%) rename airflow/{decorator_implementations => decorators}/task_group.py (100%) diff --git a/airflow/decorator_implementations/__init__.py b/airflow/decorator_implementations/__init__.py deleted file mode 100644 index 13a83393a9124..0000000000000 --- a/airflow/decorator_implementations/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. diff --git a/airflow/decorators/__init__.py b/airflow/decorators/__init__.py index c0073defa5088..72ed02d968a18 100644 --- a/airflow/decorators/__init__.py +++ b/airflow/decorators/__init__.py @@ -19,12 +19,14 @@ import importlib_metadata as metadata +from airflow.decorators.python import python_task +from airflow.decorators.python_virtualenv import _virtualenv_task from airflow.models.dag import dag # noqa class _TaskDecorator: def __init__(self): - self.store = {} + self.store = {"python": python_task, "virtualenv": _virtualenv_task} def __call__( self, python_callable: Optional[Callable] = None, multiple_outputs: Optional[bool] = None, **kwargs diff --git a/airflow/decorator_implementations/python.py b/airflow/decorators/python.py similarity index 100% rename from airflow/decorator_implementations/python.py rename to airflow/decorators/python.py diff --git a/airflow/decorator_implementations/python_virtualenv.py b/airflow/decorators/python_virtualenv.py similarity index 100% rename from airflow/decorator_implementations/python_virtualenv.py rename to airflow/decorators/python_virtualenv.py diff --git a/airflow/decorator_implementations/task_group.py b/airflow/decorators/task_group.py similarity index 100% rename from airflow/decorator_implementations/task_group.py rename to airflow/decorators/task_group.py diff --git a/setup.cfg b/setup.cfg index 6267747267e30..d3c5f574c0b7d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -180,11 +180,6 @@ airflow.serialization=*.json console_scripts= airflow=airflow.__main__:main -airflow.task_decorators= - python=airflow.decorator_implementations.python:python_task - virtualenv=airflow.decorator_implementations.python_virtualenv:_virtualenv_task - task_group=airflow.decorator_implementations.task_group:task_group - [bdist_wheel] python-tag=py3 From cd4a96f92038c87f6ffa5486bca198caf0061bad Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Wed, 28 Jul 2021 10:12:10 -0700 Subject: [PATCH 04/13] store python and virtualenv as defaults --- airflow/decorators/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow/decorators/__init__.py b/airflow/decorators/__init__.py index 72ed02d968a18..19c43a1ba1c60 100644 --- a/airflow/decorators/__init__.py +++ b/airflow/decorators/__init__.py @@ -21,6 +21,7 @@ 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.models.dag import dag # noqa From 3121eefe761639a578b921d04112d2b5830a88c6 Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Wed, 28 Jul 2021 10:46:52 -0700 Subject: [PATCH 05/13] add docs --- docs/apache-airflow/tutorial_taskflow_api.rst | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/docs/apache-airflow/tutorial_taskflow_api.rst b/docs/apache-airflow/tutorial_taskflow_api.rst index 03da9e01b69d3..0b006b6cfdb43 100644 --- a/docs/apache-airflow/tutorial_taskflow_api.rst +++ b/docs/apache-airflow/tutorial_taskflow_api.rst @@ -181,6 +181,54 @@ 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.3, 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 ``DecoratedFooOperator`` + +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 entrypoints of your setup.cfg + +Finally, use the ``options.entrypoints`` section of your ``setup.cfg`` to allow Airflow to detect your new decorator. +To allow airflow to see this endpoint, give the function URL =: under the +``airflow.task_decorators`` tag. + +.. code-block:: bash + [options.entry_points] + airflow.task_decorators= + foo=my.class.path.foo:foo_task + +Now when you install your provider, you should be able to use the ``@task.foo`` decorator and turn any python operator +into a foo task! + Multiple outputs inference -------------------------- Tasks can also infer multiple outputs by using dict python typing. From 9ddf1c9983419467c5eb9548383fd3ab675af2dc Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Mon, 9 Aug 2021 09:57:02 -0700 Subject: [PATCH 06/13] Set providers to put taskflow decorators into the provider_info.yaml --- airflow/decorators/__init__.py | 5 ++--- airflow/provider_info.schema.json | 7 ++++++ airflow/providers_manager.py | 37 +++++++++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 3 deletions(-) diff --git a/airflow/decorators/__init__.py b/airflow/decorators/__init__.py index 19c43a1ba1c60..6c15989d54e39 100644 --- a/airflow/decorators/__init__.py +++ b/airflow/decorators/__init__.py @@ -17,12 +17,11 @@ from typing import Callable, Optional -import importlib_metadata as metadata - 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.models.dag import dag # noqa +from airflow.providers_manager import ProvidersManager class _TaskDecorator: @@ -49,7 +48,7 @@ def __call__( def __getattr__(self, name): if self.store.get(name, None): return self.store[name] - connections = [e for e in metadata.entry_points()['airflow.task_decorators'] if e.name == name] + connections = ProvidersManager.taskflow_decorators[name].name mod = connections[0].load() self.store[name] = mod return mod diff --git a/airflow/provider_info.schema.json b/airflow/provider_info.schema.json index 656adbd0db16d..16bf9caa92d9a 100644 --- a/airflow/provider_info.schema.json +++ b/airflow/provider_info.schema.json @@ -27,6 +27,13 @@ "items": { "type": "string" } + }, + "task-decorators": { + "type": "array", + "description": "Apply custom decorators to the TaskFlow API. Can be accessed by users via '@task.'", + "items": { + "type": "string" + } } }, "required": [ diff --git a/airflow/providers_manager.py b/airflow/providers_manager.py index 17bcf15ce583c..e9a2b295d3a60 100644 --- a/airflow/providers_manager.py +++ b/airflow/providers_manager.py @@ -85,6 +85,12 @@ class ConnectionFormWidgetInfo(NamedTuple): field: Field +class TaskflowDecoratorInfo(NamedTuple): + """Taskflow decorator information""" + + name: str + + class ProvidersManager(LoggingMixin): """ Manages all provider packages. This is a Singleton class. The first time it is @@ -106,6 +112,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, TaskflowDecoratorInfo] = {} # 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 +126,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 +167,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 +293,15 @@ 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._taskflow_decorator_dict[taskflow_decorators] = TaskflowDecoratorInfo( + name=taskflow_decorator + ) + @staticmethod def _get_attr(obj: Any, attr_name: str): """Retrieves attributes of an object, or warns if not found""" @@ -440,6 +472,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.""" From 38dc4a4a48373d9c3a348f8669b87abba89c0241 Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Mon, 9 Aug 2021 13:23:33 -0700 Subject: [PATCH 07/13] make doc tests pass --- docs/apache-airflow/tutorial_taskflow_api.rst | 41 +++++++++++-------- docs/spelling_wordlist.txt | 1 + 2 files changed, 24 insertions(+), 18 deletions(-) diff --git a/docs/apache-airflow/tutorial_taskflow_api.rst b/docs/apache-airflow/tutorial_taskflow_api.rst index 0b006b6cfdb43..adc6af718ddf5 100644 --- a/docs/apache-airflow/tutorial_taskflow_api.rst +++ b/docs/apache-airflow/tutorial_taskflow_api.rst @@ -199,32 +199,37 @@ 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 +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 entrypoints of your setup.cfg + + 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-info.yaml of your provider + + Finally, use the ``options.entrypoints`` section of your ``setup.cfg`` to allow Airflow to detect your new decorator. To allow airflow to see this endpoint, give the function URL =: under the ``airflow.task_decorators`` tag. .. code-block:: bash - [options.entry_points] - airflow.task_decorators= - foo=my.class.path.foo:foo_task + + [options.entry_points] + airflow.task_decorators= + foo=my.class.path.foo:foo_task Now when you install your provider, you should be able to use the ``@task.foo`` decorator and turn any python operator into a foo task! 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 From 22d95dbda3327d68f37c5d88621356cee26915a9 Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Mon, 9 Aug 2021 13:24:33 -0700 Subject: [PATCH 08/13] nit --- docs/apache-airflow/tutorial_taskflow_api.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/apache-airflow/tutorial_taskflow_api.rst b/docs/apache-airflow/tutorial_taskflow_api.rst index adc6af718ddf5..124e3dfb56c20 100644 --- a/docs/apache-airflow/tutorial_taskflow_api.rst +++ b/docs/apache-airflow/tutorial_taskflow_api.rst @@ -190,7 +190,7 @@ 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 ``DecoratedFooOperator`` +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 From b9b9bd8b5ec3900b2cb8ed6cb4653c7d0e26573b Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Mon, 9 Aug 2021 13:38:31 -0700 Subject: [PATCH 09/13] final doc fix --- docs/apache-airflow/tutorial_taskflow_api.rst | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/docs/apache-airflow/tutorial_taskflow_api.rst b/docs/apache-airflow/tutorial_taskflow_api.rst index 124e3dfb56c20..db9174de1ec13 100644 --- a/docs/apache-airflow/tutorial_taskflow_api.rst +++ b/docs/apache-airflow/tutorial_taskflow_api.rst @@ -217,8 +217,21 @@ the new ``FooDecoratedOperator`` into a TaskFlow function decorator! **kwargs, ) -3. Register your new decorator in the provider-info.yaml of your provider +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 Finally, use the ``options.entrypoints`` section of your ``setup.cfg`` to allow Airflow to detect your new decorator. From 1ddbf367aebd22717930bf40cc641c114c5d2fb3 Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Mon, 9 Aug 2021 13:40:23 -0700 Subject: [PATCH 10/13] final doc fix --- docs/apache-airflow/tutorial_taskflow_api.rst | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/docs/apache-airflow/tutorial_taskflow_api.rst b/docs/apache-airflow/tutorial_taskflow_api.rst index db9174de1ec13..31ed8eb524b02 100644 --- a/docs/apache-airflow/tutorial_taskflow_api.rst +++ b/docs/apache-airflow/tutorial_taskflow_api.rst @@ -234,19 +234,6 @@ decorator! docker: airflow.providers.docker.operators.docker.docker_decorator -Finally, use the ``options.entrypoints`` section of your ``setup.cfg`` to allow Airflow to detect your new decorator. -To allow airflow to see this endpoint, give the function URL =: under the -``airflow.task_decorators`` tag. - -.. code-block:: bash - - [options.entry_points] - airflow.task_decorators= - foo=my.class.path.foo:foo_task - -Now when you install your provider, you should be able to use the ``@task.foo`` decorator and turn any python operator -into a foo task! - Multiple outputs inference -------------------------- Tasks can also infer multiple outputs by using dict python typing. From 2bc3780f958497b3785972e0871080c4644c572a Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Mon, 9 Aug 2021 19:35:21 -0700 Subject: [PATCH 11/13] fix based on testing with DockerDecorator --- airflow/decorators/__init__.py | 7 ++--- airflow/provider_info.schema.json | 7 ++++- airflow/providers_manager.py | 52 ++++++++++++++++++++++++++++--- 3 files changed, 57 insertions(+), 9 deletions(-) diff --git a/airflow/decorators/__init__.py b/airflow/decorators/__init__.py index 6c15989d54e39..859febc843215 100644 --- a/airflow/decorators/__init__.py +++ b/airflow/decorators/__init__.py @@ -48,10 +48,9 @@ def __call__( def __getattr__(self, name): if self.store.get(name, None): return self.store[name] - connections = ProvidersManager.taskflow_decorators[name].name - mod = connections[0].load() - self.store[name] = mod - return mod + decorator = ProvidersManager().taskflow_decorators[name] + self.store[name] = decorator + return decorator task = _TaskDecorator() diff --git a/airflow/provider_info.schema.json b/airflow/provider_info.schema.json index 16bf9caa92d9a..31b2f463b4db6 100644 --- a/airflow/provider_info.schema.json +++ b/airflow/provider_info.schema.json @@ -32,7 +32,12 @@ "type": "array", "description": "Apply custom decorators to the TaskFlow API. Can be accessed by users via '@task.'", "items": { - "type": "string" + "name": { + "type": "string" + }, + "path": { + "type": "string" + } } } }, diff --git a/airflow/providers_manager.py b/airflow/providers_manager.py index e9a2b295d3a60..fc421d5181fe7 100644 --- a/airflow/providers_manager.py +++ b/airflow/providers_manager.py @@ -89,6 +89,7 @@ class TaskflowDecoratorInfo(NamedTuple): """Taskflow decorator information""" name: str + decorator: object class ProvidersManager(LoggingMixin): @@ -113,7 +114,7 @@ def __init__(self): # Keeps dict of hooks keyed by connection type self._hooks_dict: Dict[str, HookInfo] = {} - self._taskflow_decorator_dict: Dict[str, TaskflowDecoratorInfo] = {} + 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 @@ -298,9 +299,7 @@ def _discover_taskflow_decorators(self) -> None: taskflow_decorators = provider[1].get("task-decorators") if taskflow_decorators: for taskflow_decorator in taskflow_decorators: - self._taskflow_decorator_dict[taskflow_decorators] = TaskflowDecoratorInfo( - name=taskflow_decorator - ) + self._add_taskflow_decorator(taskflow_decorator["name"], taskflow_decorator["path"], name) @staticmethod def _get_attr(obj: Any, attr_name: str): @@ -310,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 From 7c4851aeca12fe48f08837f70167ff8371087cf6 Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Tue, 10 Aug 2021 08:39:24 -0500 Subject: [PATCH 12/13] Update docs/apache-airflow/tutorial_taskflow_api.rst Co-authored-by: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com> --- docs/apache-airflow/tutorial_taskflow_api.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/apache-airflow/tutorial_taskflow_api.rst b/docs/apache-airflow/tutorial_taskflow_api.rst index 31ed8eb524b02..222de2e5baa73 100644 --- a/docs/apache-airflow/tutorial_taskflow_api.rst +++ b/docs/apache-airflow/tutorial_taskflow_api.rst @@ -184,7 +184,7 @@ and pythonic. Creating Custom TaskFlow Decorators ----------------------------------- -As of Airflow 2.3, users can now integrate custom decorators into their provider packages and have those 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 From 85611d98b0f2b67ceaca01f43d1e0887fcc5edb3 Mon Sep 17 00:00:00 2001 From: Daniel Imberman Date: Tue, 10 Aug 2021 07:40:52 -0700 Subject: [PATCH 13/13] cleaner import system --- airflow/decorators/__init__.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/airflow/decorators/__init__.py b/airflow/decorators/__init__.py index 859febc843215..53691725a2e4c 100644 --- a/airflow/decorators/__init__.py +++ b/airflow/decorators/__init__.py @@ -20,6 +20,7 @@ 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 @@ -27,6 +28,9 @@ 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 @@ -48,9 +52,7 @@ def __call__( def __getattr__(self, name): if self.store.get(name, None): return self.store[name] - decorator = ProvidersManager().taskflow_decorators[name] - self.store[name] = decorator - return decorator + raise AirflowException("Decorator %s not found", name) task = _TaskDecorator()