shahar1 commented on code in PR #60963:
URL: https://github.com/apache/airflow/pull/60963#discussion_r2725620313
##########
providers/imap/src/airflow/providers/imap/hooks/imap.py:
##########
@@ -163,9 +173,10 @@ def retrieve_mail_attachments(
if set to 'ignore' it won't notify you at all.
:returns: a list of tuple each containing the attachment filename and
its payload.
"""
- mail_attachments = self._retrieve_mails_attachments_by_name(
- name, check_regex, latest_only, mail_folder, mail_filter
- )
+ with self:
Review Comment:
Same comment as before
##########
providers/imap/src/airflow/providers/imap/hooks/imap.py:
##########
@@ -200,9 +213,10 @@ def download_mail_attachments(
if set to 'warn' it will only print a warning and
if set to 'ignore' it won't notify you at all.
"""
- mail_attachments = self._retrieve_mails_attachments_by_name(
- name, check_regex, latest_only, mail_folder, mail_filter
- )
+ with self:
Review Comment:
Same comment as before
##########
providers/imap/src/airflow/providers/imap/hooks/imap.py:
##########
@@ -153,6 +162,7 @@ def retrieve_mail_attachments(
:param name: The name of the attachment that will be downloaded.
:param check_regex: Checks the name for a regular expression.
:param latest_only: If set to True it will only retrieve the first
matched attachment.
+ :param max_mails: Maximum number of latest emails to process. Must be
a positive integer. Defaults to None.
Review Comment:
```suggestion
:param max_mails: Maximum number of latest emails to process. Must
be a positive integer, or None when unlimited. Defaults to None.
```
##########
providers/imap/src/airflow/providers/imap/hooks/imap.py:
##########
@@ -132,17 +139,19 @@ def has_mail_attachment(
See :py:meth:`imaplib.IMAP4.search` for details.
:returns: True if there is an attachment with the given name and False
if not.
"""
- mail_attachments = self._retrieve_mails_attachments_by_name(
- name, check_regex, True, mail_folder, mail_filter
- )
- return bool(mail_attachments)
+ with self:
Review Comment:
What is the purpose of wrapping the logic in a context block?
Using `with self:` forces `__enter__()`/`__exit__()` to run on every method
call. In this hook, `__exit__()` calls `logout()`, so each call closes the IMAP
connection. This breaks the expected usage where the caller manages the hook
lifecycle (e.g. `with ImapHook() as hook:` in `class ImapAttachmentSensor`) and
can invoke multiple methods on the same hook instance, because the connection
gets closed unexpectedly after the first call.
##########
providers/imap/src/airflow/providers/imap/hooks/imap.py:
##########
@@ -218,16 +232,25 @@ def _handle_not_found_mode(self, not_found_mode: str) ->
None:
self.log.warning("No mail attachments found!")
def _retrieve_mails_attachments_by_name(
- self, name: str, check_regex: bool, latest_only: bool, mail_folder:
str, mail_filter: str
- ) -> list:
+ self,
+ name: str,
+ check_regex: bool,
+ latest_only: bool,
+ max_mails: int | None,
+ mail_folder: str,
+ mail_filter: str,
+ ) -> list[tuple]:
if not self.mail_client:
raise RuntimeError("The 'mail_client' should be initialized
before!")
- all_matching_attachments = []
+ if max_mails is not None and max_mails <= 0:
+ raise AirflowException("max_mails must be a positive integer")
Review Comment:
Following a recent decision (see [dev list
thread](https://lists.apache.org/thread/5rv4tz0oc27bgr4khx0on0jz8fpxvh55)), the
directive is not to use `AirflowException`, but instead use native
Python/Provider-native exceptions. In this case, it should become something
like:
```suggestion
raise ValueError("max_mails must be a positive integer")
```
+ Fix related unit test
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
