124 lines
3.7 KiB
Python
124 lines
3.7 KiB
Python
import asyncpg
|
|
|
|
from utils.files import media_file_name
|
|
from utils.read.message_view import load_raw
|
|
from utils.read.models import 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"
|
|
|
|
_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 FROM messages "
|
|
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
|
|
account_id,
|
|
chat_id,
|
|
message_id,
|
|
)
|
|
if row is None:
|
|
return None
|
|
web_page = load_raw(row["raw"]).get("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 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",
|
|
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
|