You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2020/12/27 16:07:30 UTC
[GitHub] [airflow] netazvi commented on a change in pull request #13049: Add S3KeySizeSensor
netazvi commented on a change in pull request #13049:
URL: https://github.com/apache/airflow/pull/13049#discussion_r549128853
##########
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:
the delimiter can be defined here directly instead of an additional parameter
##########
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:
consider changing the check_fn to receive the size as a parameter (with default 0) and then it will be easier for users to use this function as is
----------------------------------------------------------------
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:
users@infra.apache.org