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

 ##########
 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'):
 
 Review comment:
   You are likely the heaviest user of this, so we'll just merge this and make 
it 2.0 only

----------------------------------------------------------------
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:
us...@infra.apache.org


With regards,
Apache Git Services

Reply via email to