slidge 0.2.12__py3-none-any.whl → 0.3.0a0__py3-none-any.whl

slidge/db/alembic/versions/58b98dacf819_refactor.py

@@ -0,0 +1,118 @@
+"""Add unique constraints that we should have added before
+
+Revision ID: 58b98dacf819
+Revises: 54ce3cde350c
+Create Date: 2025-04-23 22:55:11.731158
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "58b98dacf819"
+down_revision: Union[str, None] = "54ce3cde350c"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("avatar", schema=None) as batch_op:
+        batch_op.drop_constraint("avatar_unique_legacy_id", type_="unique")
+        batch_op.create_unique_constraint(batch_op.f("uq_avatar_hash"), ["hash"])
+        batch_op.create_unique_constraint(
+            batch_op.f("uq_avatar_legacy_id"), ["legacy_id"]
+        )
+
+    with op.batch_alter_table("bob", schema=None) as batch_op:
+        batch_op.create_unique_constraint(batch_op.f("uq_bob_sha_1"), ["sha_1"])
+        batch_op.create_unique_constraint(batch_op.f("uq_bob_sha_256"), ["sha_256"])
+        batch_op.create_unique_constraint(batch_op.f("uq_bob_sha_512"), ["sha_512"])
+
+    with op.batch_alter_table("contact", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            batch_op.f("uq_contact_user_account_id"), ["user_account_id", "legacy_id"]
+        )
+
+    with op.batch_alter_table("contact_sent", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            batch_op.f("uq_contact_sent_contact_id"), ["contact_id", "msg_id"]
+        )
+
+    with op.batch_alter_table("hat", schema=None) as batch_op:
+        batch_op.create_unique_constraint(batch_op.f("uq_hat_title"), ["title", "uri"])
+
+    with op.batch_alter_table("mam", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            batch_op.f("uq_mam_room_id"), ["room_id", "stanza_id"]
+        )
+
+    with op.batch_alter_table("participant", schema=None) as batch_op:
+        batch_op.alter_column("resource", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("nickname", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column(
+            "nickname_no_illegal", existing_type=sa.VARCHAR(), nullable=False
+        )
+        batch_op.create_unique_constraint(
+            batch_op.f("uq_participant_room_id"), ["room_id", "resource"]
+        )
+
+    with op.batch_alter_table("room", schema=None) as batch_op:
+        batch_op.alter_column(
+            "muc_type", existing_type=sa.VARCHAR(length=21), nullable=False
+        )
+
+    with op.batch_alter_table("user_account", schema=None) as batch_op:
+        batch_op.create_unique_constraint(batch_op.f("uq_user_account_jid"), ["jid"])
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("user_account", schema=None) as batch_op:
+        batch_op.drop_constraint(batch_op.f("uq_user_account_jid"), type_="unique")
+
+    with op.batch_alter_table("room", schema=None) as batch_op:
+        batch_op.alter_column(
+            "muc_type", existing_type=sa.VARCHAR(length=21), nullable=True
+        )
+
+    with op.batch_alter_table("participant", schema=None) as batch_op:
+        batch_op.drop_constraint(batch_op.f("uq_participant_room_id"), type_="unique")
+        batch_op.alter_column(
+            "nickname_no_illegal", existing_type=sa.VARCHAR(), nullable=True
+        )
+        batch_op.alter_column("nickname", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("resource", existing_type=sa.VARCHAR(), nullable=True)
+
+    with op.batch_alter_table("mam", schema=None) as batch_op:
+        batch_op.drop_constraint(batch_op.f("uq_mam_room_id"), type_="unique")
+
+    with op.batch_alter_table("hat", schema=None) as batch_op:
+        batch_op.drop_constraint(batch_op.f("uq_hat_title"), type_="unique")
+
+    with op.batch_alter_table("contact_sent", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            batch_op.f("uq_contact_sent_contact_id"), type_="unique"
+        )
+
+    with op.batch_alter_table("contact", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            batch_op.f("uq_contact_user_account_id"), type_="unique"
+        )
+
+    with op.batch_alter_table("bob", schema=None) as batch_op:
+        batch_op.drop_constraint(batch_op.f("uq_bob_sha_512"), type_="unique")
+        batch_op.drop_constraint(batch_op.f("uq_bob_sha_256"), type_="unique")
+        batch_op.drop_constraint(batch_op.f("uq_bob_sha_1"), type_="unique")
+
+    with op.batch_alter_table("avatar", schema=None) as batch_op:
+        batch_op.drop_constraint(batch_op.f("uq_avatar_legacy_id"), type_="unique")
+        batch_op.drop_constraint(batch_op.f("uq_avatar_hash"), type_="unique")
+        batch_op.create_unique_constraint("avatar_unique_legacy_id", ["legacy_id"])
+
+    # ### end Alembic commands ###

slidge/db/alembic/versions/75a62b74b239_ditch_hats_table.py

@@ -0,0 +1,74 @@
+"""Ditch hats table
+
+Revision ID: 75a62b74b239
+Revises: 0337c90c0b96
+Create Date: 2025-05-02 14:24:26.141034
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "75a62b74b239"
+down_revision: Union[str, None] = "0337c90c0b96"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("participant_hats")
+    op.drop_table("hat")
+    with op.batch_alter_table("contact", schema=None) as batch_op:
+        batch_op.drop_constraint("uq_contact_user_account_id", type_="unique")
+        batch_op.create_unique_constraint(
+            batch_op.f("uq_contact_user_account_id"), ["user_account_id", "jid"]
+        )
+
+    with op.batch_alter_table("participant", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("hats", sa.JSON(), nullable=False, server_default="[]")
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("participant", schema=None) as batch_op:
+        batch_op.drop_column("hats")
+
+    with op.batch_alter_table("contact", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            batch_op.f("uq_contact_user_account_id"), type_="unique"
+        )
+        batch_op.create_unique_constraint(
+            "uq_contact_user_account_id", ["user_account_id", "legacy_id"]
+        )
+
+    op.create_table(
+        "hat",
+        sa.Column("id", sa.INTEGER(), nullable=False),
+        sa.Column("title", sa.VARCHAR(), nullable=False),
+        sa.Column("uri", sa.VARCHAR(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("title", "uri", name="uq_hat_title"),
+    )
+    op.create_table(
+        "participant_hats",
+        sa.Column("participant_id", sa.INTEGER(), nullable=False),
+        sa.Column("hat_id", sa.INTEGER(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["hat_id"],
+            ["hat.id"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["participant_id"],
+            ["participant.id"],
+        ),
+        sa.PrimaryKeyConstraint("participant_id", "hat_id"),
+    )
+    # ### end Alembic commands ###

slidge/db/avatar.py

@@ -2,9 +2,7 @@ import asyncio
 import hashlib
 import io
 import logging
-import uuid
 from concurrent.futures import ThreadPoolExecutor
-from dataclasses import dataclass
 from http import HTTPStatus
 from pathlib import Path
 from typing import Optional
@@ -14,24 +12,42 @@ from multidict import CIMultiDictProxy
 from PIL.Image import Image
 from PIL.Image import open as open_image
 from sqlalchemy import select
-from sqlalchemy.orm.exc import DetachedInstanceError
 
-from slidge.core import config
-from slidge.db.models import Avatar
-from slidge.db.store import AvatarStore
-from slidge.util.types import URL, AvatarType
+from ..core import config
+from ..util.lock import NamedLockMixin
+from ..util.types import Avatar as AvatarType
+from .models import Avatar
+from .store import AvatarStore
 
 
-@dataclass
 class CachedAvatar:
-    pk: int
-    filename: str
-    hash: str
-    height: int
-    width: int
-    root: Path
-    etag: Optional[str] = None
-    last_modified: Optional[str] = None
+    def __init__(self, stored: Avatar, root_dir: Path) -> None:
+        self.stored = stored
+        self._root = root_dir
+
+    @property
+    def pk(self) -> int | None:
+        return self.stored.id
+
+    @property
+    def hash(self) -> str:
+        return self.stored.hash
+
+    @property
+    def height(self) -> int:
+        return self.stored.height
+
+    @property
+    def width(self) -> int:
+        return self.stored.width
+
+    @property
+    def etag(self) -> str | None:
+        return self.stored.etag
+
+    @property
+    def last_modified(self) -> str | None:
+        return self.stored.last_modified
 
     @property
     def data(self):
@@ -39,40 +55,32 @@ class CachedAvatar:
 
     @property
     def path(self):
-        return self.root / self.filename
-
-    @staticmethod
-    def from_store(stored: Avatar, root_dir: Path):
-        return CachedAvatar(
-            pk=stored.id,
-            filename=stored.filename,
-            hash=stored.hash,
-            height=stored.height,
-            width=stored.width,
-            etag=stored.etag,
-            root=root_dir,
-            last_modified=stored.last_modified,
-        )
+        return (self._root / self.hash).with_suffix(".png")
 
 
 class NotModified(Exception):
     pass
 
 
-class AvatarCache:
+class AvatarCache(NamedLockMixin):
     dir: Path
     http: aiohttp.ClientSession
     store: AvatarStore
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._thread_pool = ThreadPoolExecutor(config.AVATAR_RESAMPLING_THREADS)
+        super().__init__()
+
+    def get(self, stored: Avatar) -> CachedAvatar:
+        return CachedAvatar(stored, self.dir)
 
-    def set_dir(self, path: Path):
+    def set_dir(self, path: Path) -> None:
         self.dir = path
         self.dir.mkdir(exist_ok=True)
-        with self.store.session():
-            for stored in self.store.get_all():
-                avatar = CachedAvatar.from_store(stored, root_dir=path)
+        log.debug("Checking avatar files")
+        with self.store.session(expire_on_commit=False) as orm:
+            for stored in orm.query(Avatar).all():
+                avatar = CachedAvatar(stored, path)
                 if avatar.path.exists():
                     continue
                 log.warning(
@@ -80,14 +88,15 @@ class AvatarCache:
                     avatar.hash,
                     avatar.path,
                 )
-                self.store.delete_by_pk(stored.id)
+                orm.delete(stored)
+            orm.commit()
 
-    def close(self):
+    def close(self) -> None:
         self._thread_pool.shutdown(cancel_futures=True)
 
-    def __get_http_headers(self, cached: Optional[CachedAvatar | Avatar]):
+    def __get_http_headers(self, cached: Optional[CachedAvatar | Avatar] = None):
         headers = {}
-        if cached and (self.dir / cached.filename).exists():
+        if cached and (self.dir / cached.hash).with_suffix(".png").exists():
             if last_modified := cached.last_modified:
                 headers["If-Modified-Since"] = last_modified
             if etag := cached.etag:
@@ -113,105 +122,127 @@ class AvatarCache:
         async with self.http.head(url, headers=headers) as response:
             return response.status != HTTPStatus.NOT_MODIFIED
 
-    async def url_modified(self, url: URL) -> bool:
-        cached = self.store.get_by_url(url)
+    async def url_modified(self, url: str) -> bool:
+        with self.store.session() as orm:
+            cached = orm.query(Avatar).filter_by(url=url).one_or_none()
         if cached is None:
             return True
         headers = self.__get_http_headers(cached)
         return await self.__is_modified(url, headers)
 
-    def get_by_pk(self, pk: int) -> CachedAvatar:
-        stored = self.store.get_by_pk(pk)
-        assert stored is not None
-        return CachedAvatar.from_store(stored, self.dir)
-
     @staticmethod
     async def _get_image(avatar: AvatarType) -> Image:
-        if isinstance(avatar, bytes):
-            return open_image(io.BytesIO(avatar))
-        elif isinstance(avatar, Path):
-            return open_image(avatar)
+        if avatar.data is not None:
+            return open_image(io.BytesIO(avatar.data))
+        elif avatar.path is not None:
+            return open_image(avatar.path)
         raise TypeError("Avatar must be bytes or a Path", avatar)
 
     async def convert_or_get(self, avatar: AvatarType) -> CachedAvatar:
-        if isinstance(avatar, (URL, str)):
-            with self.store.session():
-                stored = self.store.get_by_url(avatar)
-                try:
-                    img, response_headers = await self.__download(
-                        avatar, self.__get_http_headers(stored)
-                    )
-                except NotModified:
-                    assert stored is not None
-                    try:
-                        return CachedAvatar.from_store(stored, self.dir)
-                    except DetachedInstanceError:
-                        # This is an awful hack to prevent errors on startup under certain conditions,
-                        # because we basically misused SQLAlchemy pretty bad in slidge.db.store.EngineMixin.session().
-                        # cf https://codeberg.org/slidge/slidge/issues/36
-                        # and https://docs.sqlalchemy.org/en/20/orm/session_basics.html#session-faq-threadsafe
-                        # It may be related to threads as we only have reports of this for slidge-whatsapp and skidge
-                        # which are the only implementations that uses threads.
-                        # In any case, a proper fix implies a major refactoring in which we spawn and close SQLAlchemy
-                        # "ORM Session"s with a reasonable, well-thought lifetime, instead of how we do it now, where we
-                        # basically just brute-forced our way into having something usable but with poor performance.
-                        # Databases, asyncio, and concurrency in general are hard… :(
-                        # Oh, and getting rid of the convoluted mess that this giant method is would also probably
-                        # be a good idea.
-                        stored = self.store.get_by_url(avatar)
-                        assert stored is not None
-                        return CachedAvatar.from_store(stored, self.dir)
-        else:
-            img = await self._get_image(avatar)
-            response_headers = None
-        with self.store.session() as orm:
-            resize = (size := config.AVATAR_SIZE) and any(x > size for x in img.size)
-            if resize:
-                await asyncio.get_event_loop().run_in_executor(
-                    self._thread_pool, img.thumbnail, (size, size)
+        if avatar.unique_id is not None:
+            with self.store.session() as orm:
+                stored = (
+                    orm.query(Avatar)
+                    .filter_by(legacy_id=str(avatar.unique_id))
+                    .one_or_none()
                 )
-                log.debug("Resampled image to %s", img.size)
-
-            filename = str(uuid.uuid1()) + ".png"
-            file_path = self.dir / filename
-
-            if (
-                not resize
-                and img.format == "PNG"
-                and isinstance(avatar, (str, Path))
-                and (path := Path(avatar))
-                and path.exists()
-            ):
-                img_bytes = path.read_bytes()
-            else:
-                with io.BytesIO() as f:
-                    img.save(f, format="PNG")
-                    img_bytes = f.getvalue()
+                if stored is not None:
+                    return self.get(stored)
+
+        if avatar.url is not None:
+            return await self.__convert_url(avatar)
+
+        return await self.convert(avatar, await self._get_image(avatar))
+
+    async def __convert_url(self, avatar: AvatarType) -> CachedAvatar:
+        assert avatar.url is not None
+        async with self.lock(avatar.unique_id or avatar.url):
+            with self.store.session() as orm:
+                if avatar.unique_id is None:
+                    stored = orm.query(Avatar).filter_by(url=avatar.url).one_or_none()
+                else:
+                    stored = (
+                        orm.query(Avatar)
+                        .filter_by(legacy_id=str(avatar.unique_id))
+                        .one_or_none()
+                    )
+                    if stored is not None:
+                        return self.get(stored)
 
-            with file_path.open("wb") as file:
-                file.write(img_bytes)
+            try:
+                img, response_headers = await self.__download(
+                    avatar.url, self.__get_http_headers(stored)
+                )
+            except NotModified:
+                assert stored is not None
+                return self.get(stored)
 
-            hash_ = hashlib.sha1(img_bytes).hexdigest()
+            return await self.convert(avatar, img, response_headers)
 
+    async def convert(
+        self,
+        avatar: AvatarType,
+        img: Image,
+        response_headers: CIMultiDictProxy[str] | None = None,
+    ) -> CachedAvatar:
+        resize = (size := config.AVATAR_SIZE) and any(x > size for x in img.size)
+        if resize:
+            await asyncio.get_event_loop().run_in_executor(
+                self._thread_pool, img.thumbnail, (size, size)
+            )
+            log.debug("Resampled image to %s", img.size)
+
+        if (
+            not resize
+            and img.format == "PNG"
+            and avatar.path is not None
+            and avatar.path.exists()
+        ):
+            img_bytes = avatar.path.read_bytes()
+        else:
+            with io.BytesIO() as f:
+                img.save(f, format="PNG")
+                img_bytes = f.getvalue()
+
+        hash_ = hashlib.sha1(img_bytes).hexdigest()
+        file_path = (self.dir / hash_).with_suffix(".png")
+        if file_path.exists():
+            log.warning("Overwriting %s", file_path)
+        with file_path.open("wb") as file:
+            file.write(img_bytes)
+
+        with self.store.session(expire_on_commit=False) as orm:
             stored = orm.execute(select(Avatar).where(Avatar.hash == hash_)).scalar()
 
             if stored is not None:
-                return CachedAvatar.from_store(stored, self.dir)
-
-            stored = Avatar(
-                filename=filename,
-                hash=hash_,
-                height=img.height,
-                width=img.width,
-                url=avatar if isinstance(avatar, (URL, str)) else None,
-            )
-            if response_headers:
-                stored.etag = response_headers.get("etag")
-                stored.last_modified = response_headers.get("last-modified")
+                if avatar.unique_id is not None:
+                    if str(avatar.unique_id) != stored.legacy_id:
+                        log.warning(
+                            "Updating the 'unique' ID of an avatar, was '%s', is now '%s'",
+                            stored.legacy_id,
+                            avatar.unique_id,
+                        )
+                        stored.legacy_id = str(avatar.unique_id)
+                        orm.add(stored)
+                        orm.commit()
+
+                return self.get(stored)
+
+        stored = Avatar(
+            hash=hash_,
+            height=img.height,
+            width=img.width,
+            url=avatar.url,
+            legacy_id=avatar.unique_id,
+        )
+        if response_headers:
+            stored.etag = response_headers.get("etag")
+            stored.last_modified = response_headers.get("last-modified")
 
+        with self.store.session(expire_on_commit=False) as orm:
             orm.add(stored)
             orm.commit()
-            return CachedAvatar.from_store(stored, self.dir)
+            return self.get(stored)
 
 
 avatar_cache = AvatarCache()

slidge/db/meta.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import json
-from typing import Union
+from typing import Any, Union
 
 import sqlalchemy as sa
 from slixmpp import JID
+from sqlalchemy import Dialect
 
 
 class JIDType(sa.TypeDecorator[JID]):
@@ -28,7 +29,11 @@ class JIDType(sa.TypeDecorator[JID]):
         return JID(value)
 
 
-class JSONEncodedDict(sa.TypeDecorator):
+JSONSerializableTypes = Union[str, float, None, "JSONSerializable"]
+JSONSerializable = dict[str, JSONSerializableTypes]
+
+
+class JSONEncodedDict(sa.TypeDecorator[JSONSerializable]):
     """
     Custom SQLAlchemy type for dictionaries stored as JSON
 
@@ -40,33 +45,39 @@ class JSONEncodedDict(sa.TypeDecorator):
 
     cache_ok = True
 
-    def process_bind_param(self, value, dialect):
-        if value is not None:
-            value = json.dumps(value)
+    def process_bind_param(
+        self, value: JSONSerializable | None, dialect: Dialect
+    ) -> str | None:
+        if value is None:
+            return None
+        return json.dumps(value)
 
-        return value
+    def process_result_value(
+        self, value: Any | None, dialect: Dialect
+    ) -> JSONSerializable | None:
+        if value is None:
+            return None
+        return json.loads(value)  # type:ignore
 
-    def process_result_value(self, value, dialect):
-        if value is not None:
-            value = json.loads(value)
-        return value
 
+class Base(sa.orm.DeclarativeBase):
+    type_annotation_map = {JSONSerializable: JSONEncodedDict, JID: JIDType}
 
-JSONSerializableTypes = Union[str, float, None, "JSONSerializable"]
-JSONSerializable = dict[str, JSONSerializableTypes]
 
+Base.metadata.naming_convention = {
+    "ix": "ix_%(column_0_label)s",
+    "uq": "uq_%(table_name)s_%(column_0_name)s",
+    "ck": "ck_%(table_name)s_`%(constraint_name)s`",
+    "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
+    "pk": "pk_%(table_name)s",
+}
 
-class Base(sa.orm.DeclarativeBase):
-    type_annotation_map = {JSONSerializable: JSONEncodedDict, JID: JIDType}
-    naming_convention = {
-        "ix": "ix_%(column_0_label)s",
-        "uq": "uq_%(table_name)s_%(column_0_name)s",
-        "ck": "ck_%(table_name)s_`%(constraint_name)s`",
-        "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
-        "pk": "pk_%(table_name)s",
-    }
 
+def get_engine(path: str, echo: bool = False) -> sa.Engine:
+    from sqlalchemy import log as sqlalchemy_log
 
-def get_engine(path: str) -> sa.Engine:
     engine = sa.create_engine(path)
+    if echo:
+        sqlalchemy_log._add_default_handler = lambda x: None  # type:ignore
+        engine.echo = True
     return engine