Save timestamps

3 files changed, 66 insertions, 26 deletions

diff --git a/models.py b/models.py
index 2c04f67..3e60640 100644
--- a/models.py
+++ b/models.py
@@ -7,6 +7,7 @@ class Message(Base):
     __tablename__ = "messages"
 
     id = Column(Integer, primary_key=True, autoincrement=True)
+    timestamp = Column(String, nullable=False)
     author_id = Column(Integer, nullable=False)
     channel_id = Column(Integer, nullable=False)
     stickers = Column(JSON, nullable=False, default=list)
diff --git a/src/cogs/archive.py b/src/cogs/archive.py
index b5572fe..196eef8 100644
--- a/src/cogs/archive.py
+++ b/src/cogs/archive.py
@@ -8,41 +8,51 @@ import os
 
 from src.utils.db import get_db
 from models import Message
-from config import BOT_COLOR, LOG
+from config import BOT_COLOR, LOG, create_embed
 
 
 class Archive(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
 
-    async def download_attachments(attachments) -> list:
+    async def download_attachments(self, attachments) -> list:
         """Download attachments and return a list of their paths."""
         paths = []
 
         for attachment in attachments:
-            async with aiohttp.ClientSession().get(attachment.url) as response:
-                content_type = response.headers.get("Content-Type")
-                if not content_type:
-                    LOG.warn(
-                        f"Failed to get content type for: {attachment.url}"
+            async with aiohttp.ClientSession() as session:
+                async with session.get(attachment.url) as response:
+                    # Check if the request was successful
+                    if response.status != 200:
+                        LOG.warn(
+                            f"Failed to download attachment: {attachment.url}"
+                        )
+                        continue
+
+                    # Check for content type
+                    content_type = response.headers.get("Content-Type")
+                    if not content_type:
+                        LOG.warn(
+                            f"Failed to get content type for: {attachment.url}"
+                        )
+                        continue
+
+                    # Create a randomized filename
+                    file_extension = content_type.split("/")[-1]
+                    filename = (
+                        "".join(
+                            random.choice(string.ascii_letters)
+                            for i in range(10)
+                        )
+                        + f".{file_extension}"
                     )
-                    continue
 
-                # Create a randomized filename
-                file_extension = content_type.split("/")[-1]
-                filename = (
-                    "".join(
-                        random.choice(string.ascii_letters) for i in range(10)
-                    )
-                    + f".{file_extension}"
-                )
-
-                # Save the image to the filesystem
-                with open(f"images/{filename}]", "wb") as file:
-                    file.write(await response.read())
+                    # Save the attachment
+                    with open(f"images/{filename}]", "wb") as file:
+                        file.write(await response.read())
 
-                # Add the path to the attachments list
-                paths.append(f"images/{filename}")
+                    # Add the path to the attachments list
+                    paths.append(f"images/{filename}")
 
         return paths
 
@@ -71,7 +81,16 @@ class Archive(commands.Cog):
                 ephemeral=True,
             )
 
-        await interaction.response.send_message("Archiving messages now.")
+        embed = discord.Embed(
+            title="Archive Beginning",
+            description=(
+                f"Archiving {amount} messages from {channel.mention}. I will"
+                " DM you a message once complete, make sure to allow messages"
+                " from me."
+            ),
+            color=BOT_COLOR,
+        )
+        await interaction.response.send_message(embed=embed, ephemeral=True)
 
         db = next(get_db())
         count = 0
@@ -79,14 +98,13 @@ class Archive(commands.Cog):
         async for message in messages:
             count += 1
 
-            attachments = []
-
             if not os.path.exists("images"):
                 os.makedirs("images")
 
             attachments = await self.download_attachments(message.attachments)
 
             db_message = Message(
+                timestamp=message.created_at.isoformat(),
                 author_id=message.author.id,
                 channel_id=channel.id,
                 stickers=[sticker.name for sticker in message.stickers],
@@ -98,7 +116,27 @@ class Archive(commands.Cog):
             )
 
             db.add(db_message)
-            db.commit()
+
+            if count > 500:
+                db.commit()
+                count = 0
+
+        db.commit()
+
+        embed = discord.Embed(
+            title="Archive Complete",
+            description=f"Archived {amount} messages from {channel.mention}.",
+            color=BOT_COLOR,
+        )
+
+        try:
+            await interaction.user.send(embed=embed)
+        except discord.Forbidden:
+            await channel.send(
+                f"{interaction.user.mention} I have completed the archive, but"
+                " was unable to DM you the final message. Please check your"
+                " DM settings to receive future messages."
+            )
 
 
 async def setup(bot):
diff --git a/src/schemas/message_schema.py b/src/schemas/message_schema.py
index 79c9659..1738f3b 100644
--- a/src/schemas/message_schema.py
+++ b/src/schemas/message_schema.py
@@ -2,6 +2,7 @@ from pydantic import BaseModel
 
 
 class Message(BaseModel):
+    timestamp: str
     author_id: int
     channel_id: int
     stickers: list[str]