[GitHub] [airflow] feluelle commented on a change in pull request #5672: [AIRFLOW-5056] Add argument to filter mails in ImapHook and related operators

[GitBox <gi...@apache.org>]

feluelle commented on a change in pull request #5672: [AIRFLOW-5056] Add argument to filter mails in ImapHook and related operators
URL: https://github.com/apache/airflow/pull/5672#discussion_r312924849
 
 

 ##########
 File path: airflow/contrib/hooks/imap_hook.py
 ##########
 @@ -30,72 +34,105 @@ class ImapHook(BaseHook):
     """
     This hook connects to a mail server by using the imap protocol.
 
+    .. note:: Please call this Hook as context manager via `with`
+        to automatically open and close the connection to the mail server.
+
     :param imap_conn_id: The connection id that contains the information used to authenticate the client.
     :type imap_conn_id: str
     """
 
     def __init__(self, imap_conn_id='imap_default'):
         super().__init__(imap_conn_id)
-        self.conn = self.get_connection(imap_conn_id)
-        self.mail_client = imaplib.IMAP4_SSL(self.conn.host)
+        self.imap_conn_id = imap_conn_id
+        self.mail_client = None
 
     def __enter__(self):
-        self.mail_client.login(self.conn.login, self.conn.password)
-        return self
+        return self.get_conn()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.mail_client.logout()
 
-    def has_mail_attachment(self, name, mail_folder='INBOX', check_regex=False):
+    def get_conn(self):
+        """
+        Login to the mail server.
+
+        .. note:: Please call this Hook as context manager via `with`
+            to automatically open and close the connection to the mail server.
+
+        :return: an authorized ImapHook object.
+        :rtype: ImapHook
+        """
+
+        if not self.mail_client:
+            conn = self.get_connection(self.imap_conn_id)
+            self.mail_client = imaplib.IMAP4_SSL(conn.host)
+            self.mail_client.login(conn.login, conn.password)
+
+        return self
+
+    def has_mail_attachment(self, name, check_regex=False, mail_folder='INBOX', mail_filter='All'):
         """
         Checks the mail folder for mails containing attachments with the given name.
 
         :param name: The name of the attachment that will be searched for.
         :type name: str
-        :param mail_folder: The mail folder where to look at.
-        :type mail_folder: str
         :param check_regex: Checks the name for a regular expression.
         :type check_regex: bool
+        :param mail_folder: The mail folder where to look at.
+        :type mail_folder: str
+        :param mail_filter: If set other than 'All' only specific mails will be checked.
+
+            .. seealso::
+                https://docs.python.org/3.6/library/imaplib.html#imaplib.IMAP4.search
+        :type mail_filter: str
         :returns: True if there is an attachment with the given name and False if not.
         :rtype: bool
         """
         mail_attachments = self._retrieve_mails_attachments_by_name(name,
-                                                                    mail_folder,
                                                                     check_regex,
-                                                                    latest_only=True)
 
 Review comment:
   No, I just passed `True` instead of `latest_only=True` to keep the same order.

----------------------------------------------------------------
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


With regards,
Apache Git Services