shahar1 commented on code in PR #66850:
URL: https://github.com/apache/airflow/pull/66850#discussion_r3344564677
##########
providers/imap/src/airflow/providers/imap/hooks/imap.py:
##########
@@ -287,14 +292,54 @@ 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:
+ seen: dict[str, int] = {}
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)
+ if not overwrite:
+ file_name = self._unique_name(name, seen,
local_output_directory)
+ file_path = self._correct_path(file_name,
local_output_directory)
+ while True:
+ try:
+ with open(file_path, "xb") as f:
+ f.write(payload)
+ break
+ except FileExistsError:
+ seen[file_name] = 0
+ file_name = self._unique_name(name, seen,
local_output_directory)
+ file_path = self._correct_path(file_name,
local_output_directory)
+ else:
+ file_name = name
+ self._create_file(file_name, payload,
local_output_directory)
+
+ def _unique_name(self, name: str, seen: dict[str, int],
local_output_directory: str) -> str:
Review Comment:
Names of methods should start with a verb (`_generate_unique_name`)
##########
providers/imap/src/airflow/providers/imap/hooks/imap.py:
##########
@@ -287,14 +292,54 @@ 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:
+ seen: dict[str, int] = {}
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)
+ if not overwrite:
+ file_name = self._unique_name(name, seen,
local_output_directory)
+ file_path = self._correct_path(file_name,
local_output_directory)
+ while True:
+ try:
+ with open(file_path, "xb") as f:
+ f.write(payload)
+ break
+ except FileExistsError:
+ seen[file_name] = 0
Review Comment:
Might be redundant considering the logic below
--
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]
