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 5aba0e2b97922..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,6 +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", + 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__}.") @@ -1242,12 +1247,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 @@ -1391,9 +1398,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, @@ -1479,16 +1487,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 @@ -1615,13 +1624,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 @@ -1635,9 +1645,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 551e773988719..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): @@ -2093,3 +2109,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()