112 lines
2.5 KiB
Python
112 lines
2.5 KiB
Python
import json
|
|
|
|
import asyncpg
|
|
|
|
_INSERT_AVATAR = """
|
|
INSERT INTO avatars
|
|
(account_id, owner_id, owner_kind, unique_id, storage_key, file_size, mime,
|
|
downloaded, raw)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9::jsonb)
|
|
ON CONFLICT (account_id, owner_id, unique_id) DO NOTHING
|
|
"""
|
|
|
|
_EXISTS = """
|
|
SELECT 1 FROM avatars
|
|
WHERE account_id = $1 AND owner_id = $2 AND unique_id = $3
|
|
"""
|
|
|
|
_GET_FILE = """
|
|
SELECT raw ->> 'file_id' AS file_id, downloaded FROM avatars
|
|
WHERE account_id = $1 AND owner_kind = $2 AND owner_id = $3 AND unique_id = $4
|
|
"""
|
|
|
|
_MARK_DOWNLOADED = """
|
|
UPDATE avatars SET downloaded = true, storage_key = $5, file_size = $6
|
|
WHERE account_id = $1 AND owner_kind = $2 AND owner_id = $3 AND unique_id = $4
|
|
"""
|
|
|
|
|
|
async def avatar_exists(
|
|
pool: asyncpg.Pool, account_id: int, owner_id: int, unique_id: str
|
|
) -> bool:
|
|
row = await pool.fetchval(_EXISTS, account_id, owner_id, unique_id)
|
|
return row is not None
|
|
|
|
|
|
async def insert_avatar( # noqa: PLR0913
|
|
pool: asyncpg.Pool,
|
|
account_id: int,
|
|
owner_id: int,
|
|
owner_kind: str,
|
|
unique_id: str,
|
|
storage_key: str | None,
|
|
file_size: int | None,
|
|
mime: str | None,
|
|
raw: str,
|
|
*,
|
|
downloaded: bool,
|
|
) -> None:
|
|
await pool.execute(
|
|
_INSERT_AVATAR,
|
|
account_id,
|
|
owner_id,
|
|
owner_kind,
|
|
unique_id,
|
|
storage_key,
|
|
file_size,
|
|
mime,
|
|
downloaded,
|
|
raw,
|
|
)
|
|
|
|
|
|
async def note_avatar( # noqa: PLR0913
|
|
pool: asyncpg.Pool,
|
|
account_id: int,
|
|
owner_id: int,
|
|
owner_kind: str,
|
|
unique_id: str,
|
|
file_id: str,
|
|
) -> None:
|
|
await insert_avatar(
|
|
pool,
|
|
account_id,
|
|
owner_id,
|
|
owner_kind,
|
|
unique_id,
|
|
None,
|
|
None,
|
|
None,
|
|
json.dumps({"file_id": file_id}),
|
|
downloaded=False,
|
|
)
|
|
|
|
|
|
async def get_avatar_file(
|
|
pool: asyncpg.Pool, account_id: int, owner_kind: str, owner_id: int, unique_id: str
|
|
) -> tuple[str | None, bool] | None:
|
|
row = await pool.fetchrow(_GET_FILE, account_id, owner_kind, owner_id, unique_id)
|
|
if row is None:
|
|
return None
|
|
return row["file_id"], row["downloaded"]
|
|
|
|
|
|
async def mark_avatar_downloaded( # noqa: PLR0913
|
|
pool: asyncpg.Pool,
|
|
account_id: int,
|
|
owner_kind: str,
|
|
owner_id: int,
|
|
unique_id: str,
|
|
storage_key: str,
|
|
file_size: int,
|
|
) -> None:
|
|
await pool.execute(
|
|
_MARK_DOWNLOADED,
|
|
account_id,
|
|
owner_kind,
|
|
owner_id,
|
|
unique_id,
|
|
storage_key,
|
|
file_size,
|
|
)
|