ImapAttachmentToS3Operator: fix it, update sample dag and update doc

airflow/providers/amazon/aws/example_dags/example_imap_attachment_to_s3.py

@@ -26,12 +26,11 @@
 from airflow import DAG
 from airflow.providers.amazon.aws.transfers.imap_attachment_to_s3 import ImapAttachmentToS3Operator
 
-# [START howto_operator_imap_attachment_to_s3_env_variables]
 IMAP_ATTACHMENT_NAME = getenv("IMAP_ATTACHMENT_NAME", "test.txt")
 IMAP_MAIL_FOLDER = getenv("IMAP_MAIL_FOLDER", "INBOX")
 IMAP_MAIL_FILTER = getenv("IMAP_MAIL_FILTER", "All")
-S3_DESTINATION_KEY = getenv("S3_DESTINATION_KEY", "s3://bucket/key.json")
-# [END howto_operator_imap_attachment_to_s3_env_variables]
+S3_BUCKET = getenv("S3_BUCKET", "test-bucket")
+S3_KEY = getenv("S3_KEY", "key")
 
 with DAG(
     dag_id="example_imap_attachment_to_s3",
@@ -40,12 +39,13 @@
     catchup=False,
     tags=['example'],
 ) as dag:
-    # [START howto_operator_imap_attachment_to_s3_task_1]
+    # [START howto_transfer_imap_attachment_to_s3]
     task_transfer_imap_attachment_to_s3 = ImapAttachmentToS3Operator(
+        task_id='transfer_imap_attachment_to_s3',
         imap_attachment_name=IMAP_ATTACHMENT_NAME,
-        s3_key=S3_DESTINATION_KEY,
+        s3_bucket=S3_BUCKET,
+        s3_key=S3_KEY,
         imap_mail_folder=IMAP_MAIL_FOLDER,
         imap_mail_filter=IMAP_MAIL_FILTER,
-        task_id='transfer_imap_attachment_to_s3',
     )
-    # [END howto_operator_imap_attachment_to_s3_task_1]
+    # [END howto_transfer_imap_attachment_to_s3]

airflow/providers/amazon/aws/transfers/imap_attachment_to_s3.py

@@ -16,7 +16,8 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module allows you to transfer mail attachments from a mail server into s3 bucket."""
-from typing import TYPE_CHECKING, Sequence
+import warnings
+from typing import TYPE_CHECKING, Optional, Sequence
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
@@ -25,6 +26,10 @@
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
+_DEPRECATION_MSG = (
+    "The s3_conn_id parameter has been deprecated. You should pass instead the aws_conn_id parameter."
+)
+
 
 class ImapAttachmentToS3Operator(BaseOperator):
     """
@@ -35,14 +40,15 @@ class ImapAttachmentToS3Operator(BaseOperator):
         :ref:`howto/operator:ImapAttachmentToS3Operator`
 
     :param imap_attachment_name: The file name of the mail attachment that you want to transfer.
+    :param s3_bucket: The targeted s3 bucket. This is the S3 bucket where the file will be downloaded.
     :param s3_key: The destination file name in the s3 bucket for the attachment.
     :param imap_check_regex: If set checks the `imap_attachment_name` for a regular expression.
     :param imap_mail_folder: The folder on the mail server to look for the attachment.
     :param imap_mail_filter: If set other than 'All' only specific mails will be checked.
         See :py:meth:`imaplib.IMAP4.search` for details.
     :param s3_overwrite: If set overwrites the s3 key if already exists.
     :param imap_conn_id: The reference to the connection details of the mail server.
-    :param s3_conn_id: The reference to the s3 connection details.
+    :param aws_conn_id: AWS connection to use.
     """
 
     template_fields: Sequence[str] = ('imap_attachment_name', 's3_key', 'imap_mail_filter')
@@ -51,24 +57,31 @@ def __init__(
         self,
         *,
         imap_attachment_name: str,
+        s3_bucket: str,
         s3_key: str,
         imap_check_regex: bool = False,
         imap_mail_folder: str = 'INBOX',
         imap_mail_filter: str = 'All',
         s3_overwrite: bool = False,
         imap_conn_id: str = 'imap_default',
-        s3_conn_id: str = 'aws_default',
+        s3_conn_id: Optional[str] = None,
+        aws_conn_id: str = 'aws_default',
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
+        if s3_conn_id:
+            warnings.warn(_DEPRECATION_MSG, DeprecationWarning, stacklevel=3)
+            aws_conn_id = s3_conn_id
+
         self.imap_attachment_name = imap_attachment_name
+        self.s3_bucket = s3_bucket
         self.s3_key = s3_key
         self.imap_check_regex = imap_check_regex
         self.imap_mail_folder = imap_mail_folder
         self.imap_mail_filter = imap_mail_filter
         self.s3_overwrite = s3_overwrite
         self.imap_conn_id = imap_conn_id
-        self.s3_conn_id = s3_conn_id
+        self.aws_conn_id = aws_conn_id
 
     def execute(self, context: 'Context') -> None:
         """
