This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5fb8637648b Add overwrite parameter to IMAP hook to prevent silent
file overwrites (#68838)
5fb8637648b is described below
commit 5fb8637648badb662d4dd812b1427dea2a996a1a
Author: Dan Kalenga <[email protected]>
AuthorDate: Mon Jul 6 18:02:45 2026 +0200
Add overwrite parameter to IMAP hook to prevent silent file overwrites
(#68838)
This adds an overwrittten paramter ( default to True, preserving the
current behavior). When set to False, duplicate filenames are renamed with a
numeric suffix( Example. report_1.xlsx , report_2.xlsx) instead of overwriting).
---
.../imap/src/airflow/providers/imap/hooks/imap.py | 32 ++++++++++++++++------
providers/imap/tests/unit/imap/hooks/test_imap.py | 27 ++++++++++++++++++
2 files changed, 50 insertions(+), 9 deletions(-)
diff --git a/providers/imap/src/airflow/providers/imap/hooks/imap.py
b/providers/imap/src/airflow/providers/imap/hooks/imap.py
index ecb99eb2061..60b2755ebd7 100644
--- a/providers/imap/src/airflow/providers/imap/hooks/imap.py
+++ b/providers/imap/src/airflow/providers/imap/hooks/imap.py
@@ -193,6 +193,7 @@ class ImapHook(BaseHook):
mail_folder: str = "INBOX",
mail_filter: str = "All",
not_found_mode: str = "raise",
+ overwrite: bool = True,
) -> None:
"""
Download mail's attachments in the mail folder by its name to the
local directory.
@@ -211,6 +212,8 @@ class ImapHook(BaseHook):
If it is set to 'raise' it will raise an exception,
if set to 'warn' it will only print a warning and
if set to 'ignore' it won't notify you at all.
+ :param overwrite: If True (default), existing files are overwritten.
If False, a unique
+ suffix is added to avoid overwriting.
"""
mail_attachments = self._retrieve_mails_attachments_by_name(
name, check_regex, latest_only, max_mails, mail_folder, mail_filter
@@ -219,7 +222,7 @@ class ImapHook(BaseHook):
if not mail_attachments:
self._handle_not_found_mode(not_found_mode)
- self._create_files(mail_attachments, local_output_directory)
+ self._create_files(mail_attachments, local_output_directory,
overwrite=overwrite)
def _handle_not_found_mode(self, not_found_mode: str) -> None:
if not_found_mode not in ("raise", "warn", "ignore"):
@@ -287,14 +290,31 @@ class ImapHook(BaseHook):
return mail.get_attachments_by_name(name, check_regex,
find_first=latest_only)
return []
- def _create_files(self, mail_attachments: list, local_output_directory:
str) -> None:
+ def _create_files(
+ self, mail_attachments: list, local_output_directory: str, overwrite:
bool = True
+ ) -> None:
for name, payload in mail_attachments:
if self._is_symlink(name):
self.log.error("Can not create file because it is a symlink!")
elif self._is_escaping_current_directory(name):
self.log.error("Can not create file because it is escaping the
current directory!")
else:
- self._create_file(name, payload, local_output_directory)
+ self._create_file(name, payload, local_output_directory,
overwrite=overwrite)
+
+ def _create_file(
+ self, name: str, payload: Any, local_output_directory: str, overwrite:
bool = True
+ ) -> None:
+ file_path = self._correct_path(name, local_output_directory)
+
+ if not overwrite:
+ base, ext = os.path.splitext(file_path)
+ counter = 1
+ while os.path.exists(file_path):
+ file_path = f"{base}_{counter}{ext}"
+ counter += 1
+
+ with open(file_path, "wb") as file:
+ file.write(payload)
def _is_symlink(self, name: str) -> bool:
# IMPORTANT NOTE: os.path.islink is not working for windows symlinks
@@ -311,12 +331,6 @@ class ImapHook(BaseHook):
else local_output_directory + os.sep + name
)
- def _create_file(self, name: str, payload: Any, local_output_directory:
str) -> None:
- file_path = self._correct_path(name, local_output_directory)
-
- with open(file_path, "wb") as file:
- file.write(payload)
-
class Mail(LoggingMixin):
"""
diff --git a/providers/imap/tests/unit/imap/hooks/test_imap.py
b/providers/imap/tests/unit/imap/hooks/test_imap.py
index 54537aa29af..0b8a7fa2f80 100644
--- a/providers/imap/tests/unit/imap/hooks/test_imap.py
+++ b/providers/imap/tests/unit/imap/hooks/test_imap.py
@@ -316,6 +316,33 @@ class TestImapHook:
mock_open_method.assert_called_once_with("test_directory/test1.csv",
"wb")
mock_open_method.return_value.write.assert_called_once_with(b"SWQsTmFtZQoxLEZlbGl4")
+ @patch("airflow.providers.imap.hooks.imap.os.path.exists")
+ @patch(open_string, new_callable=mock_open)
+ @patch(imaplib_string)
+ def test_download_mail_attachments_overwrite_true(self, mock_imaplib,
mock_open_method, mock_exists):
+ _create_fake_imap(mock_imaplib, with_mail=True)
+ mock_exists.return_value = True # pretend file already exists
+
+ with ImapHook() as imap_hook:
+ imap_hook.download_mail_attachments("test1.csv", "test_directory",
overwrite=True)
+
+ # with overwrite=True, it should still write to the original filename
+ mock_open_method.assert_called_once_with("test_directory/test1.csv",
"wb")
+
+ @patch("airflow.providers.imap.hooks.imap.os.path.exists")
+ @patch(open_string, new_callable=mock_open)
+ @patch(imaplib_string)
+ def test_download_mail_attachments_overwrite_false(self, mock_imaplib,
mock_open_method, mock_exists):
+ _create_fake_imap(mock_imaplib, with_mail=True)
+ # first call (checking original) returns True, second call (checking
_1) returns False
+ mock_exists.side_effect = [True, False]
+
+ with ImapHook() as imap_hook:
+ imap_hook.download_mail_attachments("test1.csv", "test_directory",
overwrite=False)
+
+ # with overwrite=False, it should write to a renamed file
+ mock_open_method.assert_called_once_with("test_directory/test1_1.csv",
"wb")
+
@patch(open_string, new_callable=mock_open)
@patch(imaplib_string)
def test_download_mail_attachments_not_found(self, mock_imaplib,
mock_open_method):
