256 lines
9.3 KiB
Python
256 lines
9.3 KiB
Python
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=False)
|
|
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
|