From 8caf708f261a30c74d70bb60681a9380bc244e4e Mon Sep 17 00:00:00 2001 From: Shashwati Date: Sun, 26 Apr 2026 20:24:40 +0530 Subject: [PATCH 1/3] Add overwrite parameter to download_mail_attachments to prevent silent file overwriting --- .../src/airflow/providers/imap/hooks/imap.py | 18 ++++++++++--- .../imap/tests/unit/imap/hooks/test_imap.py | 26 +++++++++++++++++++ 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/providers/imap/src/airflow/providers/imap/hooks/imap.py b/providers/imap/src/airflow/providers/imap/hooks/imap.py index b5327e81eade3..8d235f734f071 100644 --- a/providers/imap/src/airflow/providers/imap/hooks/imap.py +++ b/providers/imap/src/airflow/providers/imap/hooks/imap.py @@ -192,6 +192,7 @@ def download_mail_attachments( 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. @@ -210,6 +211,8 @@ def download_mail_attachments( 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 set to True (default), existing files with the same name will be + overwritten. If set to False, a counter suffix is appended to preserve all files. """ mail_attachments = self._retrieve_mails_attachments_by_name( name, check_regex, latest_only, max_mails, mail_folder, mail_filter @@ -218,7 +221,7 @@ def download_mail_attachments( 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"): @@ -286,14 +289,14 @@ def _check_mail_body( 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 _is_symlink(self, name: str) -> bool: # IMPORTANT NOTE: os.path.islink is not working for windows symlinks @@ -310,9 +313,16 @@ def _correct_path(self, name: str, local_output_directory: str) -> str: else local_output_directory + "/" + name ) - def _create_file(self, name: str, payload: Any, local_output_directory: str) -> None: + 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 and os.path.exists(file_path): + base, ext = os.path.splitext(file_path) + counter = 1 + while os.path.exists(f"{base}_{counter}{ext}"): + counter += 1 + file_path = f"{base}_{counter}{ext}" + with open(file_path, "wb") as file: file.write(payload) diff --git a/providers/imap/tests/unit/imap/hooks/test_imap.py b/providers/imap/tests/unit/imap/hooks/test_imap.py index 59e908a949438..8a97341130b33 100644 --- a/providers/imap/tests/unit/imap/hooks/test_imap.py +++ b/providers/imap/tests/unit/imap/hooks/test_imap.py @@ -458,3 +458,29 @@ def test_has_mail_attachment_with_max_mails(self, mock_imaplib): assert result is True assert 1 <= mock_conn.fetch.call_count <= 2 + + @patch("airflow.providers.imap.hooks.imap.os.path.exists", return_value=False) + @patch(open_string, new_callable=mock_open) + @patch(imaplib_string) + def test_download_mail_attachments_overwrite_true_default( + self, mock_imaplib, mock_open_method, mock_exists + ): + _create_fake_imap(mock_imaplib, with_mail=True) + + with ImapHook() as imap_hook: + imap_hook.download_mail_attachments("test1.csv", "test_directory") + + mock_open_method.assert_called_once_with("test_directory/test1.csv", "wb") + + @patch("airflow.providers.imap.hooks.imap.os.path.exists", side_effect=[True, False]) + @patch(open_string, new_callable=mock_open) + @patch(imaplib_string) + def test_download_mail_attachments_no_overwrite_renames_file( + self, mock_imaplib, mock_open_method, mock_exists + ): + _create_fake_imap(mock_imaplib, with_mail=True) + + with ImapHook() as imap_hook: + imap_hook.download_mail_attachments("test1.csv", "test_directory", overwrite=False) + + mock_open_method.assert_called_once_with("test_directory/test1_1.csv", "wb") From cf882836833b216fa534b11d19d18972d2cafd06 Mon Sep 17 00:00:00 2001 From: Shashwati Date: Sun, 26 Apr 2026 20:28:14 +0530 Subject: [PATCH 2/3] Fix test: overwrite=True should not call os.path.exists; rename test for clarity --- providers/imap/tests/unit/imap/hooks/test_imap.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/providers/imap/tests/unit/imap/hooks/test_imap.py b/providers/imap/tests/unit/imap/hooks/test_imap.py index 8a97341130b33..c9b5e943b3b72 100644 --- a/providers/imap/tests/unit/imap/hooks/test_imap.py +++ b/providers/imap/tests/unit/imap/hooks/test_imap.py @@ -459,18 +459,19 @@ def test_has_mail_attachment_with_max_mails(self, mock_imaplib): assert result is True assert 1 <= mock_conn.fetch.call_count <= 2 - @patch("airflow.providers.imap.hooks.imap.os.path.exists", return_value=False) + @patch("airflow.providers.imap.hooks.imap.os.path.exists", return_value=True) @patch(open_string, new_callable=mock_open) @patch(imaplib_string) - def test_download_mail_attachments_overwrite_true_default( + def test_download_mail_attachments_overwrite_true_does_not_rename( self, mock_imaplib, mock_open_method, mock_exists ): _create_fake_imap(mock_imaplib, with_mail=True) with ImapHook() as imap_hook: - imap_hook.download_mail_attachments("test1.csv", "test_directory") + imap_hook.download_mail_attachments("test1.csv", "test_directory", overwrite=True) mock_open_method.assert_called_once_with("test_directory/test1.csv", "wb") + mock_exists.assert_not_called() @patch("airflow.providers.imap.hooks.imap.os.path.exists", side_effect=[True, False]) @patch(open_string, new_callable=mock_open) From ebcd867f6794187264e89b3dbb6182e1c552ab17 Mon Sep 17 00:00:00 2001 From: Shashwati Date: Sun, 26 Apr 2026 22:17:07 +0530 Subject: [PATCH 3/3] Fix line length: reformat _create_files and _create_file signatures --- providers/imap/src/airflow/providers/imap/hooks/imap.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/providers/imap/src/airflow/providers/imap/hooks/imap.py b/providers/imap/src/airflow/providers/imap/hooks/imap.py index 8d235f734f071..2bf7dc90c5fbc 100644 --- a/providers/imap/src/airflow/providers/imap/hooks/imap.py +++ b/providers/imap/src/airflow/providers/imap/hooks/imap.py @@ -289,7 +289,9 @@ def _check_mail_body( 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, overwrite: bool = True) -> 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!") @@ -313,7 +315,9 @@ def _correct_path(self, name: str, local_output_directory: str) -> str: else local_output_directory + "/" + name ) - def _create_file(self, name: str, payload: Any, local_output_directory: str, overwrite: bool = True) -> None: + 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 and os.path.exists(file_path):