144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
import re
|
|
from collections import Counter
|
|
|
|
import asyncpg
|
|
|
|
from utils.read.chats import get_chats
|
|
from utils.read.models import AnnotationView, NoteTag, NoteView, Page
|
|
|
|
_COLS = "id, account_id, chat_id, message_id, text, created_at, updated_at"
|
|
|
|
TAG = re.compile(r"(?<!\w)#(\w+)")
|
|
WORD = re.compile(r"\w+")
|
|
|
|
_NOTES = """
|
|
SELECT a.id, a.account_id, a.chat_id, a.message_id, a.text,
|
|
a.created_at, a.updated_at,
|
|
m.text AS message_text, m.date AS message_date,
|
|
COALESCE(m.has_media, false) AS message_has_media,
|
|
m.deleted_at IS NOT NULL AS message_deleted
|
|
FROM annotations a
|
|
LEFT JOIN LATERAL (
|
|
SELECT text, date, has_media, deleted_at FROM messages
|
|
WHERE account_id = a.account_id AND chat_id = a.chat_id
|
|
AND message_id = a.message_id
|
|
LIMIT 1
|
|
) m ON true
|
|
WHERE a.account_id = $1 AND ($4::text IS NULL OR a.text ~* $4)
|
|
ORDER BY a.created_at DESC, a.id DESC
|
|
LIMIT $2 OFFSET $3
|
|
"""
|
|
|
|
|
|
async def list_annotations(
|
|
pool: asyncpg.Pool,
|
|
account_id: int,
|
|
page: Page,
|
|
*,
|
|
chat_id: int | None = None,
|
|
message_id: int | None = None,
|
|
) -> list[AnnotationView]:
|
|
params: list[object] = [account_id]
|
|
where = "account_id = $1"
|
|
if chat_id is not None:
|
|
params.append(chat_id)
|
|
where += f" AND chat_id = ${len(params)}"
|
|
if message_id is not None:
|
|
params.append(message_id)
|
|
where += f" AND message_id = ${len(params)}"
|
|
params.append(page.capped_limit)
|
|
params.append(page.offset)
|
|
rows = await pool.fetch(
|
|
f"SELECT {_COLS} FROM annotations WHERE {where} " # noqa: S608
|
|
f"ORDER BY created_at DESC LIMIT ${len(params) - 1} OFFSET ${len(params)}",
|
|
*params,
|
|
)
|
|
return [AnnotationView(**dict(row)) for row in rows]
|
|
|
|
|
|
async def get_annotation(
|
|
pool: asyncpg.Pool, annotation_id: int
|
|
) -> AnnotationView | None:
|
|
row = await pool.fetchrow(
|
|
f"SELECT {_COLS} FROM annotations WHERE id = $1", # noqa: S608
|
|
annotation_id,
|
|
)
|
|
return AnnotationView(**dict(row)) if row else None
|
|
|
|
|
|
async def create_annotation(
|
|
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int, text: str
|
|
) -> AnnotationView:
|
|
row = await pool.fetchrow(
|
|
"INSERT INTO annotations (account_id, chat_id, message_id, text) " # noqa: S608
|
|
f"VALUES ($1, $2, $3, $4) RETURNING {_COLS}",
|
|
account_id,
|
|
chat_id,
|
|
message_id,
|
|
text,
|
|
)
|
|
return AnnotationView(**dict(row))
|
|
|
|
|
|
async def update_annotation(
|
|
pool: asyncpg.Pool, annotation_id: int, text: str
|
|
) -> AnnotationView | None:
|
|
row = await pool.fetchrow(
|
|
"UPDATE annotations SET text = $2, updated_at = now() " # noqa: S608
|
|
f"WHERE id = $1 RETURNING {_COLS}",
|
|
annotation_id,
|
|
text,
|
|
)
|
|
return AnnotationView(**dict(row)) if row else None
|
|
|
|
|
|
async def delete_annotation(pool: asyncpg.Pool, annotation_id: int) -> bool:
|
|
result = await pool.execute("DELETE FROM annotations WHERE id = $1", annotation_id)
|
|
return result.endswith("1")
|
|
|
|
|
|
def _tag_pattern(tag: str | None) -> str | None:
|
|
if tag is None:
|
|
return None
|
|
return f"(^|[^[:alnum:]_])#{tag}($|[^[:alnum:]_])"
|
|
|
|
|
|
async def list_notes(
|
|
pool: asyncpg.Pool, account_id: int, page: Page, *, tag: str | None = None
|
|
) -> list[NoteView]:
|
|
if tag is not None and not WORD.fullmatch(tag):
|
|
return []
|
|
rows = await pool.fetch(
|
|
_NOTES, account_id, page.capped_limit, page.offset, _tag_pattern(tag)
|
|
)
|
|
chats = {
|
|
chat.chat_id: chat
|
|
for chat in await get_chats(
|
|
pool, account_id, list({row["chat_id"] for row in rows})
|
|
)
|
|
}
|
|
notes: list[NoteView] = []
|
|
for row in rows:
|
|
chat = chats.get(row["chat_id"])
|
|
notes.append(
|
|
NoteView(
|
|
**dict(row),
|
|
chat_title=chat.title if chat else None,
|
|
chat_has_avatar=chat.has_avatar if chat else False,
|
|
)
|
|
)
|
|
return notes
|
|
|
|
|
|
async def list_tags(pool: asyncpg.Pool, account_id: int) -> list[NoteTag]:
|
|
rows = await pool.fetch(
|
|
"SELECT text FROM annotations WHERE account_id = $1", account_id
|
|
)
|
|
counts: Counter[str] = Counter()
|
|
for row in rows:
|
|
counts.update({tag.lower() for tag in TAG.findall(row["text"])})
|
|
return [
|
|
NoteTag(tag=tag, count=count)
|
|
for tag, count in sorted(counts.items(), key=lambda item: (-item[1], item[0]))
|
|
]
|