diff --git a/db/attachment.py b/db/attachment.py index e7c80ad..46579f9 100644 --- a/db/attachment.py +++ b/db/attachment.py @@ -1,9 +1,15 @@ +from urllib.parse import urlparse from uuid import UUID import os from typing import Any, Optional +from urllib3.util import parse_url +from urllib3.util import Url + from sqlalchemy import ForeignKey +from sqlalchemy import UniqueConstraint +from sqlalchemy import BigInteger from sqlalchemy.orm import Mapped from sqlalchemy.orm import mapped_column from sqlalchemy.orm import DeclarativeBase @@ -18,7 +24,7 @@ class Attachment(Base): # 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) + msg_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True) url: Mapped[str] @@ -27,15 +33,21 @@ class Attachment(Base): channel: Mapped[dict[str, Any]] vision: Mapped[Optional[dict[str, Any]]] = mapped_column(default=None) + __table_args__ = (UniqueConstraint('msg_id', 'url', name='ix_attachments_msg_id_url_uc'),) + def file_details(self) -> tuple[str, str]: - return os.path.splitext(os.path.basename(self.url)) + u = urlparse(self.url) + if not u.path: + raise RuntimeError(f"attachment url missing path {self.url}") + return os.path.splitext(os.path.basename(u.path)) def object_name(self) -> str: + print(self.id) (filename, ext) = self.file_details() if self.convert_to_webp(): return f"{self.id}.webp" - return f"{id}{ext}" + return f"{self.id}{ext}" def convert_to_webp(self) -> bool: (filename, ext) = self.file_details() diff --git a/main.py b/main.py index e614764..053a0be 100644 --- a/main.py +++ b/main.py @@ -1,38 +1,31 @@ -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 logging import warn import aiohttp +import asyncio +import itertools +import json +import meilisearch +import os from aiobotocore.session import get_session as aiobotocore_get_session -from types_aiobotocore_s3.client import S3Client - -from PIL import Image -from dotenv import dotenv_values, load_dotenv - -from google.cloud import vision as google_vision -from google.protobuf.json_format import MessageToDict, MessageToJson -from google.cloud.vision_v1 import types -import meilisearch -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 +from db.migrate import migrate_db +from dotenv import dotenv_values, load_dotenv +from google.cloud import vision as google_vision +from google.cloud.vision_v1 import types +from io import BytesIO +from PIL import Image +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker, Session +from types_aiobotocore_s3.client import S3Client +from urllib3.util import parse_url +from urllib3.util import Url +from uuid import uuid4 # need to change the env for the google api calls load_dotenv() config = dotenv_values(".env") -DATA_DIR = './data/memestream' MEILI_INDEX = 'memes' MEILI_URL = 'http://localhost:7700' @@ -41,23 +34,13 @@ CHANNEL_IDS = [ 1095060717117374687, # aella memes ] +STORAGE_BUCKET = config['STORAGE_BUCKET'] +assert STORAGE_BUCKET + headers = { 'Authorization': config['DISCORD_AUTH'], } -# TODO -# [ ] Backfill last N images from a channel -# [ ] Set oldest_img_id for channel -# [ ] save img as json + imagefile -# [ ] check extensions (ignore mp4 maybe) - -@dataclass -class MemeDocument: - id: UUID - url: Url - text: str - tags: list[str] - def detect_text(url, vision): request = { 'image': { 'source': { 'image_uri': url } }, @@ -75,15 +58,21 @@ def filter_message_img_urls(messages): , messages) return [item for group in url_lists for item in group] -async def download_channel_attachments(session, storage, vision, channel_id): +async def process_channel( + http_session: aiohttp.ClientSession, + storage: S3Client, + vision: google_vision.ImageAnnotatorClient, + db_session_maker: sessionmaker[Session], + channel_id: int +): async def fetch_channel_messages(channel_id): url = f"https://discord.com/api/v9/channels/{channel_id}/messages" - async with session.get(url, headers=headers) as res: + async with http_session.get(url, headers=headers) as res: return await res.json() async def fetch_guild_details(guild_id): url = f"https://discord.com/api/v9/guilds/{guild_id}" - async with session.get(url, headers=headers) as res: + async with http_session.get(url, headers=headers) as res: data = await res.json() keys_to_remove = ['emojis', 'stickers', 'roles'] @@ -93,38 +82,28 @@ async def download_channel_attachments(session, storage, vision, channel_id): async def fetch_channel(channel_id): url = f"https://discord.com/api/v9/channels/{channel_id}" - async with session.get(url, headers=headers) as res: + async with http_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, - attachment: Attachment - ): - async with session.get(attachment.url) as response: + async def download_attachment(attachment: Attachment): + async with http_session.get(attachment.url) as response: if response.status == 200: - # make a webp conversion for obvious image types - res_bytes = await response.read() - (filename, ext) = attachment.file_details() object_name = attachment.object_name() - print(f"uploading object_name={object_name} filename={filename} ext={ext}") + print(f"uploading object_name={object_name} filename={filename} ext={ext}") + res_bytes = await response.read() + + # make a webp conversion for obvious image types if attachment.convert_to_webp(): + # have to read the whole img before converting 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) + 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) + 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']] @@ -164,6 +143,28 @@ async def download_channel_attachments(session, storage, vision, channel_id): # super wack flatmap and filter attachments = list(filter(None, itertools.chain.from_iterable(attachment_groups))) + db_session = db_session_maker() + + # filter attachments that already exist in the db + # dirty way is just check if msg_id exists in db, don't care about specific attachments + + # ok now we will try and upload each attachment + # for the ones that succeed we will add them to our session to be stored + async def upload_and_persist(attachment: Attachment): + try: + await download_attachment(attachment=attachment) + # this is kinda dodge, the exception will catch if we fail to add to session + db_session.add(attachment) + except Exception as e: + warn(e) + warn(f'deleting attachment upload {attachment.object_name()}') + await storage.delete_object(Bucket=STORAGE_BUCKET, Key=attachment.object_name()) + + attachments = [upload_and_persist(attachment=a) for a in attachments] + await asyncio.gather(*attachments) + + db_session.commit() + # NOW # ------------------- # insert attachment @@ -207,29 +208,33 @@ async def main(): db_engine = create_engine(db_url, echo=True) migrate_db(db_engine, db_url) + db_session_maker = sessionmaker(bind=db_engine) + 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: + ) as storage, aiohttp.ClientSession() as http_session: + storage: S3Client - storage_bucket = config['STORAGE_BUCKET'] - assert storage_bucket - # Make the bucket if it doesn't exist. - exists = await check_bucket_exists(storage, storage_bucket) + exists = await check_bucket_exists(storage, STORAGE_BUCKET) if not exists: - await storage.create_bucket(Bucket=storage_bucket) - print("Created bucket", storage_bucket) + await storage.create_bucket(Bucket=STORAGE_BUCKET) + print("Created bucket", STORAGE_BUCKET) else: - print("Bucket", storage_bucket, "already exists") + 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) + channels = [process_channel(http_session, storage, vision, db_session_maker, channel_id) for channel_id in CHANNEL_IDS] + memes = await asyncio.gather(*channels) asyncio.run(main()) + +# TODO +# [ ] Backfill last N images from a channel +# [ ] Set oldest_img_id for channel +# logger diff --git a/migrations/versions/e04bee35feab_create_memes_and_attachments_table.py b/migrations/versions/50b5b3980665_create_memes_and_attachments_table.py similarity index 75% rename from migrations/versions/e04bee35feab_create_memes_and_attachments_table.py rename to migrations/versions/50b5b3980665_create_memes_and_attachments_table.py index 918cefc..416283a 100644 --- a/migrations/versions/e04bee35feab_create_memes_and_attachments_table.py +++ b/migrations/versions/50b5b3980665_create_memes_and_attachments_table.py @@ -1,8 +1,8 @@ """Create Memes and Attachments table -Revision ID: e04bee35feab +Revision ID: 50b5b3980665 Revises: -Create Date: 2024-04-08 10:45:47.138438 +Create Date: 2024-04-08 18:30:57.222252 """ from typing import Sequence, Union @@ -12,7 +12,7 @@ import sqlalchemy as sa # revision identifiers, used by Alembic. -revision: str = 'e04bee35feab' +revision: str = '50b5b3980665' down_revision: Union[str, None] = None branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None @@ -22,17 +22,18 @@ def upgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### op.create_table('attachments', sa.Column('id', sa.Uuid(), nullable=False), - sa.Column('msg_id', sa.Integer(), nullable=False), - sa.Column('attributes', sa.JSON(), nullable=False), - sa.Column('vision', sa.JSON(), nullable=False), + sa.Column('msg_id', sa.BigInteger(), nullable=False), + sa.Column('url', sa.String(), nullable=False), sa.Column('message', sa.JSON(), nullable=False), sa.Column('guild', sa.JSON(), nullable=False), sa.Column('channel', sa.JSON(), nullable=False), - sa.PrimaryKeyConstraint('id') + sa.Column('vision', sa.JSON(), nullable=True), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('msg_id', 'url', name='ix_attachments_msg_id_url_uc') ) op.create_index(op.f('ix_attachments_msg_id'), 'attachments', ['msg_id'], unique=False) op.create_table('memes', - sa.Column('id', sa.Uuid(), nullable=False), + sa.Column('id', sa.Integer(), nullable=False), sa.Column('url', sa.String(), nullable=False), sa.Column('text', sa.String(), nullable=False), sa.Column('tags', sa.JSON(), nullable=False), diff --git a/migrations/versions/9159293e4f30_allow_vision_to_be_nullable.py b/migrations/versions/9159293e4f30_allow_vision_to_be_nullable.py deleted file mode 100644 index 35f101f..0000000 --- a/migrations/versions/9159293e4f30_allow_vision_to_be_nullable.py +++ /dev/null @@ -1,36 +0,0 @@ -"""allow vision to be nullable - -Revision ID: 9159293e4f30 -Revises: e04bee35feab -Create Date: 2024-04-08 10:50:21.447942 - -""" -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -from sqlalchemy.dialects import postgresql - -# revision identifiers, used by Alembic. -revision: str = '9159293e4f30' -down_revision: Union[str, None] = 'e04bee35feab' -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.alter_column('attachments', 'vision', - existing_type=postgresql.JSON(astext_type=sa.Text()), - nullable=True) - op.drop_column('attachments', 'attributes') - # ### end Alembic commands ### - - -def downgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.add_column('attachments', sa.Column('attributes', postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=False)) - op.alter_column('attachments', 'vision', - existing_type=postgresql.JSON(astext_type=sa.Text()), - nullable=False) - # ### end Alembic commands ###