keep downloaded files after emailed (#1322)

This commit is contained in:
LawyZheng
2024-12-05 13:35:13 +08:00
committed by GitHub
parent fc56299eac
commit 0099e78a2c

View File

@@ -1271,66 +1271,61 @@ class SendEmailBlock(Block):
file_names_by_hash: dict[str, list[str]] = defaultdict(list) file_names_by_hash: dict[str, list[str]] = defaultdict(list)
for filename in self._get_file_paths(workflow_run_context, workflow_run_id): for filename in self._get_file_paths(workflow_run_context, workflow_run_id):
path = None if filename.startswith("s3://"):
try: path = await download_from_s3(self.get_async_aws_client(), filename)
if filename.startswith("s3://"): elif filename.startswith("http://") or filename.startswith("https://"):
path = await download_from_s3(self.get_async_aws_client(), filename) path = await download_file(filename)
elif filename.startswith("http://") or filename.startswith("https://"): else:
path = await download_file(filename) LOG.info("SendEmailBlock: Looking for file locally", filename=filename)
else: if not os.path.exists(filename):
LOG.info("SendEmailBlock: Looking for file locally", filename=filename)
if not os.path.exists(filename):
raise FileNotFoundError(f"File not found: {filename}")
if not os.path.isfile(filename):
raise IsADirectoryError(f"Path is a directory: {filename}")
LOG.info("SendEmailBlock: Found file locally", path=path)
path = filename
if not path:
raise FileNotFoundError(f"File not found: {filename}") raise FileNotFoundError(f"File not found: {filename}")
if not os.path.isfile(filename):
raise IsADirectoryError(f"Path is a directory: {filename}")
# Guess the content type based on the file's extension. Encoding path = filename
# will be ignored, although we should check for simple things like LOG.info("SendEmailBlock: Found file locally", path=path)
# gzip'd or compressed files.
kind = filetype.guess(path)
if kind:
ctype = kind.mime
extension = kind.extension
else:
# No guess could be made, or the file is encoded (compressed), so
# use a generic bag-of-bits type.
ctype = "application/octet-stream"
extension = None
maintype, subtype = ctype.split("/", 1) if not path:
attachment_path = Path(path) raise FileNotFoundError(f"File not found: {filename}")
attachment_filename = attachment_path.name
# Check if the filename has an extension # Guess the content type based on the file's extension. Encoding
if not attachment_path.suffix: # will be ignored, although we should check for simple things like
# If no extension, guess it based on the MIME type # gzip'd or compressed files.
if extension: kind = filetype.guess(path)
attachment_filename += f".{extension}" if kind:
ctype = kind.mime
extension = kind.extension
else:
# No guess could be made, or the file is encoded (compressed), so
# use a generic bag-of-bits type.
ctype = "application/octet-stream"
extension = None
LOG.info( maintype, subtype = ctype.split("/", 1)
"SendEmailBlock: Adding attachment", attachment_path = Path(path)
filename=attachment_filename, attachment_filename = attachment_path.name
# Check if the filename has an extension
if not attachment_path.suffix:
# If no extension, guess it based on the MIME type
if extension:
attachment_filename += f".{extension}"
LOG.info(
"SendEmailBlock: Adding attachment",
filename=attachment_filename,
maintype=maintype,
subtype=subtype,
)
with open(path, "rb") as fp:
msg.add_attachment(
fp.read(),
maintype=maintype, maintype=maintype,
subtype=subtype, subtype=subtype,
filename=attachment_filename,
) )
with open(path, "rb") as fp: file_hash = calculate_sha256_for_file(path)
msg.add_attachment( file_names_by_hash[file_hash].append(path)
fp.read(),
maintype=maintype,
subtype=subtype,
filename=attachment_filename,
)
file_hash = calculate_sha256_for_file(path)
file_names_by_hash[file_hash].append(path)
finally:
if path:
os.unlink(path)
# Calculate file stats based on content hashes # Calculate file stats based on content hashes
total_files = sum(len(files) for files in file_names_by_hash.values()) total_files = sum(len(files) for files in file_names_by_hash.values())