ntr/memestream
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

about to filter

Browse Source
This commit is contained in:
Nathan Rashleigh committed 2024-04-08 18:44:40 +10:00
1 parent 557020a824
commit d0e2be7d12
4 files changed
+100 -118
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

db/attachment.py
+15 -3
View File
@@ -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()
main.py
+76 -71
View File
@@ -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]
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
migrations/versions/e04bee35feab_create_memes_and_attachments_table.py → migrations/versions/50b5b3980665_create_memes_and_attachments_table.py
+9 -8
View File
@@ -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),
migrations/versions/9159293e4f30_allow_vision_to_be_nullable.py
-36
View File
@@ -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 ###