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 db import attachment 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.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.ext.asyncio import async_sessionmaker from sqlalchemy.ext.asyncio import async_scoped_session from sqlalchemy.future import select 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") MEILI_INDEX = 'memes' MEILI_URL = 'http://localhost:7700' CHANNEL_IDS = [ 367936426962190336, # TPM shitposts 1095060717117374687, # aella memes ] STORAGE_BUCKET = config['STORAGE_BUCKET'] assert STORAGE_BUCKET headers = { 'Authorization': config['DISCORD_AUTH'], } async def process_channel( http_session: aiohttp.ClientSession, storage: S3Client, vision: google_vision.ImageAnnotatorClient, db_session: async_sessionmaker[AsyncSession], channel_id: int ): async def fetch_channel_messages(channel_id): url = f"https://discord.com/api/v9/channels/{channel_id}/messages" 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 http_session.get(url, headers=headers) as res: data = await res.json() keys_to_remove = ['emojis', 'stickers', 'roles'] [data.pop(key, None) for key in keys_to_remove] return data async def fetch_channel(channel_id): url = f"https://discord.com/api/v9/channels/{channel_id}" async with http_session.get(url, headers=headers) as res: return await res.json() async def download_attachment(attachment: Attachment): async with http_session.get(attachment.url) as response: if response.status == 200: (filename, ext) = attachment.file_details() object_name = attachment.object_name() 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 storage.put_object(Bucket=STORAGE_BUCKET, Key=object_name, Body=webp_output) 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']), attachment['id']) for attachment in message['attachments']] def generate_attachment(url: Url, attachment_id: str): if not url.path: return None (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(), attachment_id=int(attachment_id), msg_id=int(message['id']), message=message, guild=guild, url=url.url, channel=channel ) return [generate_attachment(*url) for url in urls] async def upload_and_persist(attachment: Attachment, session: AsyncSession): try: await download_attachment(attachment=attachment) # this is kinda dodge, the exception will catch if we fail to add to session 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()) def detect_text(url): request = { 'image': { 'source': { 'image_uri': url } }, 'features': [types.Feature(type=types.Feature.Type.TEXT_DETECTION)], } response: types.AnnotateImageResponse = vision.annotate_image(request=request) vision_json = json.loads(types.AnnotateImageResponse.to_json(response)) return request 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) # 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))) print(f"found {len(attachments)} attachments for channel guild={guild['name']} name={channel['name']} id={channel['id']}"); attachment_ids = list(map(lambda a: a.attachment_id, attachments)) async with db_session() as session: # this is a bit weird # we find the union of attachment_ids in the db and our list results = await session.execute( select(Attachment.attachment_id) .filter(Attachment.attachment_id.in_(attachment_ids))) existing_ids = [row[0] for row in results.fetchall()] # and then filter out msgs that are already processed new_attachments = list(filter(lambda a: a.attachment_id not in existing_ids, attachments)) print(f"{len(new_attachments)} attachments to download for channel guild={guild['name']} name={channel['name']} id={channel['id']}"); # ok now we will try and upload each attachment # for the ones that succeed we will add them to our session to be stored downloads = [upload_and_persist(attachment=a, session=session) for a in new_attachments] await asyncio.gather(*downloads) await session.commit() # NOW # ------------------- # 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] # client = meilisearch.Client(MEILI_URL, config['MEILI_MASTER_KEY']) # # have to vars the dataclass out for meili # documents = [vars(attachment.to_meme()) for attachment in attachments] # client.index(MEILI_INDEX).add_documents(documents) async def check_bucket_exists(storage, bucket_name): try: await storage.head_bucket(Bucket=bucket_name) return True except storage.exceptions.ClientError as e: error_code = e.response.get('Error', {}).get('Code') if error_code == '404': return False # Bucket does not exist else: raise # Reraise other exceptions async def main(): assert config['STORAGE_URL'] assert config['POSTGRES_URL'] vision = google_vision.ImageAnnotatorClient() # db setup db_url = config['POSTGRES_URL'] db_engine = create_async_engine(db_url, echo=True) await migrate_db(db_engine) print('db is up to date') db_session = async_sessionmaker(db_engine, expire_on_commit=False) 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, aiohttp.ClientSession() as http_session: storage: S3Client # 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") channels = [process_channel(http_session, storage, vision, db_session, channel_id) for channel_id in CHANNEL_IDS] memes = await asyncio.gather(*channels) await db_engine.dispose() asyncio.run(main()) # TODO # [ ] Backfill last N images from a channel # [ ] Set oldest_img_id for channel # logger