Skip to content

Commit

Permalink
Refactor mongodb client append_log store images internally in "attach…
Browse files Browse the repository at this point in the history
…ment_data field"

This change ends up being *mostly* backwards compatible since we still store the url, though this will be basically useless in the future.

Attachment size is currently hardcoded to 8MB, in the future this should be able to be set by users.
  • Loading branch information
khakers committed Nov 7, 2023
1 parent 36b7910 commit 4cffc23
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 12 deletions.
65 changes: 53 additions & 12 deletions core/clients.py
Original file line number Diff line number Diff line change
Expand Up @@ -645,6 +645,58 @@ async def edit_message(self, message_id: Union[int, str], new_content: str) -> N
{"$set": {"messages.$.content": new_content, "messages.$.edited": True}},
)

async def _handle_attachments(self, message: Message) -> list:
attachments = []
in_db_storage = self.bot.config.get("in_db_storage")
# 8 MB to bytes
image_max_size = 1024 * 1024 * 8

if in_db_storage:
for attachment in message.attachments:
# TODO this is non ideal
# ideally we would async iterate over the attachments so we don't block on io for every attachment
image_data = await attachment.read()
# Don't store images larger than
if len(image_data) > image_max_size:
logger.warning(
"Attachment %s (%s) is too large to store in the database, skipping.",
attachment.filename,
attachment.id,
)
continue
attachments.append(
{
"id": attachment.id,
"filename": attachment.filename,
"is_image": attachment.content_type.startswith("image/"),
"size": len(image_data),
"type": "internal",
# URL points to the original discord URL
"url": attachment.url,
# Attachment data is the raw bytes of the image
"attachment_data": image_data,
"content_type": attachment.content_type,
"width": attachment.width,
"height": attachment.height,
}
)
else:
for attachment in message.attachments:
attachments.append(
{
"id": attachment.id,
"filename": attachment.filename,
"is_image": attachment.content_type.startswith("image/"),
"size": attachment.size,
"url": attachment.url,
"content_type": attachment.content_type,
"width": attachment.width,
"height": attachment.height,
}
)

return attachments

async def append_log(
self,
message: Message,
Expand All @@ -668,18 +720,7 @@ async def append_log(
},
"content": message.content,
"type": type_,
"attachments": [
{
"id": a.id,
"filename": a.filename,
# In previous versions this was true for both videos and images
"is_image": a.content_type.startswith("image/"),
"size": a.size,
"url": a.url,
"content_type": a.content_type,
}
for a in message.attachments
],
"attachments": await self._handle_attachments(message),
}

return await self.logs.find_one_and_update(
Expand Down
1 change: 1 addition & 0 deletions core/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,7 @@ class ConfigManager:
"log_level": "INFO",
# data collection
"data_collection": True,
"use_in_database_image_store": True,
}

colors = {"mod_color", "recipient_color", "main_color", "error_color"}
Expand Down
9 changes: 9 additions & 0 deletions core/config_help.json
Original file line number Diff line number Diff line change
Expand Up @@ -1170,5 +1170,14 @@
"If this configuration is enabled, only roles that are hoisted (displayed seperately in member list) will be used. If a user has no hoisted roles, it will return 'None'.",
"If you would like to display the top role of a user regardless of if it's hoisted or not, disable `use_hoisted_top_role`."
]
},
"use_in_database_image_store": {
"default": "Yes",
"description": "Controls if images should be stored in the database as base64 strings.",
"examples": [],
"notes": [
"If this configuration is enabled, images will be stored in the database. If disabled, images will be stored as Discord CDN links.",
"This configuration can only be set through `.env` file or environment (config) variables."
]
}
}

0 comments on commit 4cffc23

Please sign in to comment.