From b707d72f92420b46a6f52d01e36ef10ebc8d091a Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 8 Aug 2024 16:41:11 +0800 Subject: [PATCH 01/29] Implement asset definition creating a DAG --- airflow/decorators/assets.py | 74 +++++++++++++++++++++++++++++++ airflow/serialization/schema.json | 10 ++++- airflow/utils/file.py | 4 +- 3 files changed, 85 insertions(+), 3 deletions(-) create mode 100644 airflow/decorators/assets.py diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py new file mode 100644 index 0000000000000..66d021218c12e --- /dev/null +++ b/airflow/decorators/assets.py @@ -0,0 +1,74 @@ +# 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. + +from __future__ import annotations + +import types +import typing + +import attrs + +from airflow.datasets import Dataset as Asset +from airflow.models.dag import DAG, ScheduleArg +from airflow.operators.python import PythonOperator + +if typing.TYPE_CHECKING: + from airflow.io.path import ObjectStoragePath + + +@attrs.define(kw_only=True) +class AssetDefinition: + """ + Asset representation from decorating a function with ``@asset``. + + :meta private: + """ + + name: str # TODO: This should be stored on Asset. + asset: Asset + function: types.FunctionType + schedule: ScheduleArg + + def __attrs_post_init__(self) -> None: + with DAG(dag_id=self.name, schedule=self.schedule, auto_register=True) as dag: + PythonOperator(task_id="__main__", outlets=[self.asset], python_callable=self.function) + # TODO: Currently this just gets serialized into a string. + # When we create UI for assets, we should add logic to serde so the + # serialized DAG contains appropriate asset information. + dag._wrapped_definition = self + + +@attrs.define(kw_only=True) +class asset: + """Create an asset by decorating a materialization function.""" + + schedule: ScheduleArg + uri: str | ObjectStoragePath | None + extra: dict[str, typing.Any] = attrs.field(factory=dict) + + def __call__(self, f: types.FunctionType) -> AssetDefinition: + if (name := f.__name__) != f.__qualname__: + raise ValueError("nested function not supported") + return AssetDefinition( + name=name, + asset=Asset( + uri=name if self.uri is None else str(self.uri), + extra=self.extra, + ), + function=f, + schedule=self.schedule, + ) diff --git a/airflow/serialization/schema.json b/airflow/serialization/schema.json index b26b59339816c..8af11d2e8342f 100644 --- a/airflow/serialization/schema.json +++ b/airflow/serialization/schema.json @@ -153,7 +153,7 @@ "_processor_dags_folder": { "anyOf": [ { "type": "null" }, - {"type": "string"} + { "type": "string" } ] }, "dag_display_name": { "type" : "string"}, @@ -179,7 +179,13 @@ { "$ref": "#/definitions/task_group" } ]}, "edge_info": { "$ref": "#/definitions/edge_info" }, - "dag_dependencies": { "$ref": "#/definitions/dag_dependencies" } + "dag_dependencies": { "$ref": "#/definitions/dag_dependencies" }, + "_wrapped_definition": { + "anyOf": [ + { "type": "null" }, + { "type": "string" } + ] + } }, "required": [ "dag_id", diff --git a/airflow/utils/file.py b/airflow/utils/file.py index 86b7a7891ca8f..962f97c8fcfa7 100644 --- a/airflow/utils/file.py +++ b/airflow/utils/file.py @@ -328,7 +328,9 @@ def might_contain_dag_via_default_heuristic(file_path: str, zip_file: zipfile.Zi with open(file_path, "rb") as dag_file: content = dag_file.read() content = content.lower() - return all(s in content for s in (b"dag", b"airflow")) + if b"airflow" not in content: + return False + return any(s in content for s in (b"dag", b"asset")) def _find_imported_modules(module: ast.Module) -> Generator[str, None, None]: From f26e64fce8326636cecb95c73f133beda8b01c25 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Fri, 9 Aug 2024 15:47:57 +0800 Subject: [PATCH 02/29] Basic inlet dependency --- airflow/decorators/assets.py | 50 +++++++++++++++++++++++++++++++++--- 1 file changed, 46 insertions(+), 4 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 66d021218c12e..3b7d357bd5e83 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -17,8 +17,9 @@ from __future__ import annotations +import inspect import types -import typing +from typing import TYPE_CHECKING, Any, Iterator, Mapping import attrs @@ -26,10 +27,41 @@ from airflow.models.dag import DAG, ScheduleArg from airflow.operators.python import PythonOperator -if typing.TYPE_CHECKING: +if TYPE_CHECKING: from airflow.io.path import ObjectStoragePath +@attrs.define(kw_only=True) +class AssetRef: + """Reference to an asset.""" + + name: str + + +class _AssetMainOperator(PythonOperator): + def __init__(self, *, definition_name: str, **kwargs) -> None: + super().__init__(**kwargs) + self._definition_name = definition_name + + def _iter_kwargs(self, context: Mapping[str, Any]) -> Iterator[tuple[str, Any]]: + for key in inspect.signature(self.python_callable).parameters: + if key == "self": + value: Any = AssetRef(name=self._definition_name) + elif key == "context": + value = context + else: + # TODO: This does not check if the upstream asset actually + # exists. Should we do a second pass in the DAG processor to + # raise parse-time errors if a non-existent asset is referenced? + # How? Should we also fail the task at runtime? Or should the + # dangling reference simply do nothing? + value = AssetRef(name=key) + yield key, value + + def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]: + return dict(self._iter_kwargs(context)) + + @attrs.define(kw_only=True) class AssetDefinition: """ @@ -44,8 +76,16 @@ class AssetDefinition: schedule: ScheduleArg def __attrs_post_init__(self) -> None: + parameters = inspect.signature(self.function).parameters with DAG(dag_id=self.name, schedule=self.schedule, auto_register=True) as dag: - PythonOperator(task_id="__main__", outlets=[self.asset], python_callable=self.function) + _AssetMainOperator( + task_id="__main__", + # TODO: This should use the name argument instead. + inlets=[Asset(uri=k) for k in parameters if k not in ("self", "context")], + outlets=[self.asset], + python_callable=self.function, + definition_name=self.name, + ) # TODO: Currently this just gets serialized into a string. # When we create UI for assets, we should add logic to serde so the # serialized DAG contains appropriate asset information. @@ -58,11 +98,13 @@ class asset: schedule: ScheduleArg uri: str | ObjectStoragePath | None - extra: dict[str, typing.Any] = attrs.field(factory=dict) + extra: dict[str, Any] = attrs.field(factory=dict) def __call__(self, f: types.FunctionType) -> AssetDefinition: if (name := f.__name__) != f.__qualname__: raise ValueError("nested function not supported") + if name == "self" or name == "context": + raise ValueError(f"prohibited name for asset: {name}") return AssetDefinition( name=name, asset=Asset( From 80b98734be8f7da40114973bbbc0c03838b54165 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Fri, 9 Aug 2024 16:29:19 +0800 Subject: [PATCH 03/29] Make AssetDefinition subclass Asset This seems to be the best way for 'schedule' dependencies to work. Still not entirely sure; we'll revisit this. --- airflow/decorators/assets.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 3b7d357bd5e83..5f87d489525a4 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -63,7 +63,7 @@ def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]: @attrs.define(kw_only=True) -class AssetDefinition: +class AssetDefinition(Asset): """ Asset representation from decorating a function with ``@asset``. @@ -71,7 +71,6 @@ class AssetDefinition: """ name: str # TODO: This should be stored on Asset. - asset: Asset function: types.FunctionType schedule: ScheduleArg @@ -82,7 +81,7 @@ def __attrs_post_init__(self) -> None: task_id="__main__", # TODO: This should use the name argument instead. inlets=[Asset(uri=k) for k in parameters if k not in ("self", "context")], - outlets=[self.asset], + outlets=[self], python_callable=self.function, definition_name=self.name, ) @@ -107,10 +106,8 @@ def __call__(self, f: types.FunctionType) -> AssetDefinition: raise ValueError(f"prohibited name for asset: {name}") return AssetDefinition( name=name, - asset=Asset( - uri=name if self.uri is None else str(self.uri), - extra=self.extra, - ), + uri=name if self.uri is None else str(self.uri), + extra=self.extra, function=f, schedule=self.schedule, ) From 6d33cc12c5445355b4f04367a6d671ee3c5cf492 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 25 Oct 2024 15:53:22 +0800 Subject: [PATCH 04/29] style: fix mypy error --- airflow/decorators/assets.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 5f87d489525a4..fba6eb8dfb96e 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -23,7 +23,7 @@ import attrs -from airflow.datasets import Dataset as Asset +from airflow.assets import Asset, _validate_identifier from airflow.models.dag import DAG, ScheduleArg from airflow.operators.python import PythonOperator @@ -97,6 +97,11 @@ class asset: schedule: ScheduleArg uri: str | ObjectStoragePath | None + group: str = attrs.field( + kw_only=True, + default="", + validator=[attrs.validators.max_len(1500), _validate_identifier], + ) extra: dict[str, Any] = attrs.field(factory=dict) def __call__(self, f: types.FunctionType) -> AssetDefinition: @@ -107,6 +112,7 @@ def __call__(self, f: types.FunctionType) -> AssetDefinition: return AssetDefinition( name=name, uri=name if self.uri is None else str(self.uri), + group=self.group, extra=self.extra, function=f, schedule=self.schedule, From c8eba161c286248ab95c0a296df67c026331e902 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Mon, 28 Oct 2024 18:29:07 +0800 Subject: [PATCH 05/29] feat(asset): allow uri to be None --- airflow/decorators/assets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index fba6eb8dfb96e..7b317d54533a5 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -96,7 +96,7 @@ class asset: """Create an asset by decorating a materialization function.""" schedule: ScheduleArg - uri: str | ObjectStoragePath | None + uri: str | ObjectStoragePath | None = None group: str = attrs.field( kw_only=True, default="", From c5d1ddccf330f03e73e13ba8b8e0d21b7d454be5 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Wed, 30 Oct 2024 16:31:41 +0800 Subject: [PATCH 06/29] fix: temporarily serialize AssetDefintion into a string --- airflow/decorators/assets.py | 11 ++++++++--- airflow/serialization/schema.json | 3 ++- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 7b317d54533a5..f813d13d320db 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -70,7 +70,6 @@ class AssetDefinition(Asset): :meta private: """ - name: str # TODO: This should be stored on Asset. function: types.FunctionType schedule: ScheduleArg @@ -79,8 +78,11 @@ def __attrs_post_init__(self) -> None: with DAG(dag_id=self.name, schedule=self.schedule, auto_register=True) as dag: _AssetMainOperator( task_id="__main__", - # TODO: This should use the name argument instead. - inlets=[Asset(uri=k) for k in parameters if k not in ("self", "context")], + inlets=[ + Asset(name=inlet_aset_name) + for inlet_aset_name in parameters + if inlet_aset_name not in ("self", "context") + ], outlets=[self], python_callable=self.function, definition_name=self.name, @@ -90,6 +92,9 @@ def __attrs_post_init__(self) -> None: # serialized DAG contains appropriate asset information. dag._wrapped_definition = self + def serialize(self): + return "AssetDefinition" + @attrs.define(kw_only=True) class asset: diff --git a/airflow/serialization/schema.json b/airflow/serialization/schema.json index 8af11d2e8342f..6f4386821572f 100644 --- a/airflow/serialization/schema.json +++ b/airflow/serialization/schema.json @@ -183,7 +183,8 @@ "_wrapped_definition": { "anyOf": [ { "type": "null" }, - { "type": "string" } + { "type": "string" }, + { "$ref": "#/definitions/asset"} ] } }, From 34d2b16bc16d7d472adfe477e4e47980b81085e4 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Wed, 30 Oct 2024 16:40:50 +0800 Subject: [PATCH 07/29] feat(decorators/assets): rewrite how asset definition is serialized --- airflow/decorators/assets.py | 7 ++++++- airflow/serialization/schema.json | 19 ++++++++++++++++--- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index f813d13d320db..2f74bcc4789ca 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -93,7 +93,12 @@ def __attrs_post_init__(self) -> None: dag._wrapped_definition = self def serialize(self): - return "AssetDefinition" + return { + "uri": self.uri, + "name": self.name, + "group": self.group, + "extra": self.extra, + } @attrs.define(kw_only=True) diff --git a/airflow/serialization/schema.json b/airflow/serialization/schema.json index 6f4386821572f..4ed6f1656d113 100644 --- a/airflow/serialization/schema.json +++ b/airflow/serialization/schema.json @@ -53,6 +53,21 @@ { "type": "integer" } ] }, + "asset_definition": { + "type": "object", + "properties": { + "uri": { "type": "string" }, + "name": { "type": "string" }, + "group": { "type": "string" }, + "extra": { + "anyOf": [ + {"type": "null"}, + { "$ref": "#/definitions/dict" } + ] + } + }, + "required": [ "uri", "extra" ] + }, "asset": { "type": "object", "properties": { @@ -182,9 +197,7 @@ "dag_dependencies": { "$ref": "#/definitions/dag_dependencies" }, "_wrapped_definition": { "anyOf": [ - { "type": "null" }, - { "type": "string" }, - { "$ref": "#/definitions/asset"} + { "$ref": "#/definitions/asset_definition"} ] } }, From c7e82b92b83231124c1c626cbf6f2ac3f303357f Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Wed, 30 Oct 2024 18:41:15 +0800 Subject: [PATCH 08/29] test(decorators/assets): add test cases to check whether asset decorator generate the right asset definition --- tests/decorators/test_assets.py | 107 ++++++++++++++++++++++++++++++++ 1 file changed, 107 insertions(+) create mode 100644 tests/decorators/test_assets.py diff --git a/tests/decorators/test_assets.py b/tests/decorators/test_assets.py new file mode 100644 index 0000000000000..ea7bbd6d665e3 --- /dev/null +++ b/tests/decorators/test_assets.py @@ -0,0 +1,107 @@ +# 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. +from __future__ import annotations + +import pytest + +from airflow.decorators.assets import asset + +pytestmark = pytest.mark.db_test + + +@pytest.fixture +def clear_assets(): + from tests_common.test_utils.db import clear_db_assets + + clear_db_assets() + yield + clear_db_assets() + + +@pytest.fixture +def example_asset_func(request): + name = "example_asset_func" + if getattr(request, "param", None) is not None: + name = request.param + + def _example_asset_func(): + return "This is example_asset" + + _example_asset_func.__name__ = name + _example_asset_func.__qualname__ = name + return _example_asset_func + + +@pytest.fixture +def example_asset_func_with_valid_arg_as_inlet_asset(): + def _example_asset_func(inlet_asset_1, inlet_asset_2): + return "This is example_asset" + + _example_asset_func.__name__ = "context" + _example_asset_func.__qualname__ = "context" + return _example_asset_func + + +class TestAssetDecorator: + def test_without_uri(self, example_asset_func): + asset_definition = asset(schedule=None)(example_asset_func) + + assert asset_definition.name == "example_asset_func" + assert asset_definition.uri == "example_asset_func" + assert asset_definition.group == "" + assert asset_definition.extra == {} + assert asset_definition.function == example_asset_func + assert asset_definition.schedule is None + + def test_with_uri(self, example_asset_func): + asset_definition = asset(schedule=None, uri="s3://bucket/object")(example_asset_func) + + assert asset_definition.name == "example_asset_func" + assert asset_definition.uri == "s3://bucket/object" + assert asset_definition.group == "" + assert asset_definition.extra == {} + assert asset_definition.function == example_asset_func + assert asset_definition.schedule is None + + def test_with_group_and_extra(self, example_asset_func): + asset_definition = asset(schedule=None, uri="s3://bucket/object", group="MLModel", extra={"k": "v"})( + example_asset_func + ) + assert asset_definition.name == "example_asset_func" + assert asset_definition.uri == "s3://bucket/object" + assert asset_definition.group == "MLModel" + assert asset_definition.extra == {"k": "v"} + assert asset_definition.function == example_asset_func + assert asset_definition.schedule is None + + def test_nested_function(self): + def root_func(): + @asset(schedule=None) + def asset_func(): + pass + + with pytest.raises(ValueError) as err: + root_func() + + assert err.value.args[0] == "nested function not supported" + + @pytest.mark.parametrize("example_asset_func", ("self", "context"), indirect=True) + def test_with_invalid_asset_name(self, example_asset_func): + with pytest.raises(ValueError) as err: + asset(schedule=None)(example_asset_func) + + assert err.value.args[0].startswith("prohibited name for asset: ") From 7cc485bf16a57f4b45bffe6c6e78ac90e8d9a64b Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Wed, 30 Oct 2024 19:07:44 +0800 Subject: [PATCH 09/29] test(decorators/assets): add test cases to AssetDefinition --- tests/decorators/test_assets.py | 45 +++++++++++++++++++++++++++++++-- 1 file changed, 43 insertions(+), 2 deletions(-) diff --git a/tests/decorators/test_assets.py b/tests/decorators/test_assets.py index ea7bbd6d665e3..8c90eb17f1073 100644 --- a/tests/decorators/test_assets.py +++ b/tests/decorators/test_assets.py @@ -16,8 +16,11 @@ # under the License. from __future__ import annotations +from unittest import mock + import pytest +from airflow.assets import Asset from airflow.decorators.assets import asset pytestmark = pytest.mark.db_test @@ -46,13 +49,20 @@ def _example_asset_func(): return _example_asset_func +@pytest.fixture +def example_asset_definition(example_asset_func): + return asset(schedule=None, uri="s3://bucket/object", group="MLModel", extra={"k": "v"})( + example_asset_func + ) + + @pytest.fixture def example_asset_func_with_valid_arg_as_inlet_asset(): def _example_asset_func(inlet_asset_1, inlet_asset_2): return "This is example_asset" - _example_asset_func.__name__ = "context" - _example_asset_func.__qualname__ = "context" + _example_asset_func.__name__ = "example_asset_func" + _example_asset_func.__qualname__ = "example_asset_func" return _example_asset_func @@ -105,3 +115,34 @@ def test_with_invalid_asset_name(self, example_asset_func): asset(schedule=None)(example_asset_func) assert err.value.args[0].startswith("prohibited name for asset: ") + + +class TestAssetDefinition: + def test_serialzie(self, example_asset_definition): + assert example_asset_definition.serialize() == { + "extra": {"k": "v"}, + "group": "MLModel", + "name": "example_asset_func", + "uri": "s3://bucket/object", + } + + @mock.patch("airflow.decorators.assets._AssetMainOperator") + @mock.patch("airflow.decorators.assets.DAG") + def test__attrs_post_init__( + self, DAG, _AssetMainOperator, example_asset_func_with_valid_arg_as_inlet_asset + ): + asset_definition = asset(schedule=None, uri="s3://bucket/object", group="MLModel", extra={"k": "v"})( + example_asset_func_with_valid_arg_as_inlet_asset + ) + + DAG.assert_called_once_with(dag_id="example_asset_func", schedule=None, auto_register=True) + _AssetMainOperator.assert_called_once_with( + task_id="__main__", + inlets=[ + Asset(name="inlet_asset_1", uri="inlet_asset_1", group="", extra={}), + Asset(name="inlet_asset_2", uri="inlet_asset_2", group="", extra={}), + ], + outlets=[asset_definition], + python_callable=example_asset_func_with_valid_arg_as_inlet_asset, + definition_name="example_asset_func", + ) From 1d31d4d42c315ae6e637b386ea50263caf16e3db Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Wed, 30 Oct 2024 19:15:52 +0800 Subject: [PATCH 10/29] test(decorators/asset): add test cases to Test_AssetMainOperator --- tests/decorators/test_assets.py | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/tests/decorators/test_assets.py b/tests/decorators/test_assets.py index 8c90eb17f1073..17ae6f8d8e135 100644 --- a/tests/decorators/test_assets.py +++ b/tests/decorators/test_assets.py @@ -21,7 +21,7 @@ import pytest from airflow.assets import Asset -from airflow.decorators.assets import asset +from airflow.decorators.assets import AssetRef, _AssetMainOperator, asset pytestmark = pytest.mark.db_test @@ -58,7 +58,7 @@ def example_asset_definition(example_asset_func): @pytest.fixture def example_asset_func_with_valid_arg_as_inlet_asset(): - def _example_asset_func(inlet_asset_1, inlet_asset_2): + def _example_asset_func(context, inlet_asset_1, inlet_asset_2): return "This is example_asset" _example_asset_func.__name__ = "example_asset_func" @@ -146,3 +146,23 @@ def test__attrs_post_init__( python_callable=example_asset_func_with_valid_arg_as_inlet_asset, definition_name="example_asset_func", ) + + +class Test_AssetMainOperator: + def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset): + asset_definition = asset(schedule=None, uri="s3://bucket/object", group="MLModel", extra={"k": "v"})( + example_asset_func_with_valid_arg_as_inlet_asset + ) + + op = _AssetMainOperator( + task_id="__main__", + inlets=[], + outlets=[asset_definition], + python_callable=example_asset_func_with_valid_arg_as_inlet_asset, + definition_name="example_asset_func", + ) + assert op.determine_kwargs(context={"k": "v"}) == { + "context": {"k": "v"}, + "inlet_asset_1": AssetRef(name="inlet_asset_1"), + "inlet_asset_2": AssetRef(name="inlet_asset_2"), + } From 5e933f4659c2aeddc10dedd91eac3526bff87a42 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Wed, 30 Oct 2024 19:16:55 +0800 Subject: [PATCH 11/29] test(decorators/assets): remove unused fixtures --- tests/decorators/test_assets.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/tests/decorators/test_assets.py b/tests/decorators/test_assets.py index 17ae6f8d8e135..12544bc33536c 100644 --- a/tests/decorators/test_assets.py +++ b/tests/decorators/test_assets.py @@ -23,17 +23,6 @@ from airflow.assets import Asset from airflow.decorators.assets import AssetRef, _AssetMainOperator, asset -pytestmark = pytest.mark.db_test - - -@pytest.fixture -def clear_assets(): - from tests_common.test_utils.db import clear_db_assets - - clear_db_assets() - yield - clear_db_assets() - @pytest.fixture def example_asset_func(request): From c69045abce240cae42f2456bd8accc1328b1d984 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Wed, 30 Oct 2024 19:28:17 +0800 Subject: [PATCH 12/29] docs(example_dag): add example dag for asset_decorator --- .../example_dags/example_asset_decorator.py | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 airflow/example_dags/example_asset_decorator.py diff --git a/airflow/example_dags/example_asset_decorator.py b/airflow/example_dags/example_asset_decorator.py new file mode 100644 index 0000000000000..583ffff0a007c --- /dev/null +++ b/airflow/example_dags/example_asset_decorator.py @@ -0,0 +1,46 @@ +# 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. +from __future__ import annotations + +import pendulum + +from airflow.assets import Asset +from airflow.decorators.assets import asset +from airflow.models.dag import DAG +from airflow.providers.standard.operators.bash import BashOperator + +asset1 = Asset(uri="s3://bucket/object") + + +@asset(uri="s3://bucket/asset_decorator", schedule=None) +def asset_decorator(asset1, asset2): + pass + + +with DAG( + dag_id="asset_consumes_1", + catchup=False, + start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), + schedule=[Asset(name="asset_decorator", uri="s3://bucket/asset_decorator")], + tags=["consumes", "asset-scheduled"], +) as dag3: + # [END dag_dep] + BashOperator( + outlets=[Asset("s3://consuming_1_task/asset_other.txt")], + task_id="consuming_1", + bash_command="sleep 5", + ) From 6c0ec478a6ae8126885c300f9b3193aaae65419e Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Thu, 31 Oct 2024 16:08:17 +0800 Subject: [PATCH 13/29] feat(decorators/assets): allow passing self and context into asset --- airflow/decorators/assets.py | 28 +++++++++++++---- .../example_dags/example_asset_decorator.py | 31 ++++++++++--------- airflow/models/dag.py | 8 +++-- .../core_api/routes/ui/test_assets.py | 2 +- tests/decorators/test_assets.py | 9 ++++-- 5 files changed, 51 insertions(+), 27 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 2f74bcc4789ca..322f8c122f67b 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -19,7 +19,9 @@ import inspect import types -from typing import TYPE_CHECKING, Any, Iterator, Mapping +from collections.abc import Callable +from functools import wraps +from typing import TYPE_CHECKING, Any, Iterator, Mapping, cast import attrs @@ -28,6 +30,8 @@ from airflow.operators.python import PythonOperator if TYPE_CHECKING: + from typing import Sequence + from airflow.io.path import ObjectStoragePath @@ -46,6 +50,7 @@ def __init__(self, *, definition_name: str, **kwargs) -> None: def _iter_kwargs(self, context: Mapping[str, Any]) -> Iterator[tuple[str, Any]]: for key in inspect.signature(self.python_callable).parameters: if key == "self": + key = "_self" value: Any = AssetRef(name=self._definition_name) elif key == "context": value = context @@ -62,6 +67,14 @@ def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]: return dict(self._iter_kwargs(context)) +def _handle_self_argument(func: types.FunctionType) -> types.FunctionType: + @wraps(func) + def wrapper(_self: Any, *args: Sequence[Any], **kwargs: dict[str, Any]) -> Any: + return func(_self, *args, **kwargs) + + return cast(types.FunctionType, wrapper) + + @attrs.define(kw_only=True) class AssetDefinition(Asset): """ @@ -75,13 +88,16 @@ class AssetDefinition(Asset): def __attrs_post_init__(self) -> None: parameters = inspect.signature(self.function).parameters + if "self" in parameters: + self.function: types.FunctionType = _handle_self_argument(cast(types.FunctionType, self.function)) # type: ignore[assignment] + with DAG(dag_id=self.name, schedule=self.schedule, auto_register=True) as dag: _AssetMainOperator( task_id="__main__", inlets=[ - Asset(name=inlet_aset_name) - for inlet_aset_name in parameters - if inlet_aset_name not in ("self", "context") + Asset(name=inlet_asset_name) + for inlet_asset_name in parameters + if inlet_asset_name not in ("self", "context") ], outlets=[self], python_callable=self.function, @@ -114,7 +130,7 @@ class asset: ) extra: dict[str, Any] = attrs.field(factory=dict) - def __call__(self, f: types.FunctionType) -> AssetDefinition: + def __call__(self, f: Callable) -> AssetDefinition: if (name := f.__name__) != f.__qualname__: raise ValueError("nested function not supported") if name == "self" or name == "context": @@ -124,6 +140,6 @@ def __call__(self, f: types.FunctionType) -> AssetDefinition: uri=name if self.uri is None else str(self.uri), group=self.group, extra=self.extra, - function=f, + function=cast(types.FunctionType, f), schedule=self.schedule, ) diff --git a/airflow/example_dags/example_asset_decorator.py b/airflow/example_dags/example_asset_decorator.py index 583ffff0a007c..5562b70af3bcc 100644 --- a/airflow/example_dags/example_asset_decorator.py +++ b/airflow/example_dags/example_asset_decorator.py @@ -19,28 +19,29 @@ import pendulum from airflow.assets import Asset +from airflow.decorators import dag, task from airflow.decorators.assets import asset -from airflow.models.dag import DAG -from airflow.providers.standard.operators.bash import BashOperator -asset1 = Asset(uri="s3://bucket/object") +asset1 = Asset(uri="s3://bucket/object", name="asset1") @asset(uri="s3://bucket/asset_decorator", schedule=None) -def asset_decorator(asset1, asset2): +def asset_decorator(self, context, asset1, asset2): pass -with DAG( - dag_id="asset_consumes_1", - catchup=False, - start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), +@dag( schedule=[Asset(name="asset_decorator", uri="s3://bucket/asset_decorator")], + start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), + catchup=False, tags=["consumes", "asset-scheduled"], -) as dag3: - # [END dag_dep] - BashOperator( - outlets=[Asset("s3://consuming_1_task/asset_other.txt")], - task_id="consuming_1", - bash_command="sleep 5", - ) +) +def consumes_asset_decorator(): + @task(outlets=[Asset("s3://consuming_1_task/asset_other.txt")]) + def process_nothing(): + pass + + process_nothing() + + +consumes_asset_decorator() diff --git a/airflow/models/dag.py b/airflow/models/dag.py index e6a67c6ad7e5e..2a7f9b46c7db0 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -119,6 +119,7 @@ from sqlalchemy.orm.query import Query from sqlalchemy.orm.session import Session + from airflow.decorators.assets import AssetDefinition from airflow.models.abstractoperator import TaskStateChangeCallback from airflow.models.dagbag import DagBag from airflow.models.operator import Operator @@ -437,6 +438,9 @@ class DAG(TaskSDKDag, LoggingMixin): # it's only use is for determining the relative fileloc based only on the serialize dag _processor_dags_folder: str | None = attrs.field(init=False, default=None) + # Additional information if this is created from an @asset definition. + _wrapped_definition: AssetDefinition | None = None + # Override the default from parent class to use config max_consecutive_failed_dag_runs: int = attrs.field( default=0, @@ -777,9 +781,7 @@ def get_is_paused(self, session=NEW_SESSION) -> None: @classmethod def get_serialized_fields(cls): """Stringified DAGs and operators contain exactly these fields.""" - return TaskSDKDag.get_serialized_fields() | { - "_processor_dags_folder", - } + return TaskSDKDag.get_serialized_fields() | {"_processor_dags_folder", "_wrapped_definition"} @staticmethod @internal_api_call diff --git a/tests/api_fastapi/core_api/routes/ui/test_assets.py b/tests/api_fastapi/core_api/routes/ui/test_assets.py index b71d80ae9d31b..b5c85b98ba6b9 100644 --- a/tests/api_fastapi/core_api/routes/ui/test_assets.py +++ b/tests/api_fastapi/core_api/routes/ui/test_assets.py @@ -47,5 +47,5 @@ def test_next_run_assets(test_client, dag_maker): assert response.status_code == 200 assert response.json() == { "asset_expression": {"all": ["s3://bucket/key/1"]}, - "events": [{"id": 17, "uri": "s3://bucket/key/1", "lastUpdate": None}], + "events": [{"id": 20, "uri": "s3://bucket/key/1", "lastUpdate": None}], } diff --git a/tests/decorators/test_assets.py b/tests/decorators/test_assets.py index 12544bc33536c..491e15e912f9f 100644 --- a/tests/decorators/test_assets.py +++ b/tests/decorators/test_assets.py @@ -17,6 +17,7 @@ from __future__ import annotations from unittest import mock +from unittest.mock import ANY import pytest @@ -47,7 +48,7 @@ def example_asset_definition(example_asset_func): @pytest.fixture def example_asset_func_with_valid_arg_as_inlet_asset(): - def _example_asset_func(context, inlet_asset_1, inlet_asset_2): + def _example_asset_func(self, context, inlet_asset_1, inlet_asset_2): return "This is example_asset" _example_asset_func.__name__ = "example_asset_func" @@ -132,10 +133,13 @@ def test__attrs_post_init__( Asset(name="inlet_asset_2", uri="inlet_asset_2", group="", extra={}), ], outlets=[asset_definition], - python_callable=example_asset_func_with_valid_arg_as_inlet_asset, + python_callable=ANY, definition_name="example_asset_func", ) + python_callable = _AssetMainOperator.call_args.kwargs["python_callable"] + assert python_callable.__wrapped__ == example_asset_func_with_valid_arg_as_inlet_asset + class Test_AssetMainOperator: def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset): @@ -151,6 +155,7 @@ def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset definition_name="example_asset_func", ) assert op.determine_kwargs(context={"k": "v"}) == { + "_self": AssetRef(name="example_asset_func"), "context": {"k": "v"}, "inlet_asset_1": AssetRef(name="inlet_asset_1"), "inlet_asset_2": AssetRef(name="inlet_asset_2"), From ad0c6d279eec0e5e1f34d7fbd8359fb1d14343b3 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 1 Nov 2024 15:17:43 +0800 Subject: [PATCH 14/29] feat(decorators/assets): return actual asset in asset decorator --- airflow/decorators/assets.py | 37 +++++++++++++++++++++++++++--------- 1 file changed, 28 insertions(+), 9 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 322f8c122f67b..ff9480ff8bbab 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -24,10 +24,13 @@ from typing import TYPE_CHECKING, Any, Iterator, Mapping, cast import attrs +from sqlalchemy import select from airflow.assets import Asset, _validate_identifier +from airflow.models.asset import AssetActive, AssetModel from airflow.models.dag import DAG, ScheduleArg from airflow.operators.python import PythonOperator +from airflow.utils.session import create_session if TYPE_CHECKING: from typing import Sequence @@ -43,27 +46,38 @@ class AssetRef: class _AssetMainOperator(PythonOperator): - def __init__(self, *, definition_name: str, **kwargs) -> None: + def __init__(self, *, definition_name: str, uri: str, **kwargs) -> None: super().__init__(**kwargs) self._definition_name = definition_name + self._uri = uri + self._active_assets: dict[str, Asset] = {} def _iter_kwargs(self, context: Mapping[str, Any]) -> Iterator[tuple[str, Any]]: + value: Any for key in inspect.signature(self.python_callable).parameters: if key == "self": key = "_self" - value: Any = AssetRef(name=self._definition_name) + value = self._active_assets.get(self._definition_name) elif key == "context": value = context else: - # TODO: This does not check if the upstream asset actually - # exists. Should we do a second pass in the DAG processor to - # raise parse-time errors if a non-existent asset is referenced? - # How? Should we also fail the task at runtime? Or should the - # dangling reference simply do nothing? - value = AssetRef(name=key) + value = self._active_assets.get(key, Asset(name=key)) yield key, value def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]: + asset_refs = [inlet for inlet in self.inlets if isinstance(inlet, AssetRef)] + if asset_refs: + with create_session() as session: + self._active_assets = { + asset_row[0].name: Asset( + name=asset_row[0].name, uri=asset_row[0].uri, group=asset_row[0].group + ) + for asset_row in session.execute( + select(AssetModel) + .join(AssetActive, AssetActive.name == AssetModel.name) + .where(AssetActive.name.in_(ref.name for ref in asset_refs)) + ) + } return dict(self._iter_kwargs(context)) @@ -95,19 +109,24 @@ def __attrs_post_init__(self) -> None: _AssetMainOperator( task_id="__main__", inlets=[ - Asset(name=inlet_asset_name) + AssetRef(name=inlet_asset_name) + # Asset(name=inlet_asset_name) for inlet_asset_name in parameters if inlet_asset_name not in ("self", "context") ], outlets=[self], python_callable=self.function, definition_name=self.name, + uri=self.uri, ) + # TODO: Currently this just gets serialized into a string. # When we create UI for assets, we should add logic to serde so the # serialized DAG contains appropriate asset information. dag._wrapped_definition = self + DAG.bulk_write_to_db([dag]) + def serialize(self): return { "uri": self.uri, From 290461a06e8e36139ee866824482c8cfa2aaec54 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 1 Nov 2024 15:50:37 +0800 Subject: [PATCH 15/29] refactor(decorators/assets): extract active assets fetching logic as _fetch_active_assets_by_name --- airflow/decorators/assets.py | 40 +++++++++++++++++++++--------------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index ff9480ff8bbab..1bfba250a6d63 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -29,13 +29,29 @@ from airflow.assets import Asset, _validate_identifier from airflow.models.asset import AssetActive, AssetModel from airflow.models.dag import DAG, ScheduleArg -from airflow.operators.python import PythonOperator -from airflow.utils.session import create_session +from airflow.providers.standard.operators.python import PythonOperator +from airflow.utils.session import NEW_SESSION, provide_session if TYPE_CHECKING: from typing import Sequence from airflow.io.path import ObjectStoragePath + from airflow.utils.session import Session + + +@provide_session +def _fetch_active_assets_by_name( + names: Sequence[str], + session: Session = NEW_SESSION, +) -> dict[str, Asset]: + return { + asset_row[0]: Asset(name=asset_row[0], uri=asset_row[1], group=asset_row[2], extra=asset_row[3]) + for asset_row in session.execute( + select(AssetModel.name, AssetModel.uri, AssetModel.group, AssetModel.extra) + .join(AssetActive, AssetActive.name == AssetModel.name) + .where(AssetActive.name.in_(name for name in names)) + ) + } @attrs.define(kw_only=True) @@ -65,19 +81,12 @@ def _iter_kwargs(self, context: Mapping[str, Any]) -> Iterator[tuple[str, Any]]: yield key, value def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]: - asset_refs = [inlet for inlet in self.inlets if isinstance(inlet, AssetRef)] - if asset_refs: - with create_session() as session: - self._active_assets = { - asset_row[0].name: Asset( - name=asset_row[0].name, uri=asset_row[0].uri, group=asset_row[0].group - ) - for asset_row in session.execute( - select(AssetModel) - .join(AssetActive, AssetActive.name == AssetModel.name) - .where(AssetActive.name.in_(ref.name for ref in asset_refs)) - ) - } + asset_names = [asset_ref.name for asset_ref in self.inlets if isinstance(asset_ref, AssetRef)] + if "self" in inspect.signature(self.python_callable).parameters: + asset_names.append(self._definition_name) + + if asset_names: + self._active_assets = _fetch_active_assets_by_name(asset_names) return dict(self._iter_kwargs(context)) @@ -110,7 +119,6 @@ def __attrs_post_init__(self) -> None: task_id="__main__", inlets=[ AssetRef(name=inlet_asset_name) - # Asset(name=inlet_asset_name) for inlet_asset_name in parameters if inlet_asset_name not in ("self", "context") ], From 00ecd552b21a0fdc2d91c0ca041bfd7035c85a5a Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 1 Nov 2024 17:03:04 +0800 Subject: [PATCH 16/29] feat(decorators/assets): allow fethcing inlet events through AssetRef --- airflow/assets/__init__.py | 18 ++++++++++++++++-- airflow/decorators/assets.py | 9 +-------- airflow/utils/context.py | 24 ++++++++++++++++++++++-- 3 files changed, 39 insertions(+), 12 deletions(-) diff --git a/airflow/assets/__init__.py b/airflow/assets/__init__.py index 59e0b86684493..f56ac8fc5e438 100644 --- a/airflow/assets/__init__.py +++ b/airflow/assets/__init__.py @@ -24,6 +24,7 @@ from typing import TYPE_CHECKING, Any, Callable, ClassVar, Iterable, Iterator, cast, overload import attr +import attrs from sqlalchemy import select from airflow.api_internal.internal_api_call import internal_api_call @@ -158,6 +159,13 @@ def expand_alias_to_assets(alias: str | AssetAlias, *, session: Session = NEW_SE return [] +@attrs.define(kw_only=True) +class AssetRef: + """Reference to an asset.""" + + name: str + + class BaseAsset: """ Protocol for all asset triggers to use in ``DAG(schedule=...)``. @@ -267,7 +275,7 @@ class Asset(os.PathLike, BaseAsset): group: str extra: dict[str, Any] - asset_type: ClassVar[str] = "" + asset_type: ClassVar[str] = "asset" __version__: ClassVar[int] = 1 @overload @@ -299,7 +307,7 @@ def __init__( fields = attr.fields_dict(Asset) self.name = _validate_non_empty_identifier(self, fields["name"], name) self.uri = _sanitize_uri(_validate_non_empty_identifier(self, fields["uri"], uri)) - self.group = _validate_identifier(self, fields["group"], group) if group else self.asset_type + self.group = _validate_identifier(self, fields["group"], group) self.extra = _set_extra_default(extra) def __fspath__(self) -> str: @@ -363,12 +371,18 @@ class Dataset(Asset): asset_type: ClassVar[str] = "dataset" + def __init__(self, *args, group="dataset", **kwargs) -> None: + super().__init__(*args, group=group, **kwargs) + class Model(Asset): """A representation of model dependencies between workflows.""" asset_type: ClassVar[str] = "model" + def __init__(self, *args, group="model", **kwargs) -> None: + super().__init__(*args, group=group, **kwargs) + class _AssetBooleanCondition(BaseAsset): """Base class for asset boolean logic.""" diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 1bfba250a6d63..654c1f7f2ebf1 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -26,7 +26,7 @@ import attrs from sqlalchemy import select -from airflow.assets import Asset, _validate_identifier +from airflow.assets import Asset, AssetRef, _validate_identifier from airflow.models.asset import AssetActive, AssetModel from airflow.models.dag import DAG, ScheduleArg from airflow.providers.standard.operators.python import PythonOperator @@ -54,13 +54,6 @@ def _fetch_active_assets_by_name( } -@attrs.define(kw_only=True) -class AssetRef: - """Reference to an asset.""" - - name: str - - class _AssetMainOperator(PythonOperator): def __init__(self, *, definition_name: str, uri: str, **kwargs) -> None: super().__init__(**kwargs) diff --git a/airflow/utils/context.py b/airflow/utils/context.py index 9af1486f914e0..8d872f79b878d 100644 --- a/airflow/utils/context.py +++ b/airflow/utils/context.py @@ -44,6 +44,7 @@ Asset, AssetAlias, AssetAliasEvent, + AssetRef, extract_event_key, ) from airflow.exceptions import RemovedInAirflow3Warning @@ -249,6 +250,7 @@ class InletEventsAccessors(Mapping[str, LazyAssetEventSelectSequence]): _inlets: list[Any] _assets: dict[str, Asset] _asset_aliases: dict[str, AssetAlias] + _asset_refs: dict[str, AssetRef] _session: Session def __init__(self, inlets: list, *, session: Session) -> None: @@ -256,12 +258,15 @@ def __init__(self, inlets: list, *, session: Session) -> None: self._session = session self._assets = {} self._asset_aliases = {} + self._asset_refs = {} for inlet in inlets: if isinstance(inlet, Asset): self._assets[inlet.uri] = inlet elif isinstance(inlet, AssetAlias): self._asset_aliases[inlet.name] = inlet + elif isinstance(inlet, AssetRef): + self._asset_refs[inlet.name] = inlet def __iter__(self) -> Iterator[str]: return iter(self._inlets) @@ -270,10 +275,15 @@ def __len__(self) -> int: return len(self._inlets) def __getitem__(self, key: int | str | Asset | AssetAlias) -> LazyAssetEventSelectSequence: + from airflow.decorators.assets import _fetch_active_assets_by_name + if isinstance(key, int): # Support index access; it's easier for trivial cases. obj = self._inlets[key] - if not isinstance(obj, (Asset, AssetAlias)): + if not isinstance(obj, (Asset, AssetAlias, Asset)): raise IndexError(key) + + if isinstance(obj, AssetRef): + obj = _fetch_active_assets_by_name([obj.name]) else: obj = key @@ -281,7 +291,17 @@ def __getitem__(self, key: int | str | Asset | AssetAlias) -> LazyAssetEventSele asset_alias = self._asset_aliases[obj.name] join_clause = AssetEvent.source_aliases where_clause = AssetAliasModel.name == asset_alias.name - elif isinstance(obj, (Asset, str)): + elif isinstance(obj, Asset): + join_clause = AssetEvent.asset + + asset = self._assets.get(obj.uri, None) + if not asset: + asset_name = self._asset_refs[obj.name].name + else: + asset_name = asset.name + + where_clause = AssetModel.name == asset_name + elif isinstance(obj, str): asset = self._assets[extract_event_key(obj)] join_clause = AssetEvent.asset where_clause = AssetModel.uri == asset.uri From 574b6024d4f406b658e109ddb3b504d9cc305f10 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 1 Nov 2024 17:40:33 +0800 Subject: [PATCH 17/29] feat(decorators/assets): reorder import paths --- airflow/decorators/assets.py | 20 +------------------- airflow/models/asset.py | 24 ++++++++++++++++++++++++ airflow/utils/context.py | 4 +--- 3 files changed, 26 insertions(+), 22 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 654c1f7f2ebf1..adfee4fe24f40 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -24,34 +24,16 @@ from typing import TYPE_CHECKING, Any, Iterator, Mapping, cast import attrs -from sqlalchemy import select from airflow.assets import Asset, AssetRef, _validate_identifier -from airflow.models.asset import AssetActive, AssetModel +from airflow.models.asset import _fetch_active_assets_by_name from airflow.models.dag import DAG, ScheduleArg from airflow.providers.standard.operators.python import PythonOperator -from airflow.utils.session import NEW_SESSION, provide_session if TYPE_CHECKING: from typing import Sequence from airflow.io.path import ObjectStoragePath - from airflow.utils.session import Session - - -@provide_session -def _fetch_active_assets_by_name( - names: Sequence[str], - session: Session = NEW_SESSION, -) -> dict[str, Asset]: - return { - asset_row[0]: Asset(name=asset_row[0], uri=asset_row[1], group=asset_row[2], extra=asset_row[3]) - for asset_row in session.execute( - select(AssetModel.name, AssetModel.uri, AssetModel.group, AssetModel.extra) - .join(AssetActive, AssetActive.name == AssetModel.name) - .where(AssetActive.name.in_(name for name in names)) - ) - } class _AssetMainOperator(PythonOperator): diff --git a/airflow/models/asset.py b/airflow/models/asset.py index 8ade71bd0b1f1..4cfae8fcceb0b 100644 --- a/airflow/models/asset.py +++ b/airflow/models/asset.py @@ -17,6 +17,7 @@ # under the License. from __future__ import annotations +from typing import TYPE_CHECKING from urllib.parse import urlsplit import sqlalchemy_jsonfield @@ -29,6 +30,7 @@ PrimaryKeyConstraint, String, Table, + select, text, ) from sqlalchemy.orm import relationship @@ -37,8 +39,30 @@ from airflow.models.base import Base, StringID from airflow.settings import json from airflow.utils import timezone +from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime +if TYPE_CHECKING: + from typing import Sequence + + from sqlalchemy.orm import Session + + +@provide_session +def _fetch_active_assets_by_name( + names: Sequence[str], + session: Session = NEW_SESSION, +) -> dict[str, Asset]: + return { + asset_row[0]: Asset(name=asset_row[0], uri=asset_row[1], group=asset_row[2], extra=asset_row[3]) + for asset_row in session.execute( + select(AssetModel.name, AssetModel.uri, AssetModel.group, AssetModel.extra) + .join(AssetActive, AssetActive.name == AssetModel.name) + .where(AssetActive.name.in_(name for name in names)) + ) + } + + alias_association_table = Table( "asset_alias_asset", Base.metadata, diff --git a/airflow/utils/context.py b/airflow/utils/context.py index 8d872f79b878d..8630fd23b364a 100644 --- a/airflow/utils/context.py +++ b/airflow/utils/context.py @@ -48,7 +48,7 @@ extract_event_key, ) from airflow.exceptions import RemovedInAirflow3Warning -from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel +from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel, _fetch_active_assets_by_name from airflow.utils.db import LazySelectSequence from airflow.utils.types import NOTSET @@ -275,8 +275,6 @@ def __len__(self) -> int: return len(self._inlets) def __getitem__(self, key: int | str | Asset | AssetAlias) -> LazyAssetEventSelectSequence: - from airflow.decorators.assets import _fetch_active_assets_by_name - if isinstance(key, int): # Support index access; it's easier for trivial cases. obj = self._inlets[key] if not isinstance(obj, (Asset, AssetAlias, Asset)): From 87287145643eed24ebadf0e640228ebc2573511c Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 1 Nov 2024 17:41:38 +0800 Subject: [PATCH 18/29] docs: update asset decorator example dag --- airflow/example_dags/example_asset_decorator.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/airflow/example_dags/example_asset_decorator.py b/airflow/example_dags/example_asset_decorator.py index 5562b70af3bcc..a9701a1f39c0f 100644 --- a/airflow/example_dags/example_asset_decorator.py +++ b/airflow/example_dags/example_asset_decorator.py @@ -22,22 +22,27 @@ from airflow.decorators import dag, task from airflow.decorators.assets import asset -asset1 = Asset(uri="s3://bucket/object", name="asset1") - -@asset(uri="s3://bucket/asset_decorator", schedule=None) -def asset_decorator(self, context, asset1, asset2): +@asset(uri="s3://bucket/asset1_producer", schedule=None) +def asset1_producer(): pass +@asset(uri="s3://bucket/object", schedule=None) +def asset2_producer(self, context, asset1_producer): + print(self) + print(context["inlet_events"][asset1_producer]) + + @dag( - schedule=[Asset(name="asset_decorator", uri="s3://bucket/asset_decorator")], + schedule=Asset(uri="s3://bucket/asset1_producer", name="asset1_producer") + | Asset(uri="s3://bucket/object", name="asset2_producer"), start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), catchup=False, tags=["consumes", "asset-scheduled"], ) def consumes_asset_decorator(): - @task(outlets=[Asset("s3://consuming_1_task/asset_other.txt")]) + @task(outlets=[Asset(uri="s3://consuming_1_task/asset_other.txt", name="endpoint")]) def process_nothing(): pass From 6c2539ae17a217f06c1fbbcf5974bc8a1a5c51ed Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 1 Nov 2024 17:53:35 +0800 Subject: [PATCH 19/29] test: fix tests --- airflow/decorators/assets.py | 16 ++++++++++++++-- airflow/example_dags/example_asset_decorator.py | 2 +- tests/always/test_example_dags.py | 4 ++++ tests/decorators/test_assets.py | 13 ++++++++----- tests/serialization/test_dag_serialization.py | 1 + 5 files changed, 28 insertions(+), 8 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index adfee4fe24f40..c89b089d4557b 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -37,7 +37,7 @@ class _AssetMainOperator(PythonOperator): - def __init__(self, *, definition_name: str, uri: str, **kwargs) -> None: + def __init__(self, *, definition_name: str, uri: str | None = None, **kwargs) -> None: super().__init__(**kwargs) self._definition_name = definition_name self._uri = uri @@ -49,6 +49,10 @@ def _iter_kwargs(self, context: Mapping[str, Any]) -> Iterator[tuple[str, Any]]: if key == "self": key = "_self" value = self._active_assets.get(self._definition_name) + if not value: + value = Asset(name=self._definition_name) + if self._uri is not None: + value.uri = self._uri elif key == "context": value = context else: @@ -97,7 +101,7 @@ def __attrs_post_init__(self) -> None: for inlet_asset_name in parameters if inlet_asset_name not in ("self", "context") ], - outlets=[self], + outlets=[self.to_asset()], python_callable=self.function, definition_name=self.name, uri=self.uri, @@ -110,6 +114,14 @@ def __attrs_post_init__(self) -> None: DAG.bulk_write_to_db([dag]) + def to_asset(self) -> Asset: + return Asset( + name=self.name, + uri=self.uri, + group=self.group, + extra=self.extra, + ) + def serialize(self): return { "uri": self.uri, diff --git a/airflow/example_dags/example_asset_decorator.py b/airflow/example_dags/example_asset_decorator.py index a9701a1f39c0f..b4de09c23146b 100644 --- a/airflow/example_dags/example_asset_decorator.py +++ b/airflow/example_dags/example_asset_decorator.py @@ -42,7 +42,7 @@ def asset2_producer(self, context, asset1_producer): tags=["consumes", "asset-scheduled"], ) def consumes_asset_decorator(): - @task(outlets=[Asset(uri="s3://consuming_1_task/asset_other.txt", name="endpoint")]) + @task(outlets=[Asset(name="process_nothing")]) def process_nothing(): pass diff --git a/tests/always/test_example_dags.py b/tests/always/test_example_dags.py index 7ea1ac9609bd2..4c4700af2b25b 100644 --- a/tests/always/test_example_dags.py +++ b/tests/always/test_example_dags.py @@ -199,6 +199,10 @@ def test_should_be_importable(example: str): @pytest.mark.db_test @pytest.mark.parametrize("example", example_not_excluded_dags(xfail_db_exception=True)) def test_should_not_do_database_queries(example: str): + if example.endswith("airflow/example_dags/example_asset_decorator.py"): + # The asset decorator tries to write DAG and Asset into db which is expected + return + with assert_queries_count(1, stacklevel_from_module=example.rsplit(os.sep, 1)[-1]): DagBag( dag_folder=example, diff --git a/tests/decorators/test_assets.py b/tests/decorators/test_assets.py index 491e15e912f9f..c5d4edadd024f 100644 --- a/tests/decorators/test_assets.py +++ b/tests/decorators/test_assets.py @@ -24,6 +24,8 @@ from airflow.assets import Asset from airflow.decorators.assets import AssetRef, _AssetMainOperator, asset +pytestmark = pytest.mark.db_test + @pytest.fixture def example_asset_func(request): @@ -129,12 +131,13 @@ def test__attrs_post_init__( _AssetMainOperator.assert_called_once_with( task_id="__main__", inlets=[ - Asset(name="inlet_asset_1", uri="inlet_asset_1", group="", extra={}), - Asset(name="inlet_asset_2", uri="inlet_asset_2", group="", extra={}), + AssetRef(name="inlet_asset_1"), + AssetRef(name="inlet_asset_2"), ], outlets=[asset_definition], python_callable=ANY, definition_name="example_asset_func", + uri="s3://bucket/object", ) python_callable = _AssetMainOperator.call_args.kwargs["python_callable"] @@ -155,8 +158,8 @@ def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset definition_name="example_asset_func", ) assert op.determine_kwargs(context={"k": "v"}) == { - "_self": AssetRef(name="example_asset_func"), + "_self": Asset(name="example_asset_func"), "context": {"k": "v"}, - "inlet_asset_1": AssetRef(name="inlet_asset_1"), - "inlet_asset_2": AssetRef(name="inlet_asset_2"), + "inlet_asset_1": Asset(name="inlet_asset_1"), + "inlet_asset_2": Asset(name="inlet_asset_2"), } diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py index 0111cc669cd02..eaac6116cc65a 100644 --- a/tests/serialization/test_dag_serialization.py +++ b/tests/serialization/test_dag_serialization.py @@ -568,6 +568,7 @@ def validate_deserialized_dag(self, serialized_dag: DAG, dag: DAG): "task_group", "params", "_processor_dags_folder", + "_wrapped_definition", } fields_to_check = dag.get_serialized_fields() - exclusion_list for field in fields_to_check: From d42ce8c94431f10433f54776f501fac13daf16f3 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 1 Nov 2024 18:06:33 +0800 Subject: [PATCH 20/29] test(decorators/assets): extend test_determine_kwargs to cover active asset --- tests/decorators/test_assets.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/decorators/test_assets.py b/tests/decorators/test_assets.py index c5d4edadd024f..9105a5cec520b 100644 --- a/tests/decorators/test_assets.py +++ b/tests/decorators/test_assets.py @@ -23,6 +23,7 @@ from airflow.assets import Asset from airflow.decorators.assets import AssetRef, _AssetMainOperator, asset +from airflow.models.asset import AssetActive, AssetModel pytestmark = pytest.mark.db_test @@ -145,14 +146,19 @@ def test__attrs_post_init__( class Test_AssetMainOperator: - def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset): + def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset, session): + example_asset = AssetModel(uri="s3://bucket/object1", name="inlet_asset_1") + session.add(example_asset) + session.add(AssetActive.for_asset(example_asset)) + session.commit() + asset_definition = asset(schedule=None, uri="s3://bucket/object", group="MLModel", extra={"k": "v"})( example_asset_func_with_valid_arg_as_inlet_asset ) op = _AssetMainOperator( task_id="__main__", - inlets=[], + inlets=[AssetRef(name="inlet_asset_1"), AssetRef(name="inlet_asset_2")], outlets=[asset_definition], python_callable=example_asset_func_with_valid_arg_as_inlet_asset, definition_name="example_asset_func", @@ -160,6 +166,6 @@ def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset assert op.determine_kwargs(context={"k": "v"}) == { "_self": Asset(name="example_asset_func"), "context": {"k": "v"}, - "inlet_asset_1": Asset(name="inlet_asset_1"), + "inlet_asset_1": Asset(name="inlet_asset_1", uri="s3://bucket/object1"), "inlet_asset_2": Asset(name="inlet_asset_2"), } From 17fce697b00bc48596b6ff55b8277f5fec8b5144 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Mon, 4 Nov 2024 17:25:16 +0800 Subject: [PATCH 21/29] fix: address easy to fix comments --- airflow/assets/__init__.py | 13 +++---------- airflow/decorators/assets.py | 13 +++++++------ airflow/models/asset.py | 6 +++--- 3 files changed, 13 insertions(+), 19 deletions(-) diff --git a/airflow/assets/__init__.py b/airflow/assets/__init__.py index f56ac8fc5e438..e69f819607cfc 100644 --- a/airflow/assets/__init__.py +++ b/airflow/assets/__init__.py @@ -23,7 +23,6 @@ import warnings from typing import TYPE_CHECKING, Any, Callable, ClassVar, Iterable, Iterator, cast, overload -import attr import attrs from sqlalchemy import select @@ -215,7 +214,7 @@ def iter_dag_dependencies(self, *, source: str, target: str) -> Iterator[DagDepe raise NotImplementedError -@attr.define(unsafe_hash=False) +@attrs.define(unsafe_hash=False) class AssetAlias(BaseAsset): """A represeation of asset alias which is used to create asset during the runtime.""" @@ -266,7 +265,7 @@ def _set_extra_default(extra: dict | None) -> dict: return extra -@attr.define(init=False, unsafe_hash=False) +@attrs.define(init=False, unsafe_hash=False) class Asset(os.PathLike, BaseAsset): """A representation of data asset dependencies between workflows.""" @@ -307,7 +306,7 @@ def __init__( fields = attr.fields_dict(Asset) self.name = _validate_non_empty_identifier(self, fields["name"], name) self.uri = _sanitize_uri(_validate_non_empty_identifier(self, fields["uri"], uri)) - self.group = _validate_identifier(self, fields["group"], group) + self.group = _validate_identifier(self, fields["group"], group) if group else self.asset_type self.extra = _set_extra_default(extra) def __fspath__(self) -> str: @@ -371,18 +370,12 @@ class Dataset(Asset): asset_type: ClassVar[str] = "dataset" - def __init__(self, *args, group="dataset", **kwargs) -> None: - super().__init__(*args, group=group, **kwargs) - class Model(Asset): """A representation of model dependencies between workflows.""" asset_type: ClassVar[str] = "model" - def __init__(self, *args, group="model", **kwargs) -> None: - super().__init__(*args, group=group, **kwargs) - class _AssetBooleanCondition(BaseAsset): """Base class for asset boolean logic.""" diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index c89b089d4557b..76250c9491aea 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -19,9 +19,8 @@ import inspect import types -from collections.abc import Callable from functools import wraps -from typing import TYPE_CHECKING, Any, Iterator, Mapping, cast +from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping, cast import attrs @@ -41,9 +40,10 @@ def __init__(self, *, definition_name: str, uri: str | None = None, **kwargs) -> super().__init__(**kwargs) self._definition_name = definition_name self._uri = uri - self._active_assets: dict[str, Asset] = {} - def _iter_kwargs(self, context: Mapping[str, Any]) -> Iterator[tuple[str, Any]]: + def _iter_kwargs( + self, context: Mapping[str, Any], active_assets: dict[str, Asset] + ) -> Iterator[tuple[str, Any]]: value: Any for key in inspect.signature(self.python_callable).parameters: if key == "self": @@ -60,13 +60,14 @@ def _iter_kwargs(self, context: Mapping[str, Any]) -> Iterator[tuple[str, Any]]: yield key, value def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]: + active_assets: dict[str, Asset] = {} asset_names = [asset_ref.name for asset_ref in self.inlets if isinstance(asset_ref, AssetRef)] if "self" in inspect.signature(self.python_callable).parameters: asset_names.append(self._definition_name) if asset_names: - self._active_assets = _fetch_active_assets_by_name(asset_names) - return dict(self._iter_kwargs(context)) + active_assets = _fetch_active_assets_by_name(asset_names) + return dict(self._iter_kwargs(context, active_assets)) def _handle_self_argument(func: types.FunctionType) -> types.FunctionType: diff --git a/airflow/models/asset.py b/airflow/models/asset.py index 4cfae8fcceb0b..f19174d15b1db 100644 --- a/airflow/models/asset.py +++ b/airflow/models/asset.py @@ -54,9 +54,9 @@ def _fetch_active_assets_by_name( session: Session = NEW_SESSION, ) -> dict[str, Asset]: return { - asset_row[0]: Asset(name=asset_row[0], uri=asset_row[1], group=asset_row[2], extra=asset_row[3]) - for asset_row in session.execute( - select(AssetModel.name, AssetModel.uri, AssetModel.group, AssetModel.extra) + asset_model.name: asset_model.to_public() + for asset_model in session.execute( + select(AssetModel) .join(AssetActive, AssetActive.name == AssetModel.name) .where(AssetActive.name.in_(name for name in names)) ) From de73b467db0fc2b5fab15a07361892ebd14321e5 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Mon, 4 Nov 2024 18:34:08 +0800 Subject: [PATCH 22/29] fix: fix asset serialization --- airflow/assets/__init__.py | 17 ++++++++++++----- airflow/decorators/assets.py | 4 ++-- airflow/models/asset.py | 2 +- airflow/serialization/enums.py | 1 + airflow/serialization/serialized_objects.py | 9 +++++++-- airflow/utils/context.py | 2 +- tests/decorators/test_assets.py | 2 +- tests/timetables/test_assets_timetable.py | 4 ++-- 8 files changed, 27 insertions(+), 14 deletions(-) diff --git a/airflow/assets/__init__.py b/airflow/assets/__init__.py index e69f819607cfc..60662b68630ac 100644 --- a/airflow/assets/__init__.py +++ b/airflow/assets/__init__.py @@ -123,6 +123,13 @@ def _validate_non_empty_identifier(instance, attribute, value): return value +def _validate_asset_name(instance, attribute, value): + _validate_non_empty_identifier(instance, attribute, value) + if value == "self" or value == "context": + raise ValueError(f"prohibited name for asset: {value}") + return value + + def extract_event_key(value: str | Asset | AssetAlias) -> str: """ Extract the key of an inlet or an outlet event. @@ -218,11 +225,11 @@ def iter_dag_dependencies(self, *, source: str, target: str) -> Iterator[DagDepe class AssetAlias(BaseAsset): """A represeation of asset alias which is used to create asset during the runtime.""" - name: str = attr.field(validator=_validate_non_empty_identifier) - group: str = attr.field( + name: str = attrs.field(validator=_validate_non_empty_identifier) + group: str = attrs.field( kw_only=True, default="", - validator=[attr.validators.max_len(1500), _validate_identifier], + validator=[attrs.validators.max_len(1500), _validate_identifier], ) def iter_assets(self) -> Iterator[tuple[str, Asset]]: @@ -303,8 +310,8 @@ def __init__( name = uri elif uri is None: uri = name - fields = attr.fields_dict(Asset) - self.name = _validate_non_empty_identifier(self, fields["name"], name) + fields = attrs.fields_dict(Asset) + self.name = _validate_asset_name(self, fields["name"], name) self.uri = _sanitize_uri(_validate_non_empty_identifier(self, fields["uri"], uri)) self.group = _validate_identifier(self, fields["group"], group) if group else self.asset_type self.extra = _set_extra_default(extra) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 76250c9491aea..6607504ee3099 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -48,7 +48,7 @@ def _iter_kwargs( for key in inspect.signature(self.python_callable).parameters: if key == "self": key = "_self" - value = self._active_assets.get(self._definition_name) + value = active_assets.get(self._definition_name) if not value: value = Asset(name=self._definition_name) if self._uri is not None: @@ -56,7 +56,7 @@ def _iter_kwargs( elif key == "context": value = context else: - value = self._active_assets.get(key, Asset(name=key)) + value = active_assets.get(key, Asset(name=key)) yield key, value def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]: diff --git a/airflow/models/asset.py b/airflow/models/asset.py index f19174d15b1db..1609db9a99186 100644 --- a/airflow/models/asset.py +++ b/airflow/models/asset.py @@ -54,7 +54,7 @@ def _fetch_active_assets_by_name( session: Session = NEW_SESSION, ) -> dict[str, Asset]: return { - asset_model.name: asset_model.to_public() + asset_model[0].name: asset_model[0].to_public() for asset_model in session.execute( select(AssetModel) .join(AssetActive, AssetActive.name == AssetModel.name) diff --git a/airflow/serialization/enums.py b/airflow/serialization/enums.py index dd63366b8a958..d1b946b38ef1e 100644 --- a/airflow/serialization/enums.py +++ b/airflow/serialization/enums.py @@ -59,6 +59,7 @@ class DagAttributeTypes(str, Enum): ASSET_ALIAS = "asset_alias" ASSET_ANY = "asset_any" ASSET_ALL = "asset_all" + ASSET_REF = "asset_ref" SIMPLE_TASK_INSTANCE = "simple_task_instance" BASE_JOB = "Job" TASK_INSTANCE = "task_instance" diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py index 52b0bcb1530a0..4b7ee6d0871b5 100644 --- a/airflow/serialization/serialized_objects.py +++ b/airflow/serialization/serialized_objects.py @@ -40,6 +40,7 @@ AssetAlias, AssetAll, AssetAny, + AssetRef, BaseAsset, _AssetAliasCondition, ) @@ -254,7 +255,7 @@ def encode_asset_condition(var: BaseAsset) -> dict[str, Any]: :meta private: """ if isinstance(var, Asset): - return {"__type": DAT.ASSET, "uri": var.uri, "extra": var.extra} + return {"__type": DAT.ASSET, "name": var.name, "uri": var.uri, "extra": var.extra} if isinstance(var, AssetAlias): return {"__type": DAT.ASSET_ALIAS, "name": var.name} if isinstance(var, AssetAll): @@ -272,7 +273,7 @@ def decode_asset_condition(var: dict[str, Any]) -> BaseAsset: """ dat = var["__type"] if dat == DAT.ASSET: - return Asset(var["uri"], extra=var["extra"]) + return Asset(uri=var["uri"], name=var["name"], extra=var["extra"]) if dat == DAT.ASSET_ALL: return AssetAll(*(decode_asset_condition(x) for x in var["objects"])) if dat == DAT.ASSET_ANY: @@ -743,6 +744,8 @@ def serialize( elif isinstance(var, BaseAsset): serialized_asset = encode_asset_condition(var) return cls._encode(serialized_asset, type_=serialized_asset.pop("__type")) + elif isinstance(var, AssetRef): + return cls._encode({"name": var.name}, type_=DAT.ASSET_REF) elif isinstance(var, SimpleTaskInstance): return cls._encode( cls.serialize(var.__dict__, strict=strict, use_pydantic_models=use_pydantic_models), @@ -876,6 +879,8 @@ def deserialize(cls, encoded_var: Any, use_pydantic_models=False) -> Any: return AssetAny(*(decode_asset_condition(x) for x in var["objects"])) elif type_ == DAT.ASSET_ALL: return AssetAll(*(decode_asset_condition(x) for x in var["objects"])) + elif type_ == DAT.ASSET_REF: + return AssetRef(name=var["name"]) elif type_ == DAT.SIMPLE_TASK_INSTANCE: return SimpleTaskInstance(**cls.deserialize(var)) elif type_ == DAT.CONNECTION: diff --git a/airflow/utils/context.py b/airflow/utils/context.py index 8630fd23b364a..4385178d19dca 100644 --- a/airflow/utils/context.py +++ b/airflow/utils/context.py @@ -277,7 +277,7 @@ def __len__(self) -> int: def __getitem__(self, key: int | str | Asset | AssetAlias) -> LazyAssetEventSelectSequence: if isinstance(key, int): # Support index access; it's easier for trivial cases. obj = self._inlets[key] - if not isinstance(obj, (Asset, AssetAlias, Asset)): + if not isinstance(obj, (Asset, AssetAlias)): raise IndexError(key) if isinstance(obj, AssetRef): diff --git a/tests/decorators/test_assets.py b/tests/decorators/test_assets.py index 9105a5cec520b..290d7a6b45c1b 100644 --- a/tests/decorators/test_assets.py +++ b/tests/decorators/test_assets.py @@ -135,7 +135,7 @@ def test__attrs_post_init__( AssetRef(name="inlet_asset_1"), AssetRef(name="inlet_asset_2"), ], - outlets=[asset_definition], + outlets=[asset_definition.to_asset()], python_callable=ANY, definition_name="example_asset_func", uri="s3://bucket/object", diff --git a/tests/timetables/test_assets_timetable.py b/tests/timetables/test_assets_timetable.py index f461afa31c890..bb942a4a01d43 100644 --- a/tests/timetables/test_assets_timetable.py +++ b/tests/timetables/test_assets_timetable.py @@ -134,7 +134,7 @@ def test_serialization(asset_timetable: AssetOrTimeSchedule, monkeypatch: Any) - "timetable": "mock_serialized_timetable", "asset_condition": { "__type": "asset_all", - "objects": [{"__type": "asset", "uri": "test_asset", "extra": {}}], + "objects": [{"__type": "asset", "uri": "test_asset", "name": "test_asset", "extra": {}}], }, } @@ -152,7 +152,7 @@ def test_deserialization(monkeypatch: Any) -> None: "timetable": "mock_serialized_timetable", "asset_condition": { "__type": "asset_all", - "objects": [{"__type": "asset", "uri": "test_asset", "extra": None}], + "objects": [{"__type": "asset", "name": "test_asset", "uri": "test_asset", "extra": None}], }, } deserialized = AssetOrTimeSchedule.deserialize(mock_serialized_data) From 0b76c81f32e36bba2b8e13f7ce6e2030c4872598 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Tue, 5 Nov 2024 12:05:47 +0800 Subject: [PATCH 23/29] refactor(decorators/assets): postpone the attribute check to AssetDefinition instead of asset decorator --- airflow/decorators/assets.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 6607504ee3099..e50f87c0ad5b6 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -24,7 +24,7 @@ import attrs -from airflow.assets import Asset, AssetRef, _validate_identifier +from airflow.assets import Asset, AssetRef from airflow.models.asset import _fetch_active_assets_by_name from airflow.models.dag import DAG, ScheduleArg from airflow.providers.standard.operators.python import PythonOperator @@ -138,18 +138,13 @@ class asset: schedule: ScheduleArg uri: str | ObjectStoragePath | None = None - group: str = attrs.field( - kw_only=True, - default="", - validator=[attrs.validators.max_len(1500), _validate_identifier], - ) + group: str = "" extra: dict[str, Any] = attrs.field(factory=dict) def __call__(self, f: Callable) -> AssetDefinition: if (name := f.__name__) != f.__qualname__: raise ValueError("nested function not supported") - if name == "self" or name == "context": - raise ValueError(f"prohibited name for asset: {name}") + return AssetDefinition( name=name, uri=name if self.uri is None else str(self.uri), From 69523dcdcec3cb9c1fc48a99062814b8f0f00e1d Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 7 Nov 2024 18:18:16 +0800 Subject: [PATCH 24/29] Simplify group validators The validate_identifier validator already checks the length, so we don't need an extra one doing that. --- airflow/assets/__init__.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/airflow/assets/__init__.py b/airflow/assets/__init__.py index 60662b68630ac..f1d36ac12b730 100644 --- a/airflow/assets/__init__.py +++ b/airflow/assets/__init__.py @@ -226,11 +226,7 @@ class AssetAlias(BaseAsset): """A represeation of asset alias which is used to create asset during the runtime.""" name: str = attrs.field(validator=_validate_non_empty_identifier) - group: str = attrs.field( - kw_only=True, - default="", - validator=[attrs.validators.max_len(1500), _validate_identifier], - ) + group: str = attrs.field(kw_only=True, default="", validator=_validate_identifier) def iter_assets(self) -> Iterator[tuple[str, Asset]]: return iter(()) From 81ee77bd2bea52f05e52e60286b274718d3fcd5c Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 8 Nov 2024 10:32:09 +0800 Subject: [PATCH 25/29] style(dag): remove _wrapped_definition --- airflow/decorators/assets.py | 5 ----- airflow/models/dag.py | 6 +----- airflow/serialization/schema.json | 7 +------ tests/serialization/test_dag_serialization.py | 1 - 4 files changed, 2 insertions(+), 17 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index e50f87c0ad5b6..2a54281c6e796 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -108,11 +108,6 @@ def __attrs_post_init__(self) -> None: uri=self.uri, ) - # TODO: Currently this just gets serialized into a string. - # When we create UI for assets, we should add logic to serde so the - # serialized DAG contains appropriate asset information. - dag._wrapped_definition = self - DAG.bulk_write_to_db([dag]) def to_asset(self) -> Asset: diff --git a/airflow/models/dag.py b/airflow/models/dag.py index 2a7f9b46c7db0..e48ec0a9a9c57 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -119,7 +119,6 @@ from sqlalchemy.orm.query import Query from sqlalchemy.orm.session import Session - from airflow.decorators.assets import AssetDefinition from airflow.models.abstractoperator import TaskStateChangeCallback from airflow.models.dagbag import DagBag from airflow.models.operator import Operator @@ -438,9 +437,6 @@ class DAG(TaskSDKDag, LoggingMixin): # it's only use is for determining the relative fileloc based only on the serialize dag _processor_dags_folder: str | None = attrs.field(init=False, default=None) - # Additional information if this is created from an @asset definition. - _wrapped_definition: AssetDefinition | None = None - # Override the default from parent class to use config max_consecutive_failed_dag_runs: int = attrs.field( default=0, @@ -781,7 +777,7 @@ def get_is_paused(self, session=NEW_SESSION) -> None: @classmethod def get_serialized_fields(cls): """Stringified DAGs and operators contain exactly these fields.""" - return TaskSDKDag.get_serialized_fields() | {"_processor_dags_folder", "_wrapped_definition"} + return TaskSDKDag.get_serialized_fields() | {"_processor_dags_folder"} @staticmethod @internal_api_call diff --git a/airflow/serialization/schema.json b/airflow/serialization/schema.json index 4ed6f1656d113..1e7232aa81a97 100644 --- a/airflow/serialization/schema.json +++ b/airflow/serialization/schema.json @@ -194,12 +194,7 @@ { "$ref": "#/definitions/task_group" } ]}, "edge_info": { "$ref": "#/definitions/edge_info" }, - "dag_dependencies": { "$ref": "#/definitions/dag_dependencies" }, - "_wrapped_definition": { - "anyOf": [ - { "$ref": "#/definitions/asset_definition"} - ] - } + "dag_dependencies": { "$ref": "#/definitions/dag_dependencies" } }, "required": [ "dag_id", diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py index eaac6116cc65a..0111cc669cd02 100644 --- a/tests/serialization/test_dag_serialization.py +++ b/tests/serialization/test_dag_serialization.py @@ -568,7 +568,6 @@ def validate_deserialized_dag(self, serialized_dag: DAG, dag: DAG): "task_group", "params", "_processor_dags_folder", - "_wrapped_definition", } fields_to_check = dag.get_serialized_fields() - exclusion_list for field in fields_to_check: From 7ccc9abd7043c2b26b70ee4cf3707e2e4e73f43b Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 8 Nov 2024 10:49:40 +0800 Subject: [PATCH 26/29] style(decorators/assets): change types.FunctionType to Callable --- airflow/decorators/assets.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 2a54281c6e796..01635f4839b8c 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -18,9 +18,8 @@ from __future__ import annotations import inspect -import types from functools import wraps -from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping, cast +from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping import attrs @@ -70,12 +69,12 @@ def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]: return dict(self._iter_kwargs(context, active_assets)) -def _handle_self_argument(func: types.FunctionType) -> types.FunctionType: +def _handle_self_argument(func: Callable) -> Callable: @wraps(func) def wrapper(_self: Any, *args: Sequence[Any], **kwargs: dict[str, Any]) -> Any: return func(_self, *args, **kwargs) - return cast(types.FunctionType, wrapper) + return wrapper @attrs.define(kw_only=True) @@ -86,13 +85,13 @@ class AssetDefinition(Asset): :meta private: """ - function: types.FunctionType + function: Callable schedule: ScheduleArg def __attrs_post_init__(self) -> None: parameters = inspect.signature(self.function).parameters if "self" in parameters: - self.function: types.FunctionType = _handle_self_argument(cast(types.FunctionType, self.function)) # type: ignore[assignment] + self.function: Callable = _handle_self_argument(self.function) with DAG(dag_id=self.name, schedule=self.schedule, auto_register=True) as dag: _AssetMainOperator( @@ -145,6 +144,6 @@ def __call__(self, f: Callable) -> AssetDefinition: uri=name if self.uri is None else str(self.uri), group=self.group, extra=self.extra, - function=cast(types.FunctionType, f), + function=f, schedule=self.schedule, ) From 52c8b8508fd1112fd3937d8d674e1030013e803b Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 8 Nov 2024 11:06:32 +0800 Subject: [PATCH 27/29] refactor(decorators/assets): make session in _fetch_active_assets_by_name required --- airflow/decorators/assets.py | 4 +++- airflow/models/asset.py | 4 +--- airflow/utils/context.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 01635f4839b8c..455ffc86311a5 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -27,6 +27,7 @@ from airflow.models.asset import _fetch_active_assets_by_name from airflow.models.dag import DAG, ScheduleArg from airflow.providers.standard.operators.python import PythonOperator +from airflow.utils.session import create_session if TYPE_CHECKING: from typing import Sequence @@ -65,7 +66,8 @@ def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]: asset_names.append(self._definition_name) if asset_names: - active_assets = _fetch_active_assets_by_name(asset_names) + with create_session() as session: + active_assets = _fetch_active_assets_by_name(asset_names, session) return dict(self._iter_kwargs(context, active_assets)) diff --git a/airflow/models/asset.py b/airflow/models/asset.py index 1609db9a99186..50914d51650b8 100644 --- a/airflow/models/asset.py +++ b/airflow/models/asset.py @@ -39,7 +39,6 @@ from airflow.models.base import Base, StringID from airflow.settings import json from airflow.utils import timezone -from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime if TYPE_CHECKING: @@ -48,10 +47,9 @@ from sqlalchemy.orm import Session -@provide_session def _fetch_active_assets_by_name( names: Sequence[str], - session: Session = NEW_SESSION, + session: Session, ) -> dict[str, Asset]: return { asset_model[0].name: asset_model[0].to_public() diff --git a/airflow/utils/context.py b/airflow/utils/context.py index 4385178d19dca..a02c01459c1af 100644 --- a/airflow/utils/context.py +++ b/airflow/utils/context.py @@ -281,7 +281,7 @@ def __getitem__(self, key: int | str | Asset | AssetAlias) -> LazyAssetEventSele raise IndexError(key) if isinstance(obj, AssetRef): - obj = _fetch_active_assets_by_name([obj.name]) + obj = _fetch_active_assets_by_name([obj.name], self._session) else: obj = key From 7e4ecf84cf05cc4c829ca2f03f5f4ded819478fc Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 8 Nov 2024 16:58:44 +0800 Subject: [PATCH 28/29] fix(decorators/asets): remove DAG.bulk_write_to_db and remove self handling --- airflow/decorators/assets.py | 22 +--------------------- tests/always/test_example_dags.py | 4 ---- tests/decorators/test_assets.py | 20 +++++++++++++------- 3 files changed, 14 insertions(+), 32 deletions(-) diff --git a/airflow/decorators/assets.py b/airflow/decorators/assets.py index 455ffc86311a5..2f5052c2d5c9c 100644 --- a/airflow/decorators/assets.py +++ b/airflow/decorators/assets.py @@ -18,7 +18,6 @@ from __future__ import annotations import inspect -from functools import wraps from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping import attrs @@ -30,8 +29,6 @@ from airflow.utils.session import create_session if TYPE_CHECKING: - from typing import Sequence - from airflow.io.path import ObjectStoragePath @@ -47,12 +44,7 @@ def _iter_kwargs( value: Any for key in inspect.signature(self.python_callable).parameters: if key == "self": - key = "_self" value = active_assets.get(self._definition_name) - if not value: - value = Asset(name=self._definition_name) - if self._uri is not None: - value.uri = self._uri elif key == "context": value = context else: @@ -71,14 +63,6 @@ def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]: return dict(self._iter_kwargs(context, active_assets)) -def _handle_self_argument(func: Callable) -> Callable: - @wraps(func) - def wrapper(_self: Any, *args: Sequence[Any], **kwargs: dict[str, Any]) -> Any: - return func(_self, *args, **kwargs) - - return wrapper - - @attrs.define(kw_only=True) class AssetDefinition(Asset): """ @@ -92,10 +76,8 @@ class AssetDefinition(Asset): def __attrs_post_init__(self) -> None: parameters = inspect.signature(self.function).parameters - if "self" in parameters: - self.function: Callable = _handle_self_argument(self.function) - with DAG(dag_id=self.name, schedule=self.schedule, auto_register=True) as dag: + with DAG(dag_id=self.name, schedule=self.schedule, auto_register=True): _AssetMainOperator( task_id="__main__", inlets=[ @@ -109,8 +91,6 @@ def __attrs_post_init__(self) -> None: uri=self.uri, ) - DAG.bulk_write_to_db([dag]) - def to_asset(self) -> Asset: return Asset( name=self.name, diff --git a/tests/always/test_example_dags.py b/tests/always/test_example_dags.py index 4c4700af2b25b..7ea1ac9609bd2 100644 --- a/tests/always/test_example_dags.py +++ b/tests/always/test_example_dags.py @@ -199,10 +199,6 @@ def test_should_be_importable(example: str): @pytest.mark.db_test @pytest.mark.parametrize("example", example_not_excluded_dags(xfail_db_exception=True)) def test_should_not_do_database_queries(example: str): - if example.endswith("airflow/example_dags/example_asset_decorator.py"): - # The asset decorator tries to write DAG and Asset into db which is expected - return - with assert_queries_count(1, stacklevel_from_module=example.rsplit(os.sep, 1)[-1]): DagBag( dag_folder=example, diff --git a/tests/decorators/test_assets.py b/tests/decorators/test_assets.py index 290d7a6b45c1b..a3821140e5488 100644 --- a/tests/decorators/test_assets.py +++ b/tests/decorators/test_assets.py @@ -142,20 +142,24 @@ def test__attrs_post_init__( ) python_callable = _AssetMainOperator.call_args.kwargs["python_callable"] - assert python_callable.__wrapped__ == example_asset_func_with_valid_arg_as_inlet_asset + assert python_callable == example_asset_func_with_valid_arg_as_inlet_asset class Test_AssetMainOperator: def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset, session): - example_asset = AssetModel(uri="s3://bucket/object1", name="inlet_asset_1") - session.add(example_asset) - session.add(AssetActive.for_asset(example_asset)) - session.commit() - + example_asset_model = AssetModel(uri="s3://bucket/object1", name="inlet_asset_1") asset_definition = asset(schedule=None, uri="s3://bucket/object", group="MLModel", extra={"k": "v"})( example_asset_func_with_valid_arg_as_inlet_asset ) + ad_asset_model = AssetModel.from_public(asset_definition) + + session.add(example_asset_model) + session.add(ad_asset_model) + session.add(AssetActive.for_asset(example_asset_model)) + session.add(AssetActive.for_asset(ad_asset_model)) + session.commit() + op = _AssetMainOperator( task_id="__main__", inlets=[AssetRef(name="inlet_asset_1"), AssetRef(name="inlet_asset_2")], @@ -164,7 +168,9 @@ def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset definition_name="example_asset_func", ) assert op.determine_kwargs(context={"k": "v"}) == { - "_self": Asset(name="example_asset_func"), + "self": Asset( + name="example_asset_func", uri="s3://bucket/object", group="MLModel", extra={"k": "v"} + ), "context": {"k": "v"}, "inlet_asset_1": Asset(name="inlet_asset_1", uri="s3://bucket/object1"), "inlet_asset_2": Asset(name="inlet_asset_2"), From 2060ead13d2a0be8eb5cbb1e61d631a1e8d37717 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 8 Nov 2024 17:49:46 +0800 Subject: [PATCH 29/29] feat(utils/context): fetch asset_refs all at once --- airflow/utils/context.py | 29 +++++++++++------------------ 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/airflow/utils/context.py b/airflow/utils/context.py index a02c01459c1af..b28559999e509 100644 --- a/airflow/utils/context.py +++ b/airflow/utils/context.py @@ -250,7 +250,6 @@ class InletEventsAccessors(Mapping[str, LazyAssetEventSelectSequence]): _inlets: list[Any] _assets: dict[str, Asset] _asset_aliases: dict[str, AssetAlias] - _asset_refs: dict[str, AssetRef] _session: Session def __init__(self, inlets: list, *, session: Session) -> None: @@ -258,15 +257,19 @@ def __init__(self, inlets: list, *, session: Session) -> None: self._session = session self._assets = {} self._asset_aliases = {} - self._asset_refs = {} + _asset_ref_names: list[str] = [] for inlet in inlets: if isinstance(inlet, Asset): - self._assets[inlet.uri] = inlet + self._assets[inlet.name] = inlet elif isinstance(inlet, AssetAlias): self._asset_aliases[inlet.name] = inlet elif isinstance(inlet, AssetRef): - self._asset_refs[inlet.name] = inlet + _asset_ref_names.append(inlet.name) + + if _asset_ref_names: + for asset_name, asset in _fetch_active_assets_by_name(_asset_ref_names, self._session).items(): + self._assets[asset_name] = asset def __iter__(self) -> Iterator[str]: return iter(self._inlets) @@ -277,11 +280,8 @@ def __len__(self) -> int: def __getitem__(self, key: int | str | Asset | AssetAlias) -> LazyAssetEventSelectSequence: if isinstance(key, int): # Support index access; it's easier for trivial cases. obj = self._inlets[key] - if not isinstance(obj, (Asset, AssetAlias)): + if not isinstance(obj, (Asset, AssetAlias, AssetRef)): raise IndexError(key) - - if isinstance(obj, AssetRef): - obj = _fetch_active_assets_by_name([obj.name], self._session) else: obj = key @@ -289,20 +289,13 @@ def __getitem__(self, key: int | str | Asset | AssetAlias) -> LazyAssetEventSele asset_alias = self._asset_aliases[obj.name] join_clause = AssetEvent.source_aliases where_clause = AssetAliasModel.name == asset_alias.name - elif isinstance(obj, Asset): + elif isinstance(obj, (Asset, AssetRef)): join_clause = AssetEvent.asset - - asset = self._assets.get(obj.uri, None) - if not asset: - asset_name = self._asset_refs[obj.name].name - else: - asset_name = asset.name - - where_clause = AssetModel.name == asset_name + where_clause = AssetModel.name == self._assets[obj.name].name elif isinstance(obj, str): asset = self._assets[extract_event_key(obj)] join_clause = AssetEvent.asset - where_clause = AssetModel.uri == asset.uri + where_clause = AssetModel.name == asset.name else: raise ValueError(key)