import json import asyncpg from utils.files import media_file_name from utils.read.models import MediaPlaceView, MediaVersionView, MediaView MEDIA_COLS = ( "m.id, m.account_id, m.chat_id, m.message_id, m.kind, m.storage_key, " "m.file_size, m.mime, m.ttl_seconds, m.downloaded, m.extracted_text, " "m.created_at, src.original_name" ) ORIGINAL_NAME_JOIN = """ LEFT JOIN LATERAL ( SELECT msg.raw->m.kind->>'file_name' AS original_name FROM messages msg WHERE msg.account_id = m.account_id AND msg.chat_id = m.chat_id AND msg.message_id = m.message_id ORDER BY msg.date DESC LIMIT 1 ) src ON true """ _VERSION_COLS = "id, kind, storage_key, file_size, mime, observed_at" _PLACES = """ SELECT account_id, chat_id, message_id, kind, extracted_text FROM media WHERE storage_key = $1 AND ($2::bigint IS NULL OR account_id = $2) ORDER BY created_at DESC LIMIT 100 """ _WEB_PAGE_MEDIA_KINDS = ("photo", "video", "animation", "document", "audio") def media_view(row: asyncpg.Record) -> MediaView: fields = dict(row) original = fields.pop("original_name", None) return MediaView( **fields, file_name=media_file_name( fields["kind"], fields["mime"], fields["message_id"], original ), ) async def _web_page_media_stub( pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int ) -> MediaView | None: row = await pool.fetchrow( "SELECT date, raw->'web_page' AS web_page FROM messages " "WHERE account_id = $1 AND chat_id = $2 AND message_id = $3", account_id, chat_id, message_id, ) if row is None or row["web_page"] is None: return None web_page = json.loads(row["web_page"]) if not isinstance(web_page, dict): return None kind = next( (k for k in _WEB_PAGE_MEDIA_KINDS if isinstance(web_page.get(k), dict)), None ) if kind is None: return None obj = web_page[kind] return MediaView( id=0, account_id=account_id, chat_id=chat_id, message_id=message_id, kind=kind, storage_key=None, file_size=obj.get("file_size"), mime=obj.get("mime_type"), ttl_seconds=None, downloaded=False, extracted_text=None, created_at=row["date"], file_name=media_file_name( kind, obj.get("mime_type"), message_id, obj.get("file_name") ), ) async def get_media(pool: asyncpg.Pool, media_id: int) -> MediaView | None: row = await pool.fetchrow( f"SELECT {MEDIA_COLS} FROM media m {ORIGINAL_NAME_JOIN} WHERE m.id = $1", # noqa: S608 media_id, ) return media_view(row) if row else None async def get_message_media( pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int ) -> MediaView | None: row = await pool.fetchrow( f"SELECT {MEDIA_COLS} FROM media m {ORIGINAL_NAME_JOIN} " # noqa: S608 "WHERE m.account_id = $1 AND m.chat_id = $2 AND m.message_id = $3", account_id, chat_id, message_id, ) if row is not None: return media_view(row) return await _web_page_media_stub(pool, account_id, chat_id, message_id) async def find_by_content( pool: asyncpg.Pool, storage_key: str, account_id: int | None = None ) -> list[MediaPlaceView]: rows = await pool.fetch(_PLACES, storage_key, account_id) return [MediaPlaceView(**dict(row)) for row in rows] async def get_media_versions( pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int ) -> list[MediaVersionView]: rows = await pool.fetch( f"SELECT {_VERSION_COLS} FROM media_versions " # noqa: S608 "WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 " "ORDER BY observed_at LIMIT 500", account_id, chat_id, message_id, ) return [MediaVersionView(**dict(row)) for row in rows] async def get_media_version( pool: asyncpg.Pool, version_id: int ) -> MediaVersionView | None: row = await pool.fetchrow( f"SELECT {_VERSION_COLS} FROM media_versions WHERE id = $1", # noqa: S608 version_id, ) return MediaVersionView(**dict(row)) if row else None