add new S3MultiKeySensor

airflow/providers/amazon/aws/sensors/s3.py

@@ -109,6 +109,68 @@ def get_hook(self) -> S3Hook:
         return self.hook
 
 
+class S3MultiKeySensor(BaseSensorOperator):
+    """
+    Waits for multiple keys (file-like instances on S3) to be present in a single S3 bucket.
+    Does not support multiple buckets or keys in the form of s3:// urls.
+    S3 being a key/value it does not support folders.
+
+    :param bucket_keys: The keys being waited on. Supports relative paths from root level.
+    :param bucket_name: Name of the S3 bucket.
+    :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.
+    """
+
+    template_fields: Sequence[str] = ('bucket_keys', 'bucket_name')
+
+    def __init__(
+        self,
+        *,
+        bucket_keys: List[str],
+        bucket_name: str,
+        aws_conn_id: str = 'aws_default',
+        verify: Optional[Union[str, bool]] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.bucket_name = bucket_name
+        self.bucket_keys = bucket_keys
+        self.aws_conn_id = aws_conn_id
+        self.verify = verify
+        self.hook: Optional[S3Hook] = None
+
+    def get_hook(self) -> S3Hook:
+        """Create and return an S3Hook"""
+        if self.hook:
+            return self.hook
+
+        self.hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
+        return self.hook
+
+    def poke(self, context: 'Context'):
+        self.log.info(f"Poking for keys:\n  { ', '.join(self.bucket_keys) } \nin bucket: {self.bucket_name}")
+
+        results = []
+        s3_hook = self.get_hook()
+        for s3_key in self.bucket_keys:
+            results.append(s3_hook.check_for_key(s3_key, self.bucket_name))
+
+        failed_keys = [self.bucket_keys[ix] for ix, value in enumerate(results) if value is False]
+        if failed_keys:
+            self.log.info(f"Some poke target(s) do not currently exist:\n { ', '.join(failed_keys) }")
+
+        return all(results)
+
+
 class S3KeySizeSensor(S3KeySensor):
     """
     Waits for a key (a file-like instance on S3) to be present and be more than

tests/providers/amazon/aws/sensors/test_s3_key.py

@@ -25,7 +25,7 @@
 from airflow.exceptions import AirflowException
 from airflow.models import DAG, DagRun, TaskInstance
 from airflow.models.variable import Variable
-from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor, S3KeySizeSensor
+from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor, S3KeySizeSensor, S3MultiKeySensor
 from airflow.utils import timezone
 
 
@@ -122,6 +122,53 @@ def test_poke_wildcard(self, mock_check):
         assert op.poke(None) is True
 
 
+class TestS3MultiKeySensor(unittest.TestCase):
+    @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.check_for_key')
+    def test_parse_bucket_name_and_keys_from_jinja(self, mock_check):
+        mock_check.return_value = False
+
+        Variable.set("test_bucket", "bucket")
+        Variable.set("test_key1", "key1")
+        Variable.set("test_key2", "key2")
+
+        execution_date = timezone.datetime(2020, 1, 1)
+
+        dag = DAG("test_s3_multi_key", start_date=execution_date)
+        op = S3MultiKeySensor(
+            task_id='s3_key_multi_key_sensor',
+            bucket_keys=['{{ var.value.test_key1 }}', '{{ var.value.test_key2 }}'],
+            bucket_name='{{ var.value.test_bucket }}',
+            dag=dag,
+        )
+
+        dag_run = DagRun(dag_id=dag.dag_id, execution_date=execution_date, run_id="test")
+        ti = TaskInstance(task=op)
+        ti.dag_run = dag_run
+        context = ti.get_template_context()
+        ti.render_templates(context)
+        op.poke(None)
+
+        assert op.bucket_keys == ["key1", "key2"]
+        assert op.bucket_name == "bucket"
+
+    @mock.patch('airflow.providers.amazon.aws.sensors.s3.S3Hook.check_for_key')
+    def test_poke(self, mock_check):
+        op = S3MultiKeySensor(
+            task_id='s3_multi_key_sensor', bucket_keys=['file_1', 'file_2'], bucket_name='test_bucket'
+        )
+
+        mock_check.side_effect = [False, False]
+        assert op.poke(None) is False
+        calls = [mock.call(op.bucket_keys[0], op.bucket_name), mock.call(op.bucket_keys[1], op.bucket_name)]
+        mock_check.assert_has_calls(calls)
+
+        mock_check.side_effect = [True, False]
+        assert op.poke(None) is False
+
+        mock_check.side_effect = [True, True]
+        assert op.poke(None) is True
+
+
 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):