@@ -91,5 +104,10 @@ def execute(self, context: 'Context') -> None:
                 mail_filter=self.imap_mail_filter,
             )
 
-        s3_hook = S3Hook(aws_conn_id=self.s3_conn_id)
-        s3_hook.load_bytes(bytes_data=imap_mail_attachments[0][1], key=self.s3_key, replace=self.s3_overwrite)
+        s3_hook = S3Hook(aws_conn_id=self.aws_conn_id)
+        s3_hook.load_bytes(
+            bytes_data=imap_mail_attachments[0][1],
+            bucket_name=self.s3_bucket,
+            key=self.s3_key,
+            replace=self.s3_overwrite,
+        )

airflow/providers/amazon/provider.yaml

@@ -463,7 +463,7 @@ transfers:
     python-module: airflow.providers.amazon.aws.transfers.hive_to_dynamodb
   - source-integration-name: Internet Message Access Protocol (IMAP)
     target-integration-name: Amazon Simple Storage Service (S3)
-    how-to-guide: /docs/apache-airflow-providers-amazon/operators/imap_attachment_to_s3.rst
+    how-to-guide: /docs/apache-airflow-providers-amazon/operators/transfer/imap_attachment_to_s3.rst
     python-module: airflow.providers.amazon.aws.transfers.imap_attachment_to_s3
   - source-integration-name: MongoDB
     target-integration-name: Amazon Simple Storage Service (S3)

docs/apache-airflow-providers-amazon/operators/transfer/imap_attachment_to_s3.rst

@@ -0,0 +1,44 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+ ..   http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+Imap Attachment To Amazon S3 Operator
+=====================================
+
+The ``ImapAttachmentToS3Operator`` transfers an email attachment via IMAP
+protocol from a mail server to an Amazon S3 Bucket.
+
+Prerequisite Tasks
+^^^^^^^^^^^^^^^^^^
+
+.. include:: ../_partials/prerequisite_tasks.rst
+
+.. _howto/operator:ImapAttachmentToS3Operator:
+
+Imap Attachment To Amazon S3
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_imap_attachment_to_s3.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_transfer_imap_attachment_to_s3]
+    :end-before: [END howto_transfer_imap_attachment_to_s3]
+
+Reference
+^^^^^^^^^
+
+* `IMAP Library Documentation <https://docs.python.org/3/library/imaplib.html>`__
+* `AWS boto3 Library Documentation for S3 <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html>`__

docs/apache-airflow-providers-amazon/redirects.txt

@@ -16,3 +16,4 @@
 # under the License.
 
 operators/s3_to_redshift.rst operators/transfer/s3_to_redshift.rst
+operators/imap_attachment_to_s3.rst operators/transfer/imap_attachment_to_s3.rst

tests/providers/amazon/aws/transfers/test_imap_attachment_to_s3.py

@@ -26,6 +26,7 @@ class TestImapAttachmentToS3Operator(unittest.TestCase):
     def setUp(self):
         self.kwargs = dict(
             imap_attachment_name='test_file',
+            s3_bucket='test_bucket',
             s3_key='test_file',
             imap_check_regex=False,
             imap_mail_folder='INBOX',
@@ -52,6 +53,7 @@ def test_execute(self, mock_imap_hook, mock_s3_hook):
         )
         mock_s3_hook.return_value.load_bytes.assert_called_once_with(
             bytes_data=mock_imap_hook.return_value.retrieve_mail_attachments.return_value[0][1],
+            bucket_name=self.kwargs['s3_bucket'],
             key=self.kwargs['s3_key'],
             replace=self.kwargs['s3_overwrite'],
         )

tests/providers/amazon/aws/transfers/test_imap_attachment_to_s3_system.py

@@ -17,21 +17,18 @@
 # under the License.
 import pytest
 
-from airflow.providers.amazon.aws.example_dags.example_imap_attachment_to_s3 import S3_DESTINATION_KEY
-from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.aws.example_dags.example_imap_attachment_to_s3 import S3_BUCKET
 from tests.test_utils.amazon_system_helpers import (
     AWS_DAG_FOLDER,
     AmazonSystemTest,
     provide_aws_context,
     provide_aws_s3_bucket,
 )
 
-BUCKET, _ = S3Hook.parse_s3_url(S3_DESTINATION_KEY)
-
 
 @pytest.fixture
 def provide_s3_bucket():
-    with provide_aws_s3_bucket(BUCKET):
+    with provide_aws_s3_bucket(S3_BUCKET):
         yield