You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by el...@apache.org on 2022/04/12 05:23:09 UTC
[airflow] branch main updated: Deprecate `S3PrefixSensor` and `S3KeySizeSensor` in favor of `S3KeySensor` (#22737)
This is an automated email from the ASF dual-hosted git repository.
eladkal 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 dffb0d27f2 Deprecate `S3PrefixSensor` and `S3KeySizeSensor` in favor of `S3KeySensor` (#22737)
dffb0d27f2 is described below
commit dffb0d27f25b9a1f2497535cb87de3b889aae9d0
Author: Vincent <97...@users.noreply.github.com>
AuthorDate: Tue Apr 12 01:22:51 2022 -0400
Deprecate `S3PrefixSensor` and `S3KeySizeSensor` in favor of `S3KeySensor` (#22737)
Deprecate `S3PrefixSensor` and `S3KeySizeSensor` in favor of `S3KeySensor`
Co-authored-by: Ash Berlin-Taylor <as...@firemirror.com>
---
.../amazon/aws/example_dags/example_s3_bucket.py | 58 +++++-
airflow/providers/amazon/aws/hooks/s3.py | 59 +++++-
airflow/providers/amazon/aws/sensors/s3.py | 218 +++++++++------------
airflow/sensors/s3_prefix_sensor.py | 4 +-
.../operators/s3.rst | 57 +++++-
tests/providers/amazon/aws/hooks/test_s3.py | 20 ++
tests/providers/amazon/aws/sensors/test_s3_key.py | 160 +++++++++------
.../providers/amazon/aws/sensors/test_s3_prefix.py | 41 ++--
8 files changed, 399 insertions(+), 218 deletions(-)
diff --git a/airflow/providers/amazon/aws/example_dags/example_s3_bucket.py b/airflow/providers/amazon/aws/example_dags/example_s3_bucket.py
index d04bbab048..3e1c41d862 100644
--- a/airflow/providers/amazon/aws/example_dags/example_s3_bucket.py
+++ b/airflow/providers/amazon/aws/example_dags/example_s3_bucket.py
@@ -17,6 +17,7 @@
import os
from datetime import datetime
+from typing import List
from airflow.models.baseoperator import chain
from airflow.models.dag import DAG
@@ -27,11 +28,15 @@ from airflow.providers.amazon.aws.operators.s3 import (
S3GetBucketTaggingOperator,
S3PutBucketTaggingOperator,
)
+from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
BUCKET_NAME = os.environ.get('BUCKET_NAME', 'test-airflow-12345')
+KEY = os.environ.get('KEY', 'key')
+KEY_2 = os.environ.get('KEY_2', 'key2')
TAG_KEY = os.environ.get('TAG_KEY', 'test-s3-bucket-tagging-key')
TAG_VALUE = os.environ.get('TAG_VALUE', 'test-s3-bucket-tagging-value')
+
with DAG(
dag_id='example_s3_bucket',
schedule_interval=None,
@@ -39,6 +44,22 @@ with DAG(
catchup=False,
tags=['example'],
) as dag:
+ # [START howto_sensor_s3_key_function_definition]
+ def check_fn(files: List) -> bool:
+ """
+ Example of custom check: check if all files are bigger than 1kB
+
+ :param files: List of S3 object attributes.
+ Format: [{
+ 'Size': int
+ }]
+ :return: true if the criteria is met
+ :rtype: bool
+ """
+ return all(f.get('Size', 0) > 1024 for f in files)
+
+ # [END howto_sensor_s3_key_function_definition]
+
# [START howto_operator_s3_create_bucket]
create_bucket = S3CreateBucketOperator(
task_id='s3_create_bucket',
@@ -69,10 +90,45 @@ with DAG(
)
# [END howto_operator_s3_delete_bucket_tagging]
+ # [START howto_sensor_s3_key_single_key]
+ # Check if a file exists
+ s3_sensor_one_key = S3KeySensor(
+ task_id="s3_sensor_one_key",
+ bucket_name=BUCKET_NAME,
+ bucket_key=KEY,
+ )
+ # [END howto_sensor_s3_key_single_key]
+
+ # [START howto_sensor_s3_key_multiple_keys]
+ # Check if both files exist
+ s3_sensor_two_keys = S3KeySensor(
+ task_id="s3_sensor_two_keys",
+ bucket_name=BUCKET_NAME,
+ bucket_key=[KEY, KEY_2],
+ )
+ # [END howto_sensor_s3_key_multiple_keys]
+
+ # [START howto_sensor_s3_key_function]
+ # Check if a file exists and match a certain pattern defined in check_fn
+ s3_sensor_key_function = S3KeySensor(
+ task_id="s3_sensor_key_function",
+ bucket_name=BUCKET_NAME,
+ bucket_key=KEY,
+ check_fn=check_fn,
+ )
+ # [END howto_sensor_s3_key_function]
+
# [START howto_operator_s3_delete_bucket]
delete_bucket = S3DeleteBucketOperator(
task_id='s3_delete_bucket', bucket_name=BUCKET_NAME, force_delete=True
)
# [END howto_operator_s3_delete_bucket]
- chain(create_bucket, put_tagging, get_tagging, delete_tagging, delete_bucket)
+ chain(
+ create_bucket,
+ put_tagging,
+ get_tagging,
+ delete_tagging,
+ [s3_sensor_one_key, s3_sensor_two_keys, s3_sensor_key_function],
+ delete_bucket,
+ )
diff --git a/airflow/providers/amazon/aws/hooks/s3.py b/airflow/providers/amazon/aws/hooks/s3.py
index fc23255a78..38416b4e89 100644
--- a/airflow/providers/amazon/aws/hooks/s3.py
+++ b/airflow/providers/amazon/aws/hooks/s3.py
@@ -348,26 +348,71 @@ class S3Hook(AwsBaseHook):
return self._list_key_object_filter(keys, from_datetime, to_datetime)
+ @provide_bucket_name
+ def get_file_metadata(
+ self,
+ prefix: str,
+ bucket_name: Optional[str] = None,
+ page_size: Optional[int] = None,
+ max_items: Optional[int] = None,
+ ) -> List:
+ """
+ Lists metadata objects in a bucket under prefix
+
+ :param prefix: a key prefix
+ :param bucket_name: the name of the bucket
+ :param page_size: pagination size
+ :param max_items: maximum items to return
+ :return: a list of metadata of objects
+ :rtype: list
+ """
+ config = {
+ 'PageSize': page_size,
+ 'MaxItems': max_items,
+ }
+
+ paginator = self.get_conn().get_paginator('list_objects_v2')
+ response = paginator.paginate(Bucket=bucket_name, Prefix=prefix, PaginationConfig=config)
+
+ files = []
+ for page in response:
+ if 'Contents' in page:
+ files += page['Contents']
+ return files
+
@provide_bucket_name
@unify_bucket_name_and_key
- def check_for_key(self, key: str, bucket_name: Optional[str] = None) -> bool:
+ def head_object(self, key: str, bucket_name: Optional[str] = None) -> Optional[dict]:
"""
- Checks if a key exists in a bucket
+ Retrieves metadata of an object
:param key: S3 key that will point to the file
:param bucket_name: Name of the bucket in which the file is stored
- :return: True if the key exists and False if not.
- :rtype: bool
+ :return: metadata of an object
+ :rtype: dict
"""
try:
- self.get_conn().head_object(Bucket=bucket_name, Key=key)
- return True
+ return self.get_conn().head_object(Bucket=bucket_name, Key=key)
except ClientError as e:
if e.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
- return False
+ return None
else:
raise e
+ @provide_bucket_name
+ @unify_bucket_name_and_key
+ def check_for_key(self, key: str, bucket_name: Optional[str] = None) -> bool:
+ """
+ Checks if a key exists in a bucket
+
+ :param key: S3 key that will point to the file
+ :param bucket_name: Name of the bucket in which the file is stored
+ :return: True if the key exists and False if not.
+ :rtype: bool
+ """
+ obj = self.head_object(key, bucket_name)
+ return obj is not None
+
@provide_bucket_name
@unify_bucket_name_and_key
def get_key(self, key: str, bucket_name: Optional[str] = None) -> S3Transfer:
diff --git a/airflow/providers/amazon/aws/sensors/s3.py b/airflow/providers/amazon/aws/sensors/s3.py
index 2bda06639f..3c8a210734 100644
--- a/airflow/providers/amazon/aws/sensors/s3.py
+++ b/airflow/providers/amazon/aws/sensors/s3.py
@@ -20,6 +20,7 @@
import os
import re
import sys
+import warnings
from datetime import datetime
from typing import TYPE_CHECKING, Callable, List, Optional, Sequence, Set, Union
from urllib.parse import urlparse
@@ -40,17 +41,30 @@ from airflow.sensors.base import BaseSensorOperator, poke_mode_only
class S3KeySensor(BaseSensorOperator):
"""
- Waits for a key (a file-like instance on S3) to be present in a S3 bucket.
+ Waits for one or multiple keys (a file-like instance on S3) to be present in a S3 bucket.
S3 being a key/value it does not support folders. The path is just a key
a resource.
- :param bucket_key: The key being waited on. Supports full s3:// style url
+ .. seealso::
+ For more information on how to use this sensor, take a look at the guide:
+ :ref:`howto/sensor:S3KeySensor`
+
+ :param bucket_key: The key(s) being waited on. Supports full s3:// style url
or relative path from root level. When it's specified as a full s3://
- url, please leave bucket_name as `None`.
+ url, please leave bucket_name as `None`
:param bucket_name: Name of the S3 bucket. Only needed when ``bucket_key``
- is not provided as a full s3:// url.
+ is not provided as a full s3:// url. When specified, all the keys passed to ``bucket_key``
+ refers to this bucket
:param wildcard_match: whether the bucket_key should be interpreted as a
Unix wildcard pattern
+ :param check_fn: Function that receives the list of the S3 objects,
+ and returns a boolean:
+ - ``True``: the criteria is met
+ - ``False``: the criteria isn't met
+ **Example**: Wait for any S3 object size more than 1 megabyte ::
+
+ def check_fn(files: List) -> bool:
+ return any(f.get('Size', 0) > 1048576 for f in files)
:param aws_conn_id: a reference to the s3 connection
:param verify: Whether or not to verify SSL certificates for S3 connection.
By default SSL certificates are verified.
@@ -69,36 +83,65 @@ class S3KeySensor(BaseSensorOperator):
def __init__(
self,
*,
- bucket_key: str,
+ bucket_key: Union[str, List[str]],
bucket_name: Optional[str] = None,
wildcard_match: bool = False,
+ check_fn: Optional[Callable[..., bool]] = None,
aws_conn_id: str = 'aws_default',
verify: Optional[Union[str, bool]] = None,
**kwargs,
):
super().__init__(**kwargs)
self.bucket_name = bucket_name
- self.bucket_key = bucket_key
+ self.bucket_key = [bucket_key] if isinstance(bucket_key, str) else bucket_key
self.wildcard_match = wildcard_match
+ self.check_fn = check_fn
self.aws_conn_id = aws_conn_id
self.verify = verify
self.hook: Optional[S3Hook] = None
- def _resolve_bucket_and_key(self):
+ def _resolve_bucket_and_key(self, key):
"""If key is URI, parse bucket"""
if self.bucket_name is None:
- self.bucket_name, self.bucket_key = S3Hook.parse_s3_url(self.bucket_key)
+ return S3Hook.parse_s3_url(key)
else:
- parsed_url = urlparse(self.bucket_key)
+ parsed_url = urlparse(key)
if parsed_url.scheme != '' or parsed_url.netloc != '':
raise AirflowException('If bucket_name provided, bucket_key must be relative path, not URI.')
+ return self.bucket_name, key
- def poke(self, context: 'Context'):
- self._resolve_bucket_and_key()
- self.log.info('Poking for key : s3://%s/%s', self.bucket_name, self.bucket_key)
+ def _check_key(self, key):
+ bucket_name, key = self._resolve_bucket_and_key(key)
+ self.log.info('Poking for key : s3://%s/%s', bucket_name, key)
+
+ """
+ Set variable `files` which contains a list of dict which contains only the size
+ If needed we might want to add other attributes later
+ Format: [{
+ 'Size': int
+ }]
+ """
if self.wildcard_match:
- return self.get_hook().check_for_wildcard_key(self.bucket_key, self.bucket_name)
- return self.get_hook().check_for_key(self.bucket_key, self.bucket_name)
+ prefix = re.split(r'[\[\*\?]', key, 1)[0]
+ files = self.get_hook().get_file_metadata(prefix, bucket_name)
+ if len(files) == 0:
+ return False
+
+ # Reduce the set of metadata to size only
+ files = list(map(lambda f: {'Size': f['Size']}, files))
+ else:
+ obj = self.get_hook().head_object(key, bucket_name)
+ if obj is None:
+ return False
+ files = [{'Size': obj['ContentLength']}]
+
+ if self.check_fn is not None:
+ return self.check_fn(files)
+
+ return True
+
+ def poke(self, context: 'Context'):
+ return all(self._check_key(key) for key in self.bucket_key)
def get_hook(self) -> S3Hook:
"""Create and return an S3Hook"""
@@ -111,37 +154,8 @@ class S3KeySensor(BaseSensorOperator):
class S3KeySizeSensor(S3KeySensor):
"""
- Waits for a key (a file-like instance on S3) to be present and be more than
- some size in a S3 bucket.
- S3 being a key/value it does not support folders. The path is just a key
- a resource.
-
- :param bucket_key: The key being waited on. Supports full s3:// style url
- or relative path from root level. When it's specified as a full s3://
- url, please leave bucket_name as `None`.
- :param bucket_name: Name of the S3 bucket. Only needed when ``bucket_key``
- is not provided as a full s3:// url.
- :param wildcard_match: whether the bucket_key should be interpreted as a
- Unix wildcard pattern
- :param aws_conn_id: a reference to the s3 connection
- :param verify: Whether or not to verify SSL certificates for S3 connection.
- By default SSL certificates are verified.
- You can provide the following values:
-
- - ``False``: do not validate SSL certificates. SSL will still be used
- (unless use_ssl is False), but SSL certificates will not be
- verified.
- - ``path/to/cert/bundle.pem``: A filename of the CA cert bundle to uses.
- You can specify this argument if you want to use a different
- CA cert bundle than the one used by botocore.
- :param check_fn: Function that receives the list of the S3 objects,
- and returns the boolean:
- - ``True``: a certain criteria is met
- - ``False``: the criteria isn't met
- **Example**: Wait for any S3 object size more than 1 megabyte ::
-
- def check_fn(self, data: List) -> bool:
- return any(f.get('Size', 0) > 1048576 for f in data if isinstance(f, dict))
+ This class is deprecated.
+ Please use :class:`~airflow.providers.amazon.aws.sensors.s3.S3KeySensor`.
"""
def __init__(
@@ -150,48 +164,26 @@ class S3KeySizeSensor(S3KeySensor):
check_fn: Optional[Callable[..., bool]] = None,
**kwargs,
):
- super().__init__(**kwargs)
- self.check_fn_user = check_fn
-
- def poke(self, context: 'Context'):
- if super().poke(context=context) is False:
- return False
-
- s3_objects = self.get_files(s3_hook=self.get_hook())
- if not s3_objects:
- return False
- if self.check_fn_user is not None:
- return self.check_fn_user(s3_objects)
- return self.check_fn(s3_objects)
-
- def get_files(self, s3_hook: S3Hook, delimiter: Optional[str] = '/') -> List:
- """Gets a list of files in the bucket"""
- prefix = self.bucket_key
- config = {
- 'PageSize': None,
- 'MaxItems': None,
- }
- if self.wildcard_match:
- prefix = re.split(r'[\[\*\?]', self.bucket_key, 1)[0]
+ warnings.warn(
+ """
+ S3PrefixSensor is deprecated.
+ Please use `airflow.providers.amazon.aws.sensors.s3.S3KeySensor`.
+ """,
+ DeprecationWarning,
+ stacklevel=2,
+ )
- paginator = s3_hook.get_conn().get_paginator('list_objects_v2')
- response = paginator.paginate(
- Bucket=self.bucket_name, Prefix=prefix, Delimiter=delimiter, PaginationConfig=config
+ super().__init__(
+ check_fn=check_fn if check_fn is not None else S3KeySizeSensor.default_check_fn, **kwargs
)
- keys: List = []
- for page in response:
- if 'Contents' in page:
- _temp = [k for k in page['Contents'] if isinstance(k.get('Size', None), (int, float))]
- keys = keys + _temp
- return keys
-
- def check_fn(self, data: List, object_min_size: Optional[Union[int, float]] = 0) -> bool:
+
+ @staticmethod
+ def default_check_fn(data: List) -> bool:
"""Default function for checking that S3 Objects have size more than 0
:param data: List of the objects in S3 bucket.
- :param object_min_size: Checks if the objects sizes are greater then this value.
"""
- return all(f.get('Size', 0) > object_min_size for f in data if isinstance(f, dict))
+ return all(f.get('Size', 0) > 0 for f in data)
@poke_mode_only
@@ -332,30 +324,10 @@ class S3KeysUnchangedSensor(BaseSensorOperator):
return self.is_keys_unchanged(set(self.hook.list_keys(self.bucket_name, prefix=self.prefix)))
-class S3PrefixSensor(BaseSensorOperator):
+class S3PrefixSensor(S3KeySensor):
"""
- Waits for a prefix or all prefixes to exist. A prefix is the first part of a key,
- thus enabling checking of constructs similar to glob ``airfl*`` or
- SQL LIKE ``'airfl%'``. There is the possibility to precise a delimiter to
- indicate the hierarchy or keys, meaning that the match will stop at that
- delimiter. Current code accepts sane delimiters, i.e. characters that
- are NOT special characters in the Python regex engine.
-
- :param bucket_name: Name of the S3 bucket
- :param prefix: The prefix being waited on. Relative path from bucket root level.
- :param delimiter: The delimiter intended to show hierarchy.
- Defaults to '/'.
- :param aws_conn_id: a reference to the s3 connection
- :param verify: Whether or not to verify SSL certificates for S3 connection.
- By default SSL certificates are verified.
- You can provide the following values:
-
- - ``False``: do not validate SSL certificates. SSL will still be used
- (unless use_ssl is False), but SSL certificates will not be
- verified.
- - ``path/to/cert/bundle.pem``: A filename of the CA cert bundle to uses.
- You can specify this argument if you want to use a different
- CA cert bundle than the one used by botocore.
+ This class is deprecated.
+ Please use :class:`~airflow.providers.amazon.aws.sensors.s3.S3KeySensor`.
"""
template_fields: Sequence[str] = ('prefix', 'bucket_name')
@@ -363,35 +335,21 @@ class S3PrefixSensor(BaseSensorOperator):
def __init__(
self,
*,
- bucket_name: str,
prefix: Union[str, List[str]],
delimiter: str = '/',
- aws_conn_id: str = 'aws_default',
- verify: Optional[Union[str, bool]] = None,
**kwargs,
):
- super().__init__(**kwargs)
- # Parse
- self.bucket_name = bucket_name
- self.prefix = [prefix] if isinstance(prefix, str) else prefix
- self.delimiter = delimiter
- self.aws_conn_id = aws_conn_id
- self.verify = verify
- self.hook: Optional[S3Hook] = None
-
- def poke(self, context: 'Context'):
- self.log.info('Poking for prefix : %s in bucket s3://%s', self.prefix, self.bucket_name)
- return all(self._check_for_prefix(prefix) for prefix in self.prefix)
-
- def get_hook(self) -> S3Hook:
- """Create and return an S3Hook"""
- if self.hook:
- return self.hook
+ warnings.warn(
+ """
+ S3PrefixSensor is deprecated.
+ Please use `airflow.providers.amazon.aws.sensors.s3.S3KeySensor`.
+ """,
+ DeprecationWarning,
+ stacklevel=2,
+ )
- self.hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
- return self.hook
+ self.prefix = prefix
+ prefixes = [self.prefix] if isinstance(self.prefix, str) else self.prefix
+ keys = [pref if pref.endswith(delimiter) else pref + delimiter for pref in prefixes]
- def _check_for_prefix(self, prefix: str) -> bool:
- return self.get_hook().check_for_prefix(
- prefix=prefix, delimiter=self.delimiter, bucket_name=self.bucket_name
- )
+ super().__init__(bucket_key=keys, **kwargs)
diff --git a/airflow/sensors/s3_prefix_sensor.py b/airflow/sensors/s3_prefix_sensor.py
index 9826dde3a5..d5ea72b37d 100644
--- a/airflow/sensors/s3_prefix_sensor.py
+++ b/airflow/sensors/s3_prefix_sensor.py
@@ -15,14 +15,14 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""This module is deprecated. Please use :mod:`airflow.providers.amazon.aws.sensors.s3_prefix`."""
+"""This module is deprecated. Please use :mod:`airflow.providers.amazon.aws.sensors.s3`."""
import warnings
from airflow.providers.amazon.aws.sensors.s3_prefix import S3PrefixSensor # noqa
warnings.warn(
- "This module is deprecated. Please use `airflow.providers.amazon.aws.sensors.s3_prefix`.",
+ "This module is deprecated. Please use `airflow.providers.amazon.aws.sensors.s3`.",
DeprecationWarning,
stacklevel=2,
)
diff --git a/docs/apache-airflow-providers-amazon/operators/s3.rst b/docs/apache-airflow-providers-amazon/operators/s3.rst
index 82868b6f26..e565197ce7 100644
--- a/docs/apache-airflow-providers-amazon/operators/s3.rst
+++ b/docs/apache-airflow-providers-amazon/operators/s3.rst
@@ -22,9 +22,7 @@ Amazon S3 Operators
Airflow to Amazon Simple Storage Service (S3) integration provides several operators to create and interact with S3 buckets.
- :class:`~airflow.providers.amazon.aws.sensors.s3.S3KeySensor`
- - :class:`~airflow.providers.amazon.aws.sensors.s3.S3KeySizeSensor`
- :class:`~airflow.providers.amazon.aws.sensors.s3.S3KeysUnchangedSensor`
- - :class:`~airflow.providers.amazon.aws.sensors.s3.S3PrefixSensor`
- :class:`~airflow.providers.amazon.aws.operators.s3.S3CreateBucketOperator`
- :class:`~airflow.providers.amazon.aws.operators.s3.S3DeleteBucketOperator`
- :class:`~airflow.providers.amazon.aws.operators.s3.S3DeleteBucketTaggingOperator`
@@ -111,6 +109,61 @@ To delete the tags of an Amazon S3 bucket you can use
:start-after: [START howto_operator_s3_delete_bucket_tagging]
:end-before: [END howto_operator_s3_delete_bucket_tagging]
+.. _howto/sensor:S3KeySensor:
+
+Amazon S3 Key Sensor
+--------------------
+
+To wait for one or multiple keys to be present in an Amazon S3 bucket you can use
+:class:`~airflow.providers.amazon.aws.sensors.s3.S3KeySensor`.
+For each key, it calls
+`head_object <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html#S3.Client.head_object>`__
+API (or `list_objects_v2 <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html#S3.Client.list_objects_v2>`__
+API if ``wildcard_match`` is ``True``) to check whether it is present or not.
+Please keep in mind, especially when used to check a large volume of keys, that it makes one API call per key.
+
+To check one file:
+
+.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_s3_bucket.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_sensor_s3_key_single_key]
+ :end-before: [END howto_sensor_s3_key_single_key]
+
+To check multiple files:
+
+.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_s3_bucket.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_sensor_s3_key_multiple_keys]
+ :end-before: [END howto_sensor_s3_key_multiple_keys]
+
+To check with an additional custom check you can define a function which receives a list of matched S3 object
+attributes and returns a boolean:
+
+- ``True``: a certain criteria is met
+- ``False``: the criteria isn't met
+
+This function is called for each key passed as parameter in ``bucket_key``.
+The reason why the parameter of this function is a list of objects is when ``wildcard_match`` is ``True``,
+multiple files can match one key. The list of matched S3 object attributes contain only the size and is this format:
+
+.. code-block:: python
+
+ [{"Size": int}]
+
+.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_s3_bucket.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_sensor_s3_key_function_definition]
+ :end-before: [END howto_sensor_s3_key_function_definition]
+
+.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_s3_bucket.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_sensor_s3_key_function]
+ :end-before: [END howto_sensor_s3_key_function]
+
Reference
---------
diff --git a/tests/providers/amazon/aws/hooks/test_s3.py b/tests/providers/amazon/aws/hooks/test_s3.py
index a9cbbe09f7..6e889952e4 100644
--- a/tests/providers/amazon/aws/hooks/test_s3.py
+++ b/tests/providers/amazon/aws/hooks/test_s3.py
@@ -178,6 +178,26 @@ class TestAwsS3Hook:
assert sorted(keys) == sorted(hook.list_keys(s3_bucket, delimiter='/', page_size=1))
+ def test_get_file_metadata(self, s3_bucket):
+ hook = S3Hook()
+ bucket = hook.get_bucket(s3_bucket)
+ bucket.put_object(Key='test', Body=b'a')
+
+ assert len(hook.get_file_metadata('t', s3_bucket)) == 1
+ assert hook.get_file_metadata('t', s3_bucket)[0]['Size'] is not None
+ assert len(hook.get_file_metadata('test', s3_bucket)) == 1
+ assert len(hook.get_file_metadata('a', s3_bucket)) == 0
+
+ def test_head_object(self, s3_bucket):
+ hook = S3Hook()
+ bucket = hook.get_bucket(s3_bucket)
+ bucket.put_object(Key='a', Body=b'a')
+
+ assert hook.head_object('a', s3_bucket) is not None
+ assert hook.head_object(f's3://{s3_bucket}//a') is not None
+ assert hook.head_object('b', s3_bucket) is None
+ assert hook.head_object(f's3://{s3_bucket}//b') is None
+
def test_check_for_key(self, s3_bucket):
hook = S3Hook()
bucket = hook.get_bucket(s3_bucket)
diff --git a/tests/providers/amazon/aws/sensors/test_s3_key.py b/tests/providers/amazon/aws/sensors/test_s3_key.py
index 2f6237ebff..1cb2ddc9b9 100644
--- a/tests/providers/amazon/aws/sensors/test_s3_key.py
+++ b/tests/providers/amazon/aws/sensors/test_s3_key.py
@@ -17,6 +17,7 @@
# under the License.
import unittest
+from typing import List
from unittest import mock
import pytest
@@ -40,6 +41,18 @@ class TestS3KeySensor(unittest.TestCase):
with pytest.raises(AirflowException):
op.poke(None)
+ @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.head_object')
+ def test_bucket_name_none_and_bucket_key_is_list_and_contain_relative_path(self, mock_head_object):
+ """
+ Test if exception is raised when bucket_name is None
+ and bucket_key is provided with one of the two keys as relative path rather than s3:// url.
+ :return:
+ """
+ mock_head_object.return_value = {'ContentLength': 0}
+ op = S3KeySensor(task_id='s3_key_sensor', bucket_key=["s3://test_bucket/file", "file_in_bucket"])
+ with pytest.raises(AirflowException):
+ op.poke(None)
+
def test_bucket_name_provided_and_bucket_key_is_s3_url(self):
"""
Test if exception is raised when bucket_name is provided
@@ -52,15 +65,31 @@ class TestS3KeySensor(unittest.TestCase):
with pytest.raises(AirflowException):
op.poke(None)
+ @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.head_object')
+ def test_bucket_name_provided_and_bucket_key_is_list_and_contains_s3_url(self, mock_head_object):
+ """
+ Test if exception is raised when bucket_name is provided
+ while bucket_key contains a full s3:// url.
+ :return:
+ """
+ mock_head_object.return_value = {'ContentLength': 0}
+ op = S3KeySensor(
+ task_id='s3_key_sensor',
+ bucket_key=["test_bucket", "s3://test_bucket/file"],
+ bucket_name='test_bucket',
+ )
+ with pytest.raises(AirflowException):
+ op.poke(None)
+
@parameterized.expand(
[
['s3://bucket/key', None, 'key', 'bucket'],
['key', 'bucket', 'key', 'bucket'],
]
)
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.check_for_key')
- def test_parse_bucket_key(self, key, bucket, parsed_key, parsed_bucket, mock_check):
- mock_check.return_value = False
+ @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.head_object')
+ def test_parse_bucket_key(self, key, bucket, parsed_key, parsed_bucket, mock_head_object):
+ mock_head_object.return_value = None
op = S3KeySensor(
task_id='s3_key_sensor',
@@ -70,12 +99,11 @@ class TestS3KeySensor(unittest.TestCase):
op.poke(None)
- assert op.bucket_key == parsed_key
- assert op.bucket_name == parsed_bucket
+ mock_head_object.assert_called_once_with(parsed_key, parsed_bucket)
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.check_for_key')
- def test_parse_bucket_key_from_jinja(self, mock_check):
- mock_check.return_value = False
+ @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.head_object')
+ def test_parse_bucket_key_from_jinja(self, mock_head_object):
+ mock_head_object.return_value = None
Variable.set("test_bucket_key", "s3://bucket/key")
@@ -96,76 +124,92 @@ class TestS3KeySensor(unittest.TestCase):
ti.render_templates(context)
op.poke(None)
- assert op.bucket_key == "key"
- assert op.bucket_name == "bucket"
+ mock_head_object.assert_called_once_with("key", "bucket")
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.check_for_key')
- def test_poke(self, mock_check):
+ @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.head_object')
+ def test_poke(self, mock_head_object):
op = S3KeySensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file')
- mock_check.return_value = False
+ mock_head_object.return_value = None
assert op.poke(None) is False
- mock_check.assert_called_once_with(op.bucket_key, op.bucket_name)
+ mock_head_object.assert_called_once_with("file", "test_bucket")
- mock_check.return_value = True
+ mock_head_object.return_value = {'ContentLength': 0}
assert op.poke(None) is True
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.check_for_wildcard_key')
- def test_poke_wildcard(self, mock_check):
- op = S3KeySensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file', wildcard_match=True)
+ @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.head_object')
+ def test_poke_multiple_files(self, mock_head_object):
+ op = S3KeySensor(
+ task_id='s3_key_sensor', bucket_key=['s3://test_bucket/file1', 's3://test_bucket/file2']
+ )
- mock_check.return_value = False
+ mock_head_object.side_effect = [{'ContentLength': 0}, None]
assert op.poke(None) is False
- mock_check.assert_called_once_with(op.bucket_key, op.bucket_name)
- mock_check.return_value = True
+ mock_head_object.side_effect = [{'ContentLength': 0}, {'ContentLength': 0}]
assert op.poke(None) is True
+ mock_head_object.assert_any_call("file1", "test_bucket")
+ mock_head_object.assert_any_call("file2", "test_bucket")
-class TestS3KeySizeSensor(unittest.TestCase):
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.check_for_key', return_value=False)
- def test_poke_check_for_key_false(self, mock_check_for_key):
- op = S3KeySizeSensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file')
+ @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.get_file_metadata')
+ def test_poke_wildcard(self, mock_get_file_metadata):
+ op = S3KeySensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file*', wildcard_match=True)
+
+ mock_get_file_metadata.return_value = []
+ assert op.poke(None) is False
+ mock_get_file_metadata.assert_called_once_with("file", "test_bucket")
+
+ mock_get_file_metadata.return_value = [{'Size': 0}]
+ assert op.poke(None) is True
+
+ @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.get_file_metadata')
+ def test_poke_wildcard_multiple_files(self, mock_get_file_metadata):
+ op = S3KeySensor(
+ task_id='s3_key_sensor',
+ bucket_key=['s3://test_bucket/file1*', 's3://test_bucket/file2*'],
+ wildcard_match=True,
+ )
+
+ mock_get_file_metadata.side_effect = [[{'Size': 0}], []]
assert op.poke(None) is False
- mock_check_for_key.assert_called_once_with(op.bucket_key, op.bucket_name)
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3KeySizeSensor.get_files', return_value=[])
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.check_for_key', return_value=True)
- def test_poke_get_files_false(self, mock_check_for_key, mock_get_files):
- op = S3KeySizeSensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file')
+ mock_get_file_metadata.side_effect = [[{'Size': 0}], [{'Size': 0}]]
+ assert op.poke(None) is True
+
+ mock_get_file_metadata.assert_any_call("file1", "test_bucket")
+ mock_get_file_metadata.assert_any_call("file2", "test_bucket")
+
+ @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.head_object')
+ def test_poke_with_check_function(self, mock_head_object):
+ def check_fn(files: List) -> bool:
+ return all(f.get('Size', 0) > 0 for f in files)
+
+ op = S3KeySensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file', check_fn=check_fn)
+
+ mock_head_object.return_value = {'ContentLength': 0}
assert op.poke(None) is False
- mock_check_for_key.assert_called_once_with(op.bucket_key, op.bucket_name)
- mock_get_files.assert_called_once_with(s3_hook=op.get_hook())
+
+ mock_head_object.return_value = {'ContentLength': 1}
+ assert op.poke(None) is True
+
+
+class TestS3KeySizeSensor(unittest.TestCase):
+ def test_deprecation_warnings_generated(self):
+ with pytest.warns(expected_warning=DeprecationWarning):
+ S3KeySizeSensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file')
@parameterized.expand(
[
- [{"Contents": [{"Size": 0}, {"Size": 0}]}, False],
- [{"Contents": [{"Size": 0}]}, False],
- [{"Contents": []}, False],
- [{"Contents": [{"Size": 10}]}, True],
- [{"Contents": [{"Size": 10}, {"Size": 0}]}, False],
- [{"Contents": [{"Size": 10}, {"Size": 10}]}, True],
+ [{"ContentLength": 0}, False],
+ [{"ContentLength": 10}, True],
]
)
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.get_conn')
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.check_for_key')
- def test_poke(self, paginate_return_value, poke_return_value, mock_check, mock_get_conn):
- op = S3KeySizeSensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file')
-
- mock_check.return_value = True
- mock_paginator = mock.Mock()
- mock_paginator.paginate.return_value = []
+ @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.head_object')
+ def test_poke(self, head_object_return_value, poke_return_value, mock_head_object):
+ op = S3KeySizeSensor(
+ task_id='s3_key_sensor', bucket_key=['s3://test_bucket/file', 's3://test_bucket2/file2']
+ )
- mock_get_conn.return_value.get_paginator.return_value = mock_paginator
- mock_paginator.paginate.return_value = [paginate_return_value]
+ mock_head_object.return_value = head_object_return_value
assert op.poke(None) is poke_return_value
- mock_check.assert_called_once_with(op.bucket_key, op.bucket_name)
-
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3KeySizeSensor.get_files', return_value=[])
- @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.check_for_wildcard_key')
- def test_poke_wildcard(self, mock_check, mock_get_files):
- op = S3KeySizeSensor(task_id='s3_key_sensor', bucket_key='s3://test_bucket/file', wildcard_match=True)
-
- mock_check.return_value = False
- assert op.poke(None) is False
- mock_check.assert_called_once_with(op.bucket_key, op.bucket_name)
diff --git a/tests/providers/amazon/aws/sensors/test_s3_prefix.py b/tests/providers/amazon/aws/sensors/test_s3_prefix.py
index d0124343d5..b64afe4520 100644
--- a/tests/providers/amazon/aws/sensors/test_s3_prefix.py
+++ b/tests/providers/amazon/aws/sensors/test_s3_prefix.py
@@ -19,45 +19,50 @@
from unittest import mock
from unittest.mock import call
+import pytest
+
from airflow.providers.amazon.aws.sensors.s3 import S3PrefixSensor
-@mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook')
-def test_poke(mock_hook):
+def test_deprecation_warnings_generated():
+ with pytest.warns(expected_warning=DeprecationWarning):
+ S3PrefixSensor(task_id='s3_prefix', bucket_name='bucket', prefix='prefix')
+
+
+@mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.head_object')
+def test_poke(mock_head_object):
op = S3PrefixSensor(task_id='s3_prefix', bucket_name='bucket', prefix='prefix')
- mock_hook.return_value.check_for_prefix.return_value = False
+ mock_head_object.return_value = None
assert not op.poke({})
- mock_hook.return_value.check_for_prefix.assert_called_once_with(
- prefix='prefix', delimiter='/', bucket_name='bucket'
- )
+ mock_head_object.assert_called_once_with('prefix/', 'bucket')
- mock_hook.return_value.check_for_prefix.return_value = True
+ mock_head_object.return_value = {'ContentLength': 0}
assert op.poke({})
-@mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook')
-def test_poke_should_check_multiple_prefixes(mock_hook):
- op = S3PrefixSensor(task_id='s3_prefix', bucket_name='bucket', prefix=['prefix1', 'prefix2'])
+@mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.head_object')
+def test_poke_should_check_multiple_prefixes(mock_head_object):
+ op = S3PrefixSensor(task_id='s3_prefix', bucket_name='bucket', prefix=['prefix1', 'prefix2/'])
- mock_hook.return_value.check_for_prefix.return_value = False
+ mock_head_object.return_value = None
assert not op.poke({}), "poke returns false when the prefixes do not exist"
- mock_hook.return_value.check_for_prefix.assert_has_calls(
+ mock_head_object.assert_has_calls(
calls=[
- call(prefix='prefix1', delimiter='/', bucket_name='bucket'),
+ call('prefix1/', 'bucket'),
]
)
- mock_hook.return_value.check_for_prefix.side_effect = [True, False]
+ mock_head_object.side_effect = [{'ContentLength': 0}, None]
assert not op.poke({}), "poke returns false when only some of the prefixes exist"
- mock_hook.return_value.check_for_prefix.side_effect = [True, True]
+ mock_head_object.side_effect = [{'ContentLength': 0}, {'ContentLength': 0}]
assert op.poke({}), "poke returns true when both prefixes exist"
- mock_hook.return_value.check_for_prefix.assert_has_calls(
+ mock_head_object.assert_has_calls(
calls=[
- call(prefix='prefix1', delimiter='/', bucket_name='bucket'),
- call(prefix='prefix2', delimiter='/', bucket_name='bucket'),
+ call('prefix1/', 'bucket'),
+ call('prefix2/', 'bucket'),
]
)