weldpua2008 commented on a change in pull request #13049:
URL: https://github.com/apache/airflow/pull/13049#discussion_r550321992
##########
File path: airflow/providers/amazon/aws/sensors/s3_key.py
##########
@@ -106,3 +107,91 @@ def get_hook(self) -> S3Hook:
self.hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
return self.hook
+
+
+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`.
+ :type bucket_key: str
+ :param bucket_name: Name of the S3 bucket. Only needed when ``bucket_key``
+ is not provided as a full s3:// url.
+ :type bucket_name: str
+ :param wildcard_match: whether the bucket_key should be interpreted as a
+ Unix wildcard pattern
+ :type wildcard_match: bool
+ :param aws_conn_id: a reference to the s3 connection
+ :type aws_conn_id: str
+ :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.
+ :type verify: bool or str
+ :type check_fn: Optional[Callable[..., bool]]
+ :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))
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ check_fn: Optional[Callable[..., bool]] = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.check_fn_user = check_fn
+
+ def poke(self, context):
+ if super().poke(context=context) is False:
+ return False
+
+ s3_objects = self.get_files(s3_hook=self.get_hook())
+ check_fn = self.check_fn if self.check_fn_user is None else
self.check_fn_user
+ if not s3_objects:
+ return False
+ return check_fn(s3_objects)
+
+ def get_files(self, s3_hook: S3Hook) -> List:
+ """Gets a list of files in the bucket"""
+ prefix = self.bucket_key
+ delimiter = '/'
+ config = {
+ 'PageSize': None,
+ 'MaxItems': None,
+ }
+ if self.wildcard_match:
+ prefix = re.split(r'[*]', self.bucket_key, 1)[0]
+
+ paginator = s3_hook.get_conn().get_paginator('list_objects_v2')
+ response = paginator.paginate(
+ Bucket=self.bucket_name, Prefix=prefix, Delimiter=delimiter,
PaginationConfig=config
+ )
+ keys = []
+ 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) -> bool:
Review comment:
Added
##########
File path: airflow/providers/amazon/aws/sensors/s3_key.py
##########
@@ -106,3 +107,91 @@ def get_hook(self) -> S3Hook:
self.hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
return self.hook
+
+
+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`.
+ :type bucket_key: str
+ :param bucket_name: Name of the S3 bucket. Only needed when ``bucket_key``
+ is not provided as a full s3:// url.
+ :type bucket_name: str
+ :param wildcard_match: whether the bucket_key should be interpreted as a
+ Unix wildcard pattern
+ :type wildcard_match: bool
+ :param aws_conn_id: a reference to the s3 connection
+ :type aws_conn_id: str
+ :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.
+ :type verify: bool or str
+ :type check_fn: Optional[Callable[..., bool]]
+ :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))
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ check_fn: Optional[Callable[..., bool]] = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.check_fn_user = check_fn
+
+ def poke(self, context):
+ if super().poke(context=context) is False:
+ return False
+
+ s3_objects = self.get_files(s3_hook=self.get_hook())
+ check_fn = self.check_fn if self.check_fn_user is None else
self.check_fn_user
+ if not s3_objects:
+ return False
+ return check_fn(s3_objects)
+
+ def get_files(self, s3_hook: S3Hook) -> List:
+ """Gets a list of files in the bucket"""
+ prefix = self.bucket_key
+ delimiter = '/'
+ config = {
+ 'PageSize': None,
+ 'MaxItems': None,
+ }
+ if self.wildcard_match:
+ prefix = re.split(r'[*]', self.bucket_key, 1)[0]
+
+ paginator = s3_hook.get_conn().get_paginator('list_objects_v2')
+ response = paginator.paginate(
+ Bucket=self.bucket_name, Prefix=prefix, Delimiter=delimiter,
PaginationConfig=config
Review comment:
You are right, Thanks!
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
