Allow Asset decorator to work with any TaskFlow operator

Author: yoeo

airflow-core/docs/authoring-and-scheduling/assets.rst

@@ -229,6 +229,15 @@ The other way around also applies:
     def process_example_asset(example_asset):
         """Process inlet example_asset..."""
 
+In addition, ``@asset`` can be used with ``@task`` to set initial arguments for the task or to use an operator other than ``PythonOperator``:
+
+.. code-block:: python
+
+    @asset(schedule=None)
+    @task.bash(retries=3)
+    def example_asset():
+        """Write to example_asset, from a Bash task with 3 retries..."""
+        return "echo 'run'"
 
 Output to multiple assets in one task
 -------------------------------------

task-sdk/src/airflow/sdk/definitions/asset/decorators.py

@@ -23,6 +23,7 @@
 import attrs
 
 from airflow.providers.standard.operators.python import PythonOperator
+from airflow.sdk.bases.decorator import _TaskDecorator
 from airflow.sdk.definitions.asset import Asset, AssetRef, BaseAsset
 from airflow.sdk.exceptions import AirflowRuntimeError
 
@@ -111,7 +112,12 @@ class AssetDefinition(Asset):
 
     def __attrs_post_init__(self) -> None:
         with self._source.create_dag(default_dag_id=self.name):
-            _AssetMainOperator.from_definition(self)
+            if isinstance(self._function, _TaskDecorator):
+                if "outlets" not in self._function.kwargs:
+                    self._function.kwargs["outlets"] = [v for _, v in self.iter_assets()]
+                self._function()
+            else:
+                _AssetMainOperator.from_definition(self)
 
 
 @attrs.define(kw_only=True)
@@ -131,7 +137,12 @@ class MultiAssetDefinition(BaseAsset):
 
     def __attrs_post_init__(self) -> None:
         with self._source.create_dag(default_dag_id=self._function.__name__):
-            _AssetMainOperator.from_definition(self)
+            if isinstance(self._function, _TaskDecorator):
+                if "outlets" not in self._function.kwargs:
+                    self._function.kwargs["outlets"] = [v for _, v in self.iter_assets()]
+                self._function()
+            else:
+                _AssetMainOperator.from_definition(self)
 
     def iter_assets(self) -> Iterator[tuple[AssetUniqueKey, Asset]]:
         for o in self._source.outlets:

task-sdk/tests/task_sdk/definitions/test_asset_decorators.py

@@ -22,6 +22,7 @@
 
 from airflow.sdk.definitions.asset import Asset
 from airflow.sdk.definitions.asset.decorators import _AssetMainOperator, asset
+from airflow.sdk.definitions.decorators import task
 from airflow.sdk.execution_time.comms import AssetResult, GetAssetByName
 
 
@@ -159,6 +160,17 @@ def example_asset_func(self, /):
             == "positional-only argument 'self' without a default is not supported in @asset"
         )
 
+    def test_with_task_decorator(self, func_fixer):
+        @task(retries=3)
+        @func_fixer
+        def _example_task_func():
+            return "This is example_task"
+
+        asset_definition = asset(name="asset", dag_id="dag", schedule=None)(_example_task_func)
+        assert asset_definition.name == "asset"
+        assert asset_definition._source.dag_id == "dag"
+        assert asset_definition._function == _example_task_func
+
     @pytest.mark.parametrize(
         "provided_uri, expected_uri",
         [
@@ -222,6 +234,36 @@ def test__attrs_post_init__(self, DAG, from_definition, example_asset_func_with_
         )
         from_definition.assert_called_once_with(asset_definition)
 
+    @mock.patch("airflow.sdk.bases.decorator._TaskDecorator.__call__")
+    @mock.patch("airflow.sdk.definitions.dag.DAG")
+    def test_with_task_decorator(self, DAG, __call__, func_fixer):
+        @task(retries=3)
+        @func_fixer
+        def _example_task_func():
+            return "This is example_task"
+
+        asset_definition = asset(schedule=None, uri="s3://bucket/object", group="MLModel", extra={"k": "v"})(
+            _example_task_func
+        )
+
+        DAG.assert_called_once_with(
+            dag_id="example_asset_func",
+            dag_display_name="example_asset_func",
+            description=None,
+            schedule=None,
+            catchup=False,
+            is_paused_upon_creation=None,
+            on_failure_callback=None,
+            on_success_callback=None,
+            params=None,
+            access_control=None,
+            owner_links={},
+            tags=set(),
+            auto_register=True,
+        )
+        __call__.assert_called_once_with()
+        assert asset_definition._function.kwargs["outlets"] == [asset_definition]
+
 
 class TestMultiAssetDefinition:
     @mock.patch("airflow.sdk.definitions.asset.decorators._AssetMainOperator.from_definition")
@@ -249,6 +291,37 @@ def test__attrs_post_init__(self, DAG, from_definition, example_asset_func_with_
         )
         from_definition.assert_called_once_with(definition)
 
+    @mock.patch("airflow.sdk.bases.decorator._TaskDecorator.__call__")
+    @mock.patch("airflow.sdk.definitions.dag.DAG")
+    def test_with_task_decorator(self, DAG, __call__, func_fixer):
+        @task(retries=3)
+        @func_fixer
+        def _example_task_func():
+            return "This is example_task"
+
+        definition = asset.multi(
+            schedule=None,
+            outlets=[Asset(name="a"), Asset(name="b")],
+        )(_example_task_func)
+
+        DAG.assert_called_once_with(
+            dag_id="example_asset_func",
+            dag_display_name="example_asset_func",
+            description=None,
+            schedule=None,
+            catchup=False,
+            is_paused_upon_creation=None,
+            on_failure_callback=None,
+            on_success_callback=None,
+            params=None,
+            access_control=None,
+            owner_links={},
+            tags=set(),
+            auto_register=True,
+        )
+        __call__.assert_called_once_with()
+        assert definition._function.kwargs["outlets"] == [Asset(name="a"), Asset(name="b")]
+
 
 class Test_AssetMainOperator:
     def test_from_definition(self, example_asset_func_with_valid_arg_as_inlet_asset):