From 4cffc23de269f1dffb15b2ec010bd7ab078323e2 Mon Sep 17 00:00:00 2001
From: Khakers <22665282+khakers@users.noreply.github.com>
Date: Mon, 6 Nov 2023 23:31:24 -0800
Subject: [PATCH] Refactor mongodb client append_log store images internally in
 "attachment_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.
---
 core/clients.py       | 65 +++++++++++++++++++++++++++++++++++--------
 core/config.py        |  1 +
 core/config_help.json |  9 ++++++
 3 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/core/clients.py b/core/clients.py
index e80225d9e9f..fee42525e99 100644
--- a/core/clients.py
+++ b/core/clients.py
@@ -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,
@@ -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(
diff --git a/core/config.py b/core/config.py
index b8bb35d40c2..60ab4e2fd1b 100644
--- a/core/config.py
+++ b/core/config.py
@@ -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"}
diff --git a/core/config_help.json b/core/config_help.json
index 80ebf1a48a3..f8ba05164bb 100644
--- a/core/config_help.json
+++ b/core/config_help.json
@@ -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."
+        ]
     }
 }
\ No newline at end of file