From 275e30de732d62cd68a8da7b1f87b466b0330739 Mon Sep 17 00:00:00 2001 From: shanah Date: Fri, 13 Mar 2026 09:12:27 +0000 Subject: [PATCH 1/5] feat(providers/amazon): allow disabling hook-level lineage in S3Hook (#63371) --- .../airflow/providers/amazon/aws/hooks/s3.py | 67 +++++++++++-------- .../tests/unit/amazon/aws/hooks/test_s3.py | 34 ++++++++++ 2 files changed, 72 insertions(+), 29 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py index d0f8ce8b78460..dbaf3ac04233b 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py @@ -189,6 +189,7 @@ def __init__( aws_conn_id: str | None | ArgNotSet = AwsBaseHook.default_conn_name, transfer_config_args: dict | None = None, extra_args: dict | None = None, + enable_hook_level_lineage: bool = True, *args, **kwargs, ) -> None: @@ -203,6 +204,7 @@ def __init__( if extra_args and not isinstance(extra_args, dict): raise TypeError(f"extra_args expected dict, got {type(extra_args).__name__}.") self._extra_args = extra_args or {} + self.enable_hook_level_lineage = enable_hook_level_lineage super().__init__(*args, **kwargs) @@ -1225,12 +1227,14 @@ def load_file( ExtraArgs=extra_args, Config=self.transfer_config, ) - get_hook_lineage_collector().add_input_asset( - context=self, scheme="file", asset_kwargs={"path": filename} - ) - get_hook_lineage_collector().add_output_asset( - context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} - ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_input_asset( + context=self, scheme="file", asset_kwargs={"path": filename} + ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_output_asset( + context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} + ) @unify_bucket_name_and_key @provide_bucket_name @@ -1374,9 +1378,10 @@ def _upload_file_obj( Config=self.transfer_config, ) # No input because file_obj can be anything - handle in calling function if possible - get_hook_lineage_collector().add_output_asset( - context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} - ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_output_asset( + context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} + ) def copy_object( self, @@ -1462,16 +1467,18 @@ def copy_object( CopySource=copy_source, **kwargs, ) - get_hook_lineage_collector().add_input_asset( - context=self, - scheme="s3", - asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key}, - ) - get_hook_lineage_collector().add_output_asset( - context=self, - scheme="s3", - asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key}, - ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_input_asset( + context=self, + scheme="s3", + asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key}, + ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_output_asset( + context=self, + scheme="s3", + asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key}, + ) return response @provide_bucket_name @@ -1598,13 +1605,14 @@ def download_file( file_path.parent.mkdir(exist_ok=True, parents=True) - get_hook_lineage_collector().add_output_asset( - context=self, - scheme="file", - asset_kwargs={ - "path": str(file_path) if file_path.is_absolute() else str(file_path.absolute()) - }, - ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_output_asset( + context=self, + scheme="file", + asset_kwargs={ + "path": str(file_path) if file_path.is_absolute() else str(file_path.absolute()) + }, + ) file = open(file_path, "wb") else: file = NamedTemporaryFile(dir=local_path, prefix="airflow_tmp_", delete=False) # type: ignore @@ -1618,9 +1626,10 @@ def download_file( Config=self.transfer_config, ) file.flush() - get_hook_lineage_collector().add_input_asset( - context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} - ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_input_asset( + context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} + ) return file.name def generate_presigned_url( diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py index 99fd45aff8903..471a6e388e9eb 100644 --- a/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py @@ -2041,3 +2041,37 @@ def test_unify_and_provide_ordered_properly(): matches = re.findall(r"@provide_bucket_name\s+@unify_bucket_name_and_key", code, re.MULTILINE) if matches: pytest.fail("@unify_bucket_name_and_key should be applied before @provide_bucket_name in S3Hook") + + +class TestS3HookLineageConfig: + def test_hook_lineage_enabled_by_default(self): + from airflow.providers.amazon.aws.hooks.s3 import S3Hook + + hook = S3Hook() + assert hook.enable_hook_level_lineage is True + + def test_hook_lineage_disabled_when_flag_false(self): + from airflow.providers.amazon.aws.hooks.s3 import S3Hook + + hook = S3Hook(enable_hook_level_lineage=False) + assert hook.enable_hook_level_lineage is False + + from unittest import mock + + @mock.patch("airflow.providers.amazon.aws.hooks.s3.get_hook_lineage_collector") + @mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_conn") + def test_load_string_skips_lineage_when_disabled(self, mock_conn, mock_collector): + from airflow.providers.amazon.aws.hooks.s3 import S3Hook + + hook = S3Hook(enable_hook_level_lineage=False) + hook.load_string("data", "key", bucket_name="bucket", replace=True) + mock_collector.return_value.add_output_asset.assert_not_called() + + @mock.patch("airflow.providers.amazon.aws.hooks.s3.get_hook_lineage_collector") + @mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_conn") + def test_load_string_exposes_lineage_when_enabled(self, mock_conn, mock_collector): + from airflow.providers.amazon.aws.hooks.s3 import S3Hook + + hook = S3Hook(enable_hook_level_lineage=True) + hook.load_string("data", "key", bucket_name="bucket", replace=True) + mock_collector.return_value.add_output_asset.assert_called_once() From 0b3dec6cd8a0defc511b2d76bdf3b084ce9e3663 Mon Sep 17 00:00:00 2001 From: shanah Date: Sat, 14 Mar 2026 08:00:36 +0000 Subject: [PATCH 2/5] fix: move enable_hook_level_lineage from S3Hook __init__ signature to kwargs.pop --- .../airflow/providers/amazon/aws/hooks/s3.py | 58 +++++++++---------- 1 file changed, 26 insertions(+), 32 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py index dbaf3ac04233b..3b8bdf163bcaf 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py @@ -189,13 +189,13 @@ def __init__( aws_conn_id: str | None | ArgNotSet = AwsBaseHook.default_conn_name, transfer_config_args: dict | None = None, extra_args: dict | None = None, - enable_hook_level_lineage: bool = True, *args, **kwargs, ) -> None: kwargs["client_type"] = "s3" kwargs["aws_conn_id"] = aws_conn_id self._requester_pays = kwargs.pop("requester_pays", False) + self.enable_hook_level_lineage = kwargs.pop("enable_hook_level_lineage", True) if transfer_config_args and not isinstance(transfer_config_args, dict): raise TypeError(f"transfer_config_args expected dict, got {type(transfer_config_args).__name__}.") @@ -204,7 +204,6 @@ def __init__( if extra_args and not isinstance(extra_args, dict): raise TypeError(f"extra_args expected dict, got {type(extra_args).__name__}.") self._extra_args = extra_args or {} - self.enable_hook_level_lineage = enable_hook_level_lineage super().__init__(*args, **kwargs) @@ -1227,11 +1226,10 @@ def load_file( ExtraArgs=extra_args, Config=self.transfer_config, ) - if self.enable_hook_level_lineage: - get_hook_lineage_collector().add_input_asset( - context=self, scheme="file", asset_kwargs={"path": filename} - ) - if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_input_asset( + context=self, scheme="file", asset_kwargs={"path": filename} + ) + if getattr(self, "enable_hook_level_lineage", True): get_hook_lineage_collector().add_output_asset( context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} ) @@ -1378,7 +1376,7 @@ def _upload_file_obj( Config=self.transfer_config, ) # No input because file_obj can be anything - handle in calling function if possible - if self.enable_hook_level_lineage: + if getattr(self, "enable_hook_level_lineage", True): get_hook_lineage_collector().add_output_asset( context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} ) @@ -1467,18 +1465,16 @@ def copy_object( CopySource=copy_source, **kwargs, ) - if self.enable_hook_level_lineage: - get_hook_lineage_collector().add_input_asset( - context=self, - scheme="s3", - asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key}, - ) - if self.enable_hook_level_lineage: - get_hook_lineage_collector().add_output_asset( - context=self, - scheme="s3", - asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key}, - ) + get_hook_lineage_collector().add_input_asset( + context=self, + scheme="s3", + asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key}, + ) + get_hook_lineage_collector().add_output_asset( + context=self, + scheme="s3", + asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key}, + ) return response @provide_bucket_name @@ -1605,14 +1601,13 @@ def download_file( file_path.parent.mkdir(exist_ok=True, parents=True) - if self.enable_hook_level_lineage: - get_hook_lineage_collector().add_output_asset( - context=self, - scheme="file", - asset_kwargs={ - "path": str(file_path) if file_path.is_absolute() else str(file_path.absolute()) - }, - ) + get_hook_lineage_collector().add_output_asset( + context=self, + scheme="file", + asset_kwargs={ + "path": str(file_path) if file_path.is_absolute() else str(file_path.absolute()) + }, + ) file = open(file_path, "wb") else: file = NamedTemporaryFile(dir=local_path, prefix="airflow_tmp_", delete=False) # type: ignore @@ -1626,10 +1621,9 @@ def download_file( Config=self.transfer_config, ) file.flush() - if self.enable_hook_level_lineage: - get_hook_lineage_collector().add_input_asset( - context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} - ) + get_hook_lineage_collector().add_input_asset( + context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} + ) return file.name def generate_presigned_url( From 563e565f5ca718be5da1402c513ae565f0cf9327 Mon Sep 17 00:00:00 2001 From: shnhdan Date: Mon, 16 Mar 2026 20:33:50 +0530 Subject: [PATCH 3/5] Update providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com> --- providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py index 3b8bdf163bcaf..f30dff4229ac4 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py @@ -1229,7 +1229,7 @@ def load_file( get_hook_lineage_collector().add_input_asset( context=self, scheme="file", asset_kwargs={"path": filename} ) - if getattr(self, "enable_hook_level_lineage", True): + if self.enable_hook_level_lineage: get_hook_lineage_collector().add_output_asset( context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} ) From 927f7ab1673ee5202d86b26a78f0f43c1aaaf56e Mon Sep 17 00:00:00 2001 From: shnhdan Date: Mon, 8 Jun 2026 08:12:34 +0000 Subject: [PATCH 4/5] fix: gate all remaining hook lineage calls behind enable_hook_level_lineage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously load_file (input asset), copy_object (both assets), and download_file (both assets) were ungated — firing even when enable_hook_level_lineage=False. All get_hook_lineage_collector() calls in S3Hook are now consistently gated behind the flag. Also normalized getattr() guard to direct attribute access. --- .../airflow/providers/amazon/aws/hooks/s3.py | 52 ++++++++++--------- 1 file changed, 28 insertions(+), 24 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py index a30266f458931..6f5ea0fc20e97 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py @@ -1243,9 +1243,10 @@ def load_file( ExtraArgs=extra_args, Config=self.transfer_config, ) - get_hook_lineage_collector().add_input_asset( - context=self, scheme="file", asset_kwargs={"path": filename} - ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_input_asset( + context=self, scheme="file", asset_kwargs={"path": filename} + ) if self.enable_hook_level_lineage: get_hook_lineage_collector().add_output_asset( context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} @@ -1393,7 +1394,7 @@ def _upload_file_obj( Config=self.transfer_config, ) # No input because file_obj can be anything - handle in calling function if possible - if getattr(self, "enable_hook_level_lineage", True): + if self.enable_hook_level_lineage: get_hook_lineage_collector().add_output_asset( context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} ) @@ -1482,16 +1483,17 @@ def copy_object( CopySource=copy_source, **kwargs, ) - get_hook_lineage_collector().add_input_asset( - context=self, - scheme="s3", - asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key}, - ) - get_hook_lineage_collector().add_output_asset( - context=self, - scheme="s3", - asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key}, - ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_input_asset( + context=self, + scheme="s3", + asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key}, + ) + get_hook_lineage_collector().add_output_asset( + context=self, + scheme="s3", + asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key}, + ) return response @provide_bucket_name @@ -1618,13 +1620,14 @@ def download_file( file_path.parent.mkdir(exist_ok=True, parents=True) - get_hook_lineage_collector().add_output_asset( - context=self, - scheme="file", - asset_kwargs={ - "path": str(file_path) if file_path.is_absolute() else str(file_path.absolute()) - }, - ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_output_asset( + context=self, + scheme="file", + asset_kwargs={ + "path": str(file_path) if file_path.is_absolute() else str(file_path.absolute()) + }, + ) file = open(file_path, "wb") else: file = NamedTemporaryFile(dir=local_path, prefix="airflow_tmp_", delete=False) # type: ignore @@ -1638,9 +1641,10 @@ def download_file( Config=self.transfer_config, ) file.flush() - get_hook_lineage_collector().add_input_asset( - context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} - ) + if self.enable_hook_level_lineage: + get_hook_lineage_collector().add_input_asset( + context=self, scheme="s3", asset_kwargs={"bucket": bucket_name, "key": key} + ) return file.name def generate_presigned_url( From b1c78866d8d40d49bfdca3ed7f435f27af2fcd79 Mon Sep 17 00:00:00 2001 From: shnhdan Date: Sun, 14 Jun 2026 14:35:17 +0000 Subject: [PATCH 5/5] feat: add default_hook_lineage config for S3Hook lineage control Following the default_deferrable pattern, add a new [lineage] section to airflow.cfg with default_hook_lineage boolean (default: True). S3Hook.__init__ now reads this as the default for enable_hook_level_lineage instead of hardcoding True. Allows global disabling via config while permitting per-instance override. Usage: [lineage] default_hook_lineage = False # airflow.cfg S3Hook(enable_hook_level_lineage=False) # per-instance As suggested by @kacpermuda in review. Pattern can be extended to other hooks in future PRs. Related: #66992 --- .../airflow/config_templates/default_airflow.cfg | 9 +++++++++ .../src/airflow/providers/amazon/aws/hooks/s3.py | 6 +++++- .../tests/unit/amazon/aws/hooks/test_s3.py | 16 ++++++++++++++++ 3 files changed, 30 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/config_templates/default_airflow.cfg b/airflow-core/src/airflow/config_templates/default_airflow.cfg index f0f9c1b60c823..51c7cf534f622 100644 --- a/airflow-core/src/airflow/config_templates/default_airflow.cfg +++ b/airflow-core/src/airflow/config_templates/default_airflow.cfg @@ -43,3 +43,12 @@ # # airflow config list --defaults > "${AIRFLOW_HOME}/airflow.cfg" # + +[lineage] + +# When enabled, hooks will collect lineage data by default. +# Individual hooks can still override this via enable_hook_level_lineage kwarg. +# Set to False to disable hook-level lineage collection globally (e.g. to suppress +# intermediate asset spam during multipart uploads). +default_hook_lineage = True + diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py index 6f5ea0fc20e97..1bf3088d7b3dd 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3.py @@ -61,6 +61,7 @@ from airflow.providers.amazon.aws.exceptions import S3HookUriParseFailure from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook from airflow.providers.amazon.aws.utils.tags import format_tags +from airflow.configuration import conf from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException from airflow.utils.helpers import chunks @@ -196,7 +197,10 @@ def __init__( kwargs["client_type"] = "s3" kwargs["aws_conn_id"] = aws_conn_id self._requester_pays = kwargs.pop("requester_pays", False) - self.enable_hook_level_lineage = kwargs.pop("enable_hook_level_lineage", True) + self.enable_hook_level_lineage = kwargs.pop( + "enable_hook_level_lineage", + conf.getboolean("lineage", "default_hook_lineage", fallback=True), + ) if transfer_config_args and not isinstance(transfer_config_args, dict): raise TypeError(f"transfer_config_args expected dict, got {type(transfer_config_args).__name__}.") diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py index e31260f195189..8426fd4b99be4 100644 --- a/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py @@ -77,6 +77,22 @@ def s3_bucket(mocked_s3_res): return bucket + def test_hook_lineage_reads_from_conf(self): + """Test that enable_hook_level_lineage defaults to conf value.""" + from unittest.mock import patch + from airflow.configuration import conf + + # When conf sets default_hook_lineage=False, hook should default to False + with patch.object(conf, "getboolean", return_value=False) as mock_conf: + hook = S3Hook() + mock_conf.assert_called_with("lineage", "default_hook_lineage", fallback=True) + assert hook.enable_hook_level_lineage is False + + # Per-instance override still works even when conf says False + hook_override = S3Hook(enable_hook_level_lineage=True) + assert hook_override.enable_hook_level_lineage is True + + class TestAwsS3Hook: @mock_aws def test_get_conn(self):