This is an automated email from the ASF dual-hosted git repository. vincbeck pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push: new 21c1e1c3e66 Add Key to default meta keys in S3KeySensor (#50122) 21c1e1c3e66 is described below commit 21c1e1c3e6623983d11f04400a47c12f001a67d7 Author: Dmitry Dandykin <49510873+dukast...@users.noreply.github.com> AuthorDate: Mon May 12 15:59:29 2025 +0200 Add Key to default meta keys in S3KeySensor (#50122) Co-authored-by: Dmitry Dandykin <dmitry.dandy...@axenix.pro> --- .../src/airflow/providers/amazon/aws/sensors/s3.py | 5 ++-- .../tests/unit/amazon/aws/sensors/test_s3.py | 35 ++++++++-------------- 2 files changed, 16 insertions(+), 24 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/s3.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/s3.py index 1bea70ae255..8a34812eca6 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/s3.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/s3.py @@ -107,7 +107,7 @@ class S3KeySensor(AwsBaseSensor[S3Hook]): self.verify = verify self.deferrable = deferrable self.use_regex = use_regex - self.metadata_keys = metadata_keys if metadata_keys else ["Size"] + self.metadata_keys = metadata_keys if metadata_keys else ["Size", "Key"] def _check_key(self, key, context: Context): bucket_name, key = self.hook.get_s3_bucket_key(self.bucket_name, key, "bucket_name", "bucket_key") @@ -116,7 +116,8 @@ class S3KeySensor(AwsBaseSensor[S3Hook]): """ Set variable `files` which contains a list of dict which contains attributes defined by the user Format: [{ - 'Size': int + 'Size': int, + 'Key': str, }] """ if self.wildcard_match: diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_s3.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_s3.py index b9f8fb284bc..4ada09d0923 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_s3.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_s3.py @@ -320,11 +320,19 @@ class TestS3KeySensor: with pytest.raises(AirflowException, match=message): op.execute_complete(context={}, event={"status": "error", "message": message}) + @pytest.mark.parametrize( + "metadata_keys, expected", + [ + (["Size", "Key"], True), + (["Content"], False), + (None, True), + ], + ) @mock_aws - def test_custom_metadata_default_return_vals(self): + def test_custom_metadata_default_return_vals(self, metadata_keys, expected): def check_fn(files: list) -> bool: for f in files: - if "Size" not in f: + if "Size" not in f or "Key" not in f: return False return True @@ -335,31 +343,14 @@ class TestS3KeySensor: key="test-key", string_data="test-body", ) - - op = S3KeySensor( - task_id="test-metadata", - bucket_key="test-key", - bucket_name="test-bucket", - metadata_keys=["Size"], - check_fn=check_fn, - ) - assert op.poke(None) is True op = S3KeySensor( task_id="test-metadata", bucket_key="test-key", bucket_name="test-bucket", - metadata_keys=["Content"], + metadata_keys=metadata_keys, check_fn=check_fn, ) - assert op.poke(None) is False - - op = S3KeySensor( - task_id="test-metadata", - bucket_key="test-key", - bucket_name="test-bucket", - check_fn=check_fn, - ) - assert op.poke(None) is True + assert op.poke(None) is expected @mock_aws def test_custom_metadata_default_custom_vals(self): @@ -391,7 +382,7 @@ class TestS3KeySensor: def check_fn(files: list) -> bool: hook = S3Hook() metadata_keys = set(hook.head_object(bucket_name="test-bucket", key="test-key").keys()) - test_data_keys = set(files[0].keys()) + test_data_keys = set(files[0].keys()) - {"Key"} return test_data_keys == metadata_keys