From a69b15a3e84a2a50354ffd1db4d408953babc7d5 Mon Sep 17 00:00:00 2001 From: Nathan Rashleigh Date: Mon, 8 Apr 2024 17:01:47 +1000 Subject: [PATCH] attachment class --- db/attachment.py | 45 ++++++++++++--- db/base.py | 12 +++- db/meme.py | 16 ++++-- main.py | 146 +++++++++++++++++++++++++++-------------------- requirements.txt | 4 +- 5 files changed, 144 insertions(+), 79 deletions(-) diff --git a/db/attachment.py b/db/attachment.py index bb431c6..e7c80ad 100644 --- a/db/attachment.py +++ b/db/attachment.py @@ -1,4 +1,13 @@ -from sqlalchemy import create_engine, Column, Integer, String, JSON, Uuid +from uuid import UUID +import os + +from typing import Any, Optional + +from sqlalchemy import ForeignKey +from sqlalchemy.orm import Mapped +from sqlalchemy.orm import mapped_column +from sqlalchemy.orm import DeclarativeBase +from sqlalchemy.orm import relationship from db.base import Base @@ -6,10 +15,30 @@ from db.base import Base class Attachment(Base): __tablename__ = 'attachments' - id = Column(Uuid, primary_key=True) - msg_id = Column(Integer, nullable=False, index=True) - attributes = Column(JSON, nullable=False) - vision = Column(JSON, nullable=False) - message = Column(JSON, nullable=False) - guild = Column(JSON, nullable=False) - channel = Column(JSON, nullable=False) + # can be multiple attachments on a msg + # so we make our own ids for attachments + id: Mapped[UUID] = mapped_column(primary_key=True) + msg_id: Mapped[int] = mapped_column(nullable=False, index=True) + + url: Mapped[str] + + message: Mapped[dict[str, Any]] + guild: Mapped[dict[str, Any]] + channel: Mapped[dict[str, Any]] + vision: Mapped[Optional[dict[str, Any]]] = mapped_column(default=None) + + def file_details(self) -> tuple[str, str]: + return os.path.splitext(os.path.basename(self.url)) + + def object_name(self) -> str: + (filename, ext) = self.file_details() + if self.convert_to_webp(): + return f"{self.id}.webp" + + return f"{id}{ext}" + + def convert_to_webp(self) -> bool: + (filename, ext) = self.file_details() + convert_extensions = ['jpg', 'png', 'jpeg', 'webp'] + return ext.lower()[1:] in convert_extensions + diff --git a/db/base.py b/db/base.py index 860e542..5532fdb 100644 --- a/db/base.py +++ b/db/base.py @@ -1,3 +1,11 @@ -from sqlalchemy.ext.declarative import declarative_base +from typing import Any -Base = declarative_base() +from sqlalchemy.orm import DeclarativeBase +from sqlalchemy.orm import MappedAsDataclass +from sqlalchemy.types import JSON + +class Base(DeclarativeBase, MappedAsDataclass): + type_annotation_map = { + dict[str, Any]: JSON, + list[str]: JSON, + } diff --git a/db/meme.py b/db/meme.py index 6b09459..1f914a7 100644 --- a/db/meme.py +++ b/db/meme.py @@ -1,4 +1,10 @@ -from sqlalchemy import Column, Integer, String, JSON, Uuid +from typing import List + +from sqlalchemy import ForeignKey +from sqlalchemy.orm import Mapped +from sqlalchemy.orm import mapped_column +from sqlalchemy.orm import DeclarativeBase +from sqlalchemy.orm import relationship from db.base import Base @@ -6,7 +12,7 @@ from db.base import Base class Meme(Base): __tablename__ = 'memes' - id = Column(Uuid, primary_key=True) - url = Column(String, nullable=False) - text = Column(String, nullable=False) - tags = Column(JSON, nullable=False) \ No newline at end of file + id: Mapped[int] = mapped_column(primary_key=True) + url: Mapped[str] = mapped_column(nullable=False) + text: Mapped[str] = mapped_column(nullable=False) + tags: Mapped[list[str]] = mapped_column(nullable=False) diff --git a/main.py b/main.py index ee18c4c..e614764 100644 --- a/main.py +++ b/main.py @@ -1,12 +1,13 @@ import os import asyncio import json +import itertools + from urllib3.util import parse_url from pathlib import Path from io import BytesIO from dataclasses import dataclass from uuid import uuid4, UUID -from typing import Dict, Any, List import aiohttp @@ -24,6 +25,8 @@ from urllib3.util import Url, url from sqlalchemy import create_engine from db.migrate import migrate_db +from db.attachment import Attachment +from db.meme import Meme # need to change the env for the google api calls load_dotenv() @@ -53,7 +56,7 @@ class MemeDocument: id: UUID url: Url text: str - tags: List[str] + tags: list[str] def detect_text(url, vision): request = { @@ -93,68 +96,85 @@ async def download_channel_attachments(session, storage, vision, channel_id): async with session.get(url, headers=headers) as res: return await res.json() + + async def upload_object(object_name, body): + storage_bucket = config['STORAGE_BUCKET'] + if not storage_bucket: + raise RuntimeError(f"config missing STORAGE_BUCKET") + return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=body) + async def download_attachment( session: aiohttp.ClientSession, storage: S3Client, - url: Url, - message, + attachment: Attachment ): - # there's a big param of some kind after the image url that we remove - if not url.path: - raise RuntimeError(f"url missing path url={url}") - - storage_bucket = config['UPLOAD_BUCKET'] - if not storage_bucket: - raise RuntimeError(f"config missing UPLOAD_BUCKET") - - (filename, ext) = os.path.splitext(os.path.basename(url.path)) - # TODO - # FILTER WANTED TYPES - - id = uuid4() - - async with session.get(url.url) as response: + async with session.get(attachment.url) as response: if response.status == 200: - # make a webp conversion for obvious image types - convert_extensions = ['jpg', 'png', 'jpeg', 'webp'] res_bytes = await response.read() - object_name: str - # trim the . - if ext.lower()[1:] in convert_extensions: + (filename, ext) = attachment.file_details() + object_name = attachment.object_name() + print(f"uploading object_name={object_name} filename={filename} ext={ext}") + + if attachment.convert_to_webp(): image = Image.open(BytesIO(res_bytes)) - webp_output = BytesIO() image.save(webp_output, format='WebP') webp_output.seek(0) + return await upload_object(object_name=object_name, body=webp_output) - object_name = f"{id}.webp" - print(f"uploading object_name={object_name} filename={filename} ext={ext}") - return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=webp_output) + return await upload_object(object_name=object_name, body=res_bytes) - object_name = f"{id}{ext}" - print(f"uploading object_name={object_name} filename={filename} ext={ext}") - return await storage.put_object(Bucket=storage_bucket, Key=object_name, Body=res_bytes) + def create_attachments(message, channel, guild): + urls = [parse_url(attachment['url']) for attachment in message['attachments']] - # add discord id to server processed msgs list - # if there then nothing to do + def generate_attachment(url: Url): + if not url.path: + return None - # return MemeDocument( - # id=id, - # url=self.upload_url, - # text=self.text or '', - # tags=[self.guild_name, self.channel_name], - # ) + (filename, ext) = os.path.splitext(os.path.basename(url.path)) + extensions = ['jpg', 'png', 'jpeg', 'webp', 'mov', 'gif', 'mp4'] + if ext.lower()[1:] not in extensions: + return None + + return Attachment( + id=uuid4(), + msg_id=int(message['id']), + message=message, + guild=guild, + url=url.url, + channel=channel + ) + + return [generate_attachment(url) for url in urls] channel = await fetch_channel(channel_id) guild = await fetch_guild_details(channel['guild_id']) print(f"Fetching messages for channel guild={guild['name']} name={channel['name']} id={channel['id']}") messages = await fetch_channel_messages(channel_id) - urls_and_messages = filter_message_img_urls(messages) - tasks = [download_attachment(session, storage, url, message) for (url, message) in urls_and_messages] - await asyncio.gather(*tasks) + # at this point we have the channel messages, the guild (server) and channel info + + # first thing we have to do is parse out the urls in the msg attachments + # once we have those we can generate where we should put them in our storage + # should only persist them later once the processing is done + attachment_groups = [create_attachments(message, channel, guild) for message in messages] + + # super wack flatmap and filter + attachments = list(filter(None, itertools.chain.from_iterable(attachment_groups))) + + # NOW + # ------------------- + # insert attachment + # upload to s3 + # fetch vision for relevant attachments + # create meme from attachment and vision + # insert meme + # add documents to meili + + # tasks = [download_attachment(session, storage, url, message) for (url, message) in urls_and_messages] + # await asyncio.gather(*tasks) # [attachment.detect_text(vision) for attachment in attachments] # [attachment.persist() for attachment in attachments] @@ -184,32 +204,32 @@ async def main(): # db setup db_url = config['POSTGRES_URL'] - db_engine = create_engine(db_url) + db_engine = create_engine(db_url, echo=True) migrate_db(db_engine, db_url) - # aioboto = aiobotocore_get_session() - # async with aioboto.create_client( - # 's3', - # endpoint_url=config['STORAGE_URL'], - # aws_access_key_id=config['STORAGE_ACCESS_KEY'], - # aws_secret_access_key=config['STORAGE_SECRET_KEY'], - # ) as storage: - # storage: S3Client + aioboto = aiobotocore_get_session() + async with aioboto.create_client( + 's3', + endpoint_url=config['STORAGE_URL'], + aws_access_key_id=config['STORAGE_ACCESS_KEY'], + aws_secret_access_key=config['STORAGE_SECRET_KEY'], + ) as storage: + storage: S3Client - # storage_bucket = config['UPLOAD_BUCKET'] - # assert storage_bucket + storage_bucket = config['STORAGE_BUCKET'] + assert storage_bucket - # # Make the bucket if it doesn't exist. - # exists = await check_bucket_exists(storage, storage_bucket) - # if not exists: - # await storage.create_bucket(Bucket=storage_bucket) - # print("Created bucket", storage_bucket) - # else: - # print("Bucket", storage_bucket, "already exists") + # Make the bucket if it doesn't exist. + exists = await check_bucket_exists(storage, storage_bucket) + if not exists: + await storage.create_bucket(Bucket=storage_bucket) + print("Created bucket", storage_bucket) + else: + print("Bucket", storage_bucket, "already exists") - # async with aiohttp.ClientSession() as session: - # channels = [download_channel_attachments(session, storage, vision, channel_id) for channel_id in CHANNEL_IDS] - # memes = await asyncio.gather(*channels) + async with aiohttp.ClientSession() as session: + channels = [download_channel_attachments(session, storage, vision, channel_id) for channel_id in CHANNEL_IDS] + memes = await asyncio.gather(*channels) asyncio.run(main()) diff --git a/requirements.txt b/requirements.txt index cc4f917..ba6e93a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,8 @@ google-cloud-vision meilisearch pillow aiobotocore +types-aiobotocore-s3 + sqlalchemy alembic -psycopg2-binary +psycopg2-binary \ No newline at end of file