diff --git a/providers/imap/src/airflow/providers/imap/hooks/imap.py b/providers/imap/src/airflow/providers/imap/hooks/imap.py index b5327e81eade3..2bf7dc90c5fbc 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,16 @@ 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 +315,18 @@ 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..c9b5e943b3b72 100644 --- a/providers/imap/tests/unit/imap/hooks/test_imap.py +++ b/providers/imap/tests/unit/imap/hooks/test_imap.py @@ -458,3 +458,30 @@ 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=True) + @patch(open_string, new_callable=mock_open) + @patch(imaplib_string) + 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", 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) + @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")