Add 'overwrite' option to ImapHook.download_mail_attachments

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)
 

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")