mirror of
https://github.com/maubot/rss.git
synced 2025-09-26 19:40:56 -04:00
Switch to asyncpg for database
This commit is contained in:
Tulir Asokan
parent
428b471fec
commit
18ef939a04
7 changed files with 470 additions and 310 deletions
374
rss/db.py
374
rss/db.py
|
|
@ -1,5 +1,5 @@
|
|||
# rss - A maubot plugin to subscribe to RSS/Atom feeds.
|
||||
# Copyright (C) 2020 Tulir Asokan
|
||||
# Copyright (C) 2022 Tulir Asokan
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published by
|
||||
|
|
@ -13,221 +13,207 @@
|
|||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
from typing import Iterable, NamedTuple, List, Optional, Dict, Tuple
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from string import Template
|
||||
|
||||
from sqlalchemy import (Column, String, Integer, DateTime, Text, Boolean, ForeignKey,
|
||||
Table, MetaData,
|
||||
select, and_, true)
|
||||
from sqlalchemy.engine.base import Engine
|
||||
from asyncpg import Record
|
||||
from attr import dataclass
|
||||
import attr
|
||||
|
||||
from mautrix.types import UserID, RoomID
|
||||
|
||||
Subscription = NamedTuple("Subscription", feed_id=int, room_id=RoomID, user_id=UserID,
|
||||
notification_template=Template, send_notice=bool)
|
||||
Feed = NamedTuple("Feed", id=int, url=str, title=str, subtitle=str, link=str, next_retry=int,
|
||||
error_count=int, subscriptions=List[Subscription])
|
||||
Entry = NamedTuple("Entry", feed_id=int, id=str, date=datetime, title=str, summary=str, link=str)
|
||||
from mautrix.types import RoomID, UserID
|
||||
from mautrix.util.async_db import Database, Scheme
|
||||
|
||||
|
||||
class Database:
|
||||
db: Engine
|
||||
feed: Table
|
||||
subscription: Table
|
||||
entry: Table
|
||||
version: Table
|
||||
@dataclass
|
||||
class Subscription:
|
||||
feed_id: int
|
||||
room_id: RoomID
|
||||
user_id: UserID
|
||||
notification_template: Template
|
||||
send_notice: bool
|
||||
|
||||
def __init__(self, db: Engine) -> None:
|
||||
@classmethod
|
||||
def from_row(cls, row: Record | None) -> Subscription | None:
|
||||
if not row:
|
||||
return None
|
||||
feed_id = row["id"]
|
||||
room_id = row["room_id"]
|
||||
user_id = row["user_id"]
|
||||
if not room_id or not user_id:
|
||||
return None
|
||||
send_notice = bool(row["send_notice"])
|
||||
tpl = Template(row["notification_template"])
|
||||
return cls(
|
||||
feed_id=feed_id,
|
||||
room_id=room_id,
|
||||
user_id=user_id,
|
||||
notification_template=tpl,
|
||||
send_notice=send_notice,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Feed:
|
||||
id: int
|
||||
url: str
|
||||
title: str
|
||||
subtitle: str
|
||||
link: str
|
||||
|
||||
next_retry: int = 0
|
||||
error_count: int = 0
|
||||
|
||||
subscriptions: list[Subscription] = attr.ib(factory=lambda: [])
|
||||
|
||||
@classmethod
|
||||
def from_row(cls, row: Record | None) -> Feed | None:
|
||||
if not row:
|
||||
return None
|
||||
data = {**row}
|
||||
data.pop("room_id", None)
|
||||
data.pop("user_id", None)
|
||||
data.pop("send_notice", None)
|
||||
data.pop("notification_template", None)
|
||||
return cls(**data, subscriptions=[])
|
||||
|
||||
|
||||
date_fmt = "%Y-%m-%d %H:%M:%S"
|
||||
date_fmt_microseconds = "%Y-%m-%d %H:%M:%S.%f"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Entry:
|
||||
feed_id: int
|
||||
id: str
|
||||
date: datetime
|
||||
title: str
|
||||
summary: str
|
||||
link: str
|
||||
|
||||
@classmethod
|
||||
def from_row(cls, row: Record | None) -> Entry | None:
|
||||
if not row:
|
||||
return None
|
||||
data = {**row}
|
||||
date = data.pop("date")
|
||||
if not isinstance(date, datetime):
|
||||
try:
|
||||
date = datetime.strptime(date, date_fmt_microseconds if "." in date else date_fmt)
|
||||
except ValueError:
|
||||
date = datetime.now()
|
||||
return cls(**data, date=date)
|
||||
|
||||
|
||||
class DBManager:
|
||||
db: Database
|
||||
|
||||
def __init__(self, db: Database) -> None:
|
||||
self.db = db
|
||||
metadata = MetaData()
|
||||
self.feed = Table("feed", metadata,
|
||||
Column("id", Integer, primary_key=True, autoincrement=True),
|
||||
Column("url", Text, nullable=False, unique=True),
|
||||
Column("title", Text, nullable=False),
|
||||
Column("subtitle", Text, nullable=False),
|
||||
Column("link", Text, nullable=False),
|
||||
Column("next_retry", Integer, nullable=False),
|
||||
Column("error_count", Integer, nullable=False))
|
||||
self.subscription = Table("subscription", metadata,
|
||||
Column("feed_id", Integer, ForeignKey("feed.id"),
|
||||
primary_key=True),
|
||||
Column("room_id", String(255), primary_key=True),
|
||||
Column("user_id", String(255), nullable=False),
|
||||
Column("notification_template", String(255), nullable=True),
|
||||
Column("send_notice", Boolean, nullable=False,
|
||||
server_default=true()))
|
||||
self.entry = Table("entry", metadata,
|
||||
Column("feed_id", Integer, ForeignKey("feed.id"), primary_key=True),
|
||||
Column("id", String(255), primary_key=True),
|
||||
Column("date", DateTime, nullable=False),
|
||||
Column("title", Text, nullable=False),
|
||||
Column("summary", Text, nullable=False),
|
||||
Column("link", Text, nullable=False))
|
||||
self.version = Table("version", metadata,
|
||||
Column("version", Integer, primary_key=True))
|
||||
self.upgrade()
|
||||
|
||||
def upgrade(self) -> None:
|
||||
self.db.execute("CREATE TABLE IF NOT EXISTS version (version INTEGER PRIMARY KEY)")
|
||||
try:
|
||||
version, = next(self.db.execute(select([self.version.c.version])))
|
||||
except (StopIteration, IndexError):
|
||||
version = 0
|
||||
if version == 0:
|
||||
self.db.execute("""CREATE TABLE IF NOT EXISTS feed (
|
||||
id INTEGER NOT NULL,
|
||||
url TEXT NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
subtitle TEXT NOT NULL,
|
||||
link TEXT NOT NULL,
|
||||
PRIMARY KEY (id),
|
||||
UNIQUE (url)
|
||||
)""")
|
||||
self.db.execute("""CREATE TABLE IF NOT EXISTS subscription (
|
||||
feed_id INTEGER NOT NULL,
|
||||
room_id VARCHAR(255) NOT NULL,
|
||||
user_id VARCHAR(255) NOT NULL,
|
||||
notification_template VARCHAR(255),
|
||||
PRIMARY KEY (feed_id, room_id),
|
||||
FOREIGN KEY(feed_id) REFERENCES feed (id)
|
||||
)""")
|
||||
self.db.execute("""CREATE TABLE IF NOT EXISTS entry (
|
||||
feed_id INTEGER NOT NULL,
|
||||
id VARCHAR(255) NOT NULL,
|
||||
date DATETIME NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
summary TEXT NOT NULL,
|
||||
link TEXT NOT NULL,
|
||||
PRIMARY KEY (feed_id, id),
|
||||
FOREIGN KEY(feed_id) REFERENCES feed (id)
|
||||
)""")
|
||||
version = 1
|
||||
if version == 1:
|
||||
self.db.execute("ALTER TABLE subscription ADD COLUMN send_notice BOOLEAN DEFAULT true")
|
||||
version = 2
|
||||
if version == 2:
|
||||
self.db.execute("ALTER TABLE feed ADD COLUMN next_retry BIGINT DEFAULT 0")
|
||||
self.db.execute("ALTER TABLE feed ADD COLUMN error_count BIGINT DEFAULT 0")
|
||||
version = 3
|
||||
self.db.execute(self.version.delete())
|
||||
self.db.execute(self.version.insert().values(version=version))
|
||||
|
||||
def get_feeds(self) -> Iterable[Feed]:
|
||||
rows = self.db.execute(select([self.feed,
|
||||
self.subscription.c.room_id,
|
||||
self.subscription.c.user_id,
|
||||
self.subscription.c.notification_template,
|
||||
self.subscription.c.send_notice])
|
||||
.where(self.subscription.c.feed_id == self.feed.c.id))
|
||||
map: Dict[int, Feed] = {}
|
||||
async def get_feeds(self) -> list[Feed]:
|
||||
q = """
|
||||
SELECT id, url, title, subtitle, link, next_retry, error_count,
|
||||
room_id, user_id, notification_template, send_notice
|
||||
FROM feed INNER JOIN subscription ON feed.id = subscription.feed_id
|
||||
"""
|
||||
rows = await self.db.fetch(q)
|
||||
feeds: dict[int, Feed] = {}
|
||||
for row in rows:
|
||||
(feed_id, url, title, subtitle, link, next_retry, error_count,
|
||||
room_id, user_id, notification_template, send_notice) = row
|
||||
map.setdefault(feed_id, Feed(feed_id, url, title, subtitle, link, next_retry,
|
||||
error_count, subscriptions=[]))
|
||||
map[feed_id].subscriptions.append(
|
||||
Subscription(feed_id=feed_id, room_id=room_id, user_id=user_id,
|
||||
notification_template=Template(notification_template),
|
||||
send_notice=send_notice))
|
||||
return map.values()
|
||||
try:
|
||||
feed = feeds[row["id"]]
|
||||
except KeyError:
|
||||
feed = feeds[row["id"]] = Feed.from_row(row)
|
||||
feed.subscriptions.append(Subscription.from_row(row))
|
||||
return list(feeds.values())
|
||||
|
||||
def get_feeds_by_room(self, room_id: RoomID) -> Iterable[Tuple[Feed, UserID]]:
|
||||
return ((Feed(feed_id, url, title, subtitle, link, next_retry, error_count,
|
||||
subscriptions=[]),
|
||||
user_id)
|
||||
for (feed_id, url, title, subtitle, link, next_retry, error_count, user_id) in
|
||||
self.db.execute(select([self.feed, self.subscription.c.user_id])
|
||||
.where(and_(self.subscription.c.room_id == room_id,
|
||||
self.subscription.c.feed_id == self.feed.c.id))))
|
||||
async def get_feeds_by_room(self, room_id: RoomID) -> list[tuple[Feed, UserID]]:
|
||||
q = """
|
||||
SELECT id, url, title, subtitle, link, next_retry, error_count, user_id FROM feed
|
||||
INNER JOIN subscription ON feed.id = subscription.feed_id AND subscription.room_id = $1
|
||||
"""
|
||||
rows = await self.db.fetch(q, room_id)
|
||||
return [(Feed.from_row(row), row["user_id"]) for row in rows]
|
||||
|
||||
def get_rooms_by_feed(self, feed_id: int) -> Iterable[RoomID]:
|
||||
return (row[0] for row in
|
||||
self.db.execute(select([self.subscription.c.room_id])
|
||||
.where(self.subscription.c.feed_id == feed_id)))
|
||||
async def get_entries(self, feed_id: int) -> list[Entry]:
|
||||
q = "SELECT feed_id, id, date, title, summary, link FROM entry WHERE feed_id = $1"
|
||||
return [Entry.from_row(row) for row in await self.db.fetch(q, feed_id)]
|
||||
|
||||
def get_entries(self, feed_id: int) -> Iterable[Entry]:
|
||||
return (Entry(*row) for row in
|
||||
self.db.execute(select([self.entry]).where(self.entry.c.feed_id == feed_id)))
|
||||
|
||||
def add_entries(self, entries: Iterable[Entry], override_feed_id: Optional[int] = None) -> None:
|
||||
async def add_entries(self, entries: list[Entry], override_feed_id: int | None = None) -> None:
|
||||
if not entries:
|
||||
return
|
||||
entries = [entry._asdict() for entry in entries]
|
||||
if override_feed_id is not None:
|
||||
if override_feed_id:
|
||||
for entry in entries:
|
||||
entry["feed_id"] = override_feed_id
|
||||
self.db.execute(self.entry.insert(), entries)
|
||||
entry.feed_id = override_feed_id
|
||||
records = [attr.astuple(entry) for entry in entries]
|
||||
columns = ("feed_id", "id", "date", "title", "summary", "link")
|
||||
async with self.db.acquire() as conn:
|
||||
if self.db.scheme == Scheme.POSTGRES:
|
||||
await conn.copy_records_to_table("entry", records=records, columns=columns)
|
||||
else:
|
||||
q = (
|
||||
"INSERT INTO entry (feed_id, id, date, title, summary, link) "
|
||||
"VALUES ($1, $2, $3, $4, $5, $6)"
|
||||
)
|
||||
await conn.executemany(q, records)
|
||||
|
||||
def get_feed_by_url(self, url: str) -> Optional[Feed]:
|
||||
rows = self.db.execute(select([self.feed]).where(self.feed.c.url == url))
|
||||
try:
|
||||
row = next(rows)
|
||||
return Feed(*row, subscriptions=[])
|
||||
except (ValueError, StopIteration):
|
||||
return None
|
||||
async def get_feed_by_url(self, url: str) -> Feed | None:
|
||||
q = "SELECT id, url, title, subtitle, link, next_retry, error_count FROM feed WHERE url=$1"
|
||||
return Feed.from_row(await self.db.fetchrow(q, url))
|
||||
|
||||
def get_feed_by_id(self, feed_id: int) -> Optional[Feed]:
|
||||
rows = self.db.execute(select([self.feed]).where(self.feed.c.id == feed_id))
|
||||
try:
|
||||
row = next(rows)
|
||||
return Feed(*row, subscriptions=[])
|
||||
except (ValueError, StopIteration):
|
||||
return None
|
||||
async def get_subscription(
|
||||
self, feed_id: int, room_id: RoomID
|
||||
) -> tuple[Subscription | None, Feed | None]:
|
||||
q = """
|
||||
SELECT id, url, title, subtitle, link, next_retry, error_count,
|
||||
room_id, user_id, notification_template, send_notice
|
||||
FROM feed LEFT JOIN subscription ON feed.id = subscription.feed_id AND room_id = $2
|
||||
WHERE feed.id = $1
|
||||
"""
|
||||
row = await self.db.fetchrow(q, feed_id, room_id)
|
||||
return Subscription.from_row(row), Feed.from_row(row)
|
||||
|
||||
def get_subscription(self, feed_id: int, room_id: RoomID) -> Tuple[Optional[Subscription],
|
||||
Optional[Feed]]:
|
||||
tbl = self.subscription
|
||||
rows = self.db.execute(select([self.feed, tbl.c.room_id, tbl.c.user_id,
|
||||
tbl.c.notification_template, tbl.c.send_notice])
|
||||
.where(and_(tbl.c.feed_id == feed_id, tbl.c.room_id == room_id,
|
||||
self.feed.c.id == feed_id)))
|
||||
try:
|
||||
(feed_id, url, title, subtitle, link, next_retry, error_count,
|
||||
room_id, user_id, template, send_notice) = next(rows)
|
||||
notification_template = Template(template)
|
||||
return (Subscription(feed_id, room_id, user_id, notification_template, send_notice)
|
||||
if room_id else None,
|
||||
Feed(feed_id, url, title, subtitle, link, next_retry, error_count, []))
|
||||
except (ValueError, StopIteration):
|
||||
return None, None
|
||||
async def update_room_id(self, old: RoomID, new: RoomID) -> None:
|
||||
await self.db.execute("UPDATE subscription SET room_id = $1 WHERE room_id = $2", new, old)
|
||||
|
||||
def update_room_id(self, old: RoomID, new: RoomID) -> None:
|
||||
self.db.execute(self.subscription.update()
|
||||
.where(self.subscription.c.room_id == old)
|
||||
.values(room_id=new))
|
||||
async def create_feed(self, info: Feed) -> Feed:
|
||||
q = (
|
||||
"INSERT INTO feed (url, title, subtitle, link, next_retry) "
|
||||
"VALUES ($1, $2, $3, $4, $5) RETURNING (id)"
|
||||
)
|
||||
info.id = await self.db.fetchval(
|
||||
q, info.url, info.title, info.subtitle, info.link, info.next_retry
|
||||
)
|
||||
return info
|
||||
|
||||
def create_feed(self, info: Feed) -> Feed:
|
||||
res = self.db.execute(self.feed.insert().values(url=info.url, title=info.title,
|
||||
subtitle=info.subtitle, link=info.link,
|
||||
next_retry=info.next_retry))
|
||||
return Feed(id=res.inserted_primary_key[0], url=info.url, title=info.title,
|
||||
subtitle=info.subtitle, link=info.link, next_retry=info.next_retry,
|
||||
error_count=info.error_count, subscriptions=[])
|
||||
async def set_backoff(self, info: Feed, error_count: int, next_retry: int) -> None:
|
||||
q = "UPDATE feed SET error_count = $2, next_retry = $3 WHERE id = $1"
|
||||
await self.db.execute(q, info.id, error_count, next_retry)
|
||||
|
||||
def set_backoff(self, info: Feed, error_count: int, next_retry: int) -> None:
|
||||
self.db.execute(self.feed.update()
|
||||
.where(self.feed.c.id == info.id)
|
||||
.values(error_count=error_count, next_retry=next_retry))
|
||||
async def subscribe(
|
||||
self,
|
||||
feed_id: int,
|
||||
room_id: RoomID,
|
||||
user_id: UserID,
|
||||
template: str | None = None,
|
||||
send_notice: bool = True,
|
||||
) -> None:
|
||||
q = """
|
||||
INSERT INTO subscription (feed_id, room_id, user_id, notification_template, send_notice)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
"""
|
||||
template = template or "New post in $feed_title: [$title]($link)"
|
||||
await self.db.execute(q, feed_id, room_id, user_id, template, send_notice)
|
||||
|
||||
def subscribe(self, feed_id: int, room_id: RoomID, user_id: UserID) -> None:
|
||||
self.db.execute(self.subscription.insert().values(
|
||||
feed_id=feed_id, room_id=room_id, user_id=user_id,
|
||||
notification_template="New post in $feed_title: [$title]($link)"))
|
||||
async def unsubscribe(self, feed_id: int, room_id: RoomID) -> None:
|
||||
q = "DELETE FROM subscription WHERE feed_id = $1 AND room_id = $2"
|
||||
await self.db.execute(q, feed_id, room_id)
|
||||
|
||||
def unsubscribe(self, feed_id: int, room_id: RoomID) -> None:
|
||||
tbl = self.subscription
|
||||
self.db.execute(tbl.delete().where(and_(tbl.c.feed_id == feed_id,
|
||||
tbl.c.room_id == room_id)))
|
||||
async def update_template(self, feed_id: int, room_id: RoomID, template: str) -> None:
|
||||
q = "UPDATE subscription SET notification_template=$3 WHERE feed_id=$1 AND room_id=$2"
|
||||
await self.db.execute(q, feed_id, room_id, template)
|
||||
|
||||
def update_template(self, feed_id: int, room_id: RoomID, template: str) -> None:
|
||||
tbl = self.subscription
|
||||
self.db.execute(tbl.update()
|
||||
.where(and_(tbl.c.feed_id == feed_id, tbl.c.room_id == room_id))
|
||||
.values(notification_template=template))
|
||||
|
||||
def set_send_notice(self, feed_id: int, room_id: RoomID, send_notice: bool) -> None:
|
||||
tbl = self.subscription
|
||||
self.db.execute(tbl.update()
|
||||
.where(and_(tbl.c.feed_id == feed_id, tbl.c.room_id == room_id))
|
||||
.values(send_notice=send_notice))
|
||||
async def set_send_notice(self, feed_id: int, room_id: RoomID, send_notice: bool) -> None:
|
||||
q = "UPDATE subscription SET send_notice=$3 WHERE feed_id=$1 AND room_id=$2"
|
||||
await self.db.execute(q, feed_id, room_id, send_notice)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue