diff --git a/airflow/providers/imap/hooks/imap.py b/airflow/providers/imap/hooks/imap.py index 64e07ad444d6e..53bb190e6f0e9 100644 --- a/airflow/providers/imap/hooks/imap.py +++ b/airflow/providers/imap/hooks/imap.py @@ -24,10 +24,11 @@ import imaplib import os import re -from typing import Any, Iterable, List, Optional, Tuple +from typing import Any, Iterable, List, Optional, Tuple, Union from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook +from airflow.models.connection import Connection from airflow.utils.log.logging_mixin import LoggingMixin @@ -51,7 +52,7 @@ class ImapHook(BaseHook): def __init__(self, imap_conn_id: str = default_conn_name) -> None: super().__init__() self.imap_conn_id = imap_conn_id - self.mail_client: Optional[imaplib.IMAP4_SSL] = None + self.mail_client: Optional[Union[imaplib.IMAP4_SSL, imaplib.IMAP4]] = None def __enter__(self) -> 'ImapHook': return self.get_conn() @@ -71,14 +72,24 @@ def get_conn(self) -> 'ImapHook': """ if not self.mail_client: conn = self.get_connection(self.imap_conn_id) - if conn.port: - self.mail_client = imaplib.IMAP4_SSL(conn.host, conn.port) - else: - self.mail_client = imaplib.IMAP4_SSL(conn.host) + self.mail_client = self._build_client(conn) self.mail_client.login(conn.login, conn.password) return self + def _build_client(self, conn: Connection) -> Union[imaplib.IMAP4_SSL, imaplib.IMAP4]: + if conn.extra_dejson.get('use_ssl', True): + IMAP = imaplib.IMAP4_SSL + else: + IMAP = imaplib.IMAP4 + + if conn.port: + mail_client = IMAP(conn.host, conn.port) + else: + mail_client = IMAP(conn.host) + + return mail_client + def has_mail_attachment( self, name: str, *, check_regex: bool = False, mail_folder: str = 'INBOX', mail_filter: str = 'All' ) -> bool: diff --git a/docs/apache-airflow-providers-imap/connections/imap.rst b/docs/apache-airflow-providers-imap/connections/imap.rst index e06c483169f24..18f303be73caf 100644 --- a/docs/apache-airflow-providers-imap/connections/imap.rst +++ b/docs/apache-airflow-providers-imap/connections/imap.rst @@ -49,7 +49,12 @@ Host Specify the IMAP host url. Port - Specify the IMAP port to connect to. + Specify the IMAP port to connect to. The default depends on the whether you use ssl or not. + +Extra (optional) + Specify the extra parameters (as json dictionary) + + * ``use_ssl``: If set to false, then a non-ssl connection is being used. Default is true. Also note that changing the ssl option also influences the default port being used. When specifying the connection in environment variable you should specify it using URI syntax. @@ -60,4 +65,12 @@ For example: .. code-block:: bash - export AIRFLOW_CONN_IMAP_DEFAULT='imap://username:password@myimap.com:993' + export AIRFLOW_CONN_IMAP_DEFAULT='imap://username:password@myimap.com:993?use_ssl=true' + +Another example for connecting via a non-SSL connection. + +.. code-block:: bash + + export AIRFLOW_CONN_IMAP_NONSSL='imap://username:password@myimap.com:143?use_ssl=false' + +Note that you can set the port regardless of whether you choose to use ssl or not. The above examples show default ports for SSL and Non-SSL connections. diff --git a/tests/providers/imap/hooks/test_imap.py b/tests/providers/imap/hooks/test_imap.py index 784edf2a71463..d78b9de5d91db 100644 --- a/tests/providers/imap/hooks/test_imap.py +++ b/tests/providers/imap/hooks/test_imap.py @@ -15,8 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - import imaplib +import json import unittest from unittest.mock import Mock, mock_open, patch @@ -31,9 +31,13 @@ open_string = 'airflow.providers.imap.hooks.imap.open' -def _create_fake_imap(mock_imaplib, with_mail=False, attachment_name='test1.csv'): - mock_conn = Mock(spec=imaplib.IMAP4_SSL) - mock_imaplib.IMAP4_SSL.return_value = mock_conn +def _create_fake_imap(mock_imaplib, with_mail=False, attachment_name='test1.csv', use_ssl=True): + if use_ssl: + mock_conn = Mock(spec=imaplib.IMAP4_SSL) + mock_imaplib.IMAP4_SSL.return_value = mock_conn + else: + mock_conn = Mock(spec=imaplib.IMAP4) + mock_imaplib.IMAP4.return_value = mock_conn mock_conn.login.return_value = ('OK', []) @@ -63,8 +67,19 @@ def setUp(self): conn_type='imap', host='imap_server_address', login='imap_user', + password='imap_password', port=1993, + ) + ) + db.merge_conn( + Connection( + conn_id='imap_nonssl', + conn_type='imap', + host='imap_server_address', + login='imap_user', password='imap_password', + port=1143, + extra=json.dumps(dict(use_ssl=False)), ) ) @@ -79,6 +94,17 @@ def test_connect_and_disconnect(self, mock_imaplib): mock_conn.login.assert_called_once_with('imap_user', 'imap_password') assert mock_conn.logout.call_count == 1 + @patch(imaplib_string) + def test_connect_and_disconnect_via_nonssl(self, mock_imaplib): + mock_conn = _create_fake_imap(mock_imaplib, use_ssl=False) + + with ImapHook(imap_conn_id='imap_nonssl'): + pass + + mock_imaplib.IMAP4.assert_called_once_with('imap_server_address', 1143) + mock_conn.login.assert_called_once_with('imap_user', 'imap_password') + assert mock_conn.logout.call_count == 1 + @patch(imaplib_string) def test_has_mail_attachment_found(self, mock_imaplib): _create_fake_imap(mock_imaplib, with_mail=True)