slidge 0.1.2__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

slidge/db/alembic/versions/b33993e87db3_move_everything_to_persistent_db.py

@@ -0,0 +1,214 @@
+"""Move everything to persistent DB
+
+Revision ID: b33993e87db3
+Revises: e91195719c2c
+Create Date: 2024-06-25 16:09:36.663953
+
+"""
+
+import shutil
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+import slidge.db.meta
+from slidge import global_config
+
+# revision identifiers, used by Alembic.
+revision: str = "b33993e87db3"
+down_revision: Union[str, None] = "e91195719c2c"
+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.create_table(
+        "avatar",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("filename", sa.String(), nullable=False),
+        sa.Column("hash", sa.String(), nullable=False),
+        sa.Column("height", sa.Integer(), nullable=False),
+        sa.Column("width", sa.Integer(), nullable=False),
+        sa.Column("legacy_id", sa.String(), nullable=True),
+        sa.Column("url", sa.String(), nullable=True),
+        sa.Column("etag", sa.String(), nullable=True),
+        sa.Column("last_modified", sa.String(), nullable=True),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("filename"),
+        sa.UniqueConstraint("hash"),
+        sa.UniqueConstraint("legacy_id"),
+    )
+    op.create_table(
+        "attachment",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("user_account_id", sa.Integer(), nullable=False),
+        sa.Column("legacy_file_id", sa.String(), nullable=True),
+        sa.Column("url", sa.String(), nullable=False),
+        sa.Column("sims", sa.String(), nullable=True),
+        sa.Column("sfs", sa.String(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["user_account_id"],
+            ["user_account.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(
+        op.f("ix_attachment_legacy_file_id"),
+        "attachment",
+        ["legacy_file_id"],
+        unique=False,
+    )
+    op.create_index(op.f("ix_attachment_url"), "attachment", ["url"], unique=False)
+    op.create_table(
+        "contact",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("user_account_id", sa.Integer(), nullable=False),
+        sa.Column("legacy_id", sa.String(), nullable=False),
+        sa.Column("jid", slidge.db.meta.JIDType(), nullable=False),
+        sa.Column("avatar_id", sa.Integer(), nullable=True),
+        sa.Column("nick", sa.String(), nullable=True),
+        sa.Column("cached_presence", sa.Boolean(), nullable=False),
+        sa.Column("last_seen", sa.DateTime(), nullable=True),
+        sa.Column("ptype", sa.String(), nullable=True),
+        sa.Column("pstatus", sa.String(), nullable=True),
+        sa.Column("pshow", sa.String(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["avatar_id"],
+            ["avatar.id"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_account_id"],
+            ["user_account.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("user_account_id", "jid"),
+        sa.UniqueConstraint("user_account_id", "legacy_id"),
+    )
+    op.create_table(
+        "legacy_ids_multi",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("user_account_id", sa.Integer(), nullable=False),
+        sa.Column("legacy_id", sa.String(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["user_account_id"],
+            ["user_account.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(
+        "legacy_ids_multi_user_account_id_legacy_id",
+        "legacy_ids_multi",
+        ["user_account_id", "legacy_id"],
+        unique=True,
+    )
+    op.create_table(
+        "room",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("user_account_id", sa.Integer(), nullable=False),
+        sa.Column("legacy_id", sa.String(), nullable=False),
+        sa.Column("jid", slidge.db.meta.JIDType(), nullable=False),
+        sa.Column("avatar_id", sa.Integer(), nullable=True),
+        sa.Column("name", sa.String(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["avatar_id"],
+            ["avatar.id"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_account_id"],
+            ["user_account.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("jid"),
+        sa.UniqueConstraint("legacy_id"),
+    )
+    op.create_table(
+        "xmpp_to_legacy_ids",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("user_account_id", sa.Integer(), nullable=False),
+        sa.Column("xmpp_id", sa.String(), nullable=False),
+        sa.Column("legacy_id", sa.String(), nullable=False),
+        sa.Column(
+            "type",
+            sa.Enum("DM", "GROUP_CHAT", "THREAD", name="xmpptolegacyenum"),
+            nullable=False,
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_account_id"],
+            ["user_account.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(
+        "xmpp_legacy",
+        "xmpp_to_legacy_ids",
+        ["user_account_id", "xmpp_id", "legacy_id"],
+        unique=True,
+    )
+    op.create_table(
+        "mam",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("room_id", sa.Integer(), nullable=False),
+        sa.Column("stanza_id", sa.String(), nullable=False),
+        sa.Column("timestamp", sa.DateTime(), nullable=False),
+        sa.Column("author_jid", slidge.db.meta.JIDType(), nullable=False),
+        sa.Column("stanza", sa.String(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["room_id"],
+            ["room.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("room_id", "stanza_id"),
+    )
+    op.create_table(
+        "xmpp_ids_multi",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("user_account_id", sa.Integer(), nullable=False),
+        sa.Column("xmpp_id", sa.String(), nullable=False),
+        sa.Column("legacy_ids_multi_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["legacy_ids_multi_id"],
+            ["legacy_ids_multi.id"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_account_id"],
+            ["user_account.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(
+        "legacy_ids_multi_user_account_id_xmpp_id",
+        "xmpp_ids_multi",
+        ["user_account_id", "xmpp_id"],
+        unique=True,
+    )
+
+    try:
+        shutil.rmtree(global_config.HOME_DIR / "slidge_avatars_v2")
+    except (FileNotFoundError, AttributeError):
+        pass
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(
+        "legacy_ids_multi_user_account_id_xmpp_id", table_name="xmpp_ids_multi"
+    )
+    op.drop_table("xmpp_ids_multi")
+    op.drop_table("mam")
+    op.drop_index("xmpp_legacy", table_name="xmpp_to_legacy_ids")
+    op.drop_table("xmpp_to_legacy_ids")
+    op.drop_table("room")
+    op.drop_index(
+        "legacy_ids_multi_user_account_id_legacy_id", table_name="legacy_ids_multi"
+    )
+    op.drop_table("legacy_ids_multi")
+    op.drop_table("contact")
+    op.drop_index(op.f("ix_attachment_url"), table_name="attachment")
+    op.drop_index(op.f("ix_attachment_legacy_file_id"), table_name="attachment")
+    op.drop_table("attachment")
+    op.drop_table("avatar")
+    # ### end Alembic commands ###

slidge/db/alembic/versions/e91195719c2c_store_users_avatars_persistently.py

@@ -0,0 +1,26 @@
+"""Store users' avatars' hashes persistently
+
+Revision ID: e91195719c2c
+Revises: aa9d82a7f6ef
+Create Date: 2024-06-01 14:14:51.984943
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "e91195719c2c"
+down_revision: Union[str, None] = "aa9d82a7f6ef"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.add_column("user_account", sa.Column("avatar_hash", sa.String(), nullable=True))
+
+
+def downgrade() -> None:
+    op.drop_column("user_account", "avatar_hash")

slidge/db/avatar.py

@@ -0,0 +1,224 @@
+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 Any, Callable, Optional
+
+import aiohttp
+from multidict import CIMultiDictProxy
+from PIL.Image import Image
+from PIL.Image import open as open_image
+from sqlalchemy import select
+
+from slidge.core import config
+from slidge.db.models import Avatar
+from slidge.db.store import AvatarStore
+from slidge.util.types import URL, AvatarType, LegacyFileIdType
+
+
+@dataclass
+class CachedAvatar:
+    pk: int
+    filename: str
+    hash: str
+    height: int
+    width: int
+    root: Path
+    etag: Optional[str] = None
+    last_modified: Optional[str] = None
+
+    @property
+    def data(self):
+        return self.path.read_bytes()
+
+    @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,
+        )
+
+
+class NotModified(Exception):
+    pass
+
+
+class AvatarCache:
+    dir: Path
+    http: aiohttp.ClientSession
+    store: AvatarStore
+    legacy_avatar_type: Callable[[str], Any] = str
+
+    def __init__(self):
+        self._thread_pool = ThreadPoolExecutor(config.AVATAR_RESAMPLING_THREADS)
+
+    def set_dir(self, path: Path):
+        self.dir = path
+        self.dir.mkdir(exist_ok=True)
+
+    def close(self):
+        self._thread_pool.shutdown(cancel_futures=True)
+
+    def __get_http_headers(self, cached: Optional[CachedAvatar | Avatar]):
+        headers = {}
+        if cached and (self.dir / cached.filename).exists():
+            if last_modified := cached.last_modified:
+                headers["If-Modified-Since"] = last_modified
+            if etag := cached.etag:
+                headers["If-None-Match"] = etag
+        return headers
+
+    async def __download(
+        self,
+        url: str,
+        headers: dict[str, str],
+    ) -> tuple[Image, CIMultiDictProxy[str]]:
+        async with self.http.get(url, headers=headers) as response:
+            if response.status == HTTPStatus.NOT_MODIFIED:
+                log.debug("Using avatar cache for %s", url)
+                raise NotModified
+            return (
+                open_image(io.BytesIO(await response.read())),
+                response.headers,
+            )
+
+    async def __is_modified(self, url, headers) -> bool:
+        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)
+        if cached is None:
+            return True
+        headers = self.__get_http_headers(cached)
+        return await self.__is_modified(url, headers)
+
+    def get(self, unique_id: LegacyFileIdType | URL) -> Optional[CachedAvatar]:
+        if isinstance(unique_id, URL):
+            stored = self.store.get_by_url(unique_id)
+        else:
+            stored = self.store.get_by_legacy_id(str(unique_id))
+        if stored is None:
+            return None
+        return CachedAvatar.from_store(stored, self.dir)
+
+    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)
+        raise TypeError("Avatar must be bytes or a Path", avatar)
+
+    async def convert_or_get(
+        self,
+        avatar: AvatarType,
+        unique_id: Optional[LegacyFileIdType],
+    ) -> CachedAvatar:
+        if unique_id is not None:
+            cached = self.get(str(unique_id))
+            if cached is not None:
+                return cached
+
+        if isinstance(avatar, (URL, str)):
+            if unique_id is None:
+                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
+                    return CachedAvatar.from_store(stored, self.dir)
+
+            else:
+                img, _ = await self.__download(avatar, {})
+                response_headers = None
+        else:
+            img = await self._get_image(avatar)
+            response_headers = None
+        with self.store.session() as orm:
+            stored = orm.execute(
+                select(Avatar).where(Avatar.legacy_id == str(unique_id))
+            ).scalar()
+            if stored is not None and stored.url is None:
+                return CachedAvatar.from_store(stored, self.dir)
+
+            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)
+
+            filename = str(uuid.uuid1()) + ".png"
+            file_path = self.dir / filename
+
+            if (
+                not resize
+                and img.format == "PNG"
+                and isinstance(unique_id, str)
+                and (path := Path(unique_id))
+                and path.exists()
+            ):
+                img_bytes = path.read_bytes()
+            else:
+                with io.BytesIO() as f:
+                    img.save(f, format="PNG")
+                    img_bytes = f.getvalue()
+
+            with file_path.open("wb") as file:
+                file.write(img_bytes)
+
+            hash_ = hashlib.sha1(img_bytes).hexdigest()
+
+            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,
+                legacy_id=None if unique_id is None else str(unique_id),
+                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")
+
+            orm.add(stored)
+            orm.commit()
+            return CachedAvatar.from_store(stored, self.dir)
+
+
+avatar_cache = AvatarCache()
+log = logging.getLogger(__name__)
+_download_lock = asyncio.Lock()
+
+__all__ = (
+    "CachedAvatar",
+    "avatar_cache",
+)

slidge/db/meta.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import json
+from typing import Union
+
+import sqlalchemy as sa
+from slixmpp import JID
+
+
+class JIDType(sa.TypeDecorator[JID]):
+    """
+    Custom SQLAlchemy type for JIDs
+    """
+
+    impl = sa.types.TEXT
+    cache_ok = True
+
+    def process_bind_param(self, value: JID | None, dialect: sa.Dialect) -> str | None:
+        if value is None:
+            return value
+        return str(value)
+
+    def process_result_value(
+        self, value: str | None, dialect: sa.Dialect
+    ) -> JID | None:
+        if value is None:
+            return value
+        return JID(value)
+
+
+class JSONEncodedDict(sa.TypeDecorator):
+    """
+    Custom SQLAlchemy type for dictionaries stored as JSON
+
+    Note that mutations of the dictionary are not detected by SQLAlchemy,
+    which is why use ``attributes.flag_modified()`` in ``UserStore.update()``
+    """
+
+    impl = sa.VARCHAR
+
+    cache_ok = True
+
+    def process_bind_param(self, value, dialect):
+        if value is not None:
+            value = json.dumps(value)
+
+        return value
+
+    def process_result_value(self, value, dialect):
+        if value is not None:
+            value = json.loads(value)
+        return value
+
+
+JSONSerializableTypes = Union[str, float, None, "JSONSerializable"]
+JSONSerializable = dict[str, JSONSerializableTypes]
+
+
+class Base(sa.orm.DeclarativeBase):
+    type_annotation_map = {JSONSerializable: JSONEncodedDict, JID: JIDType}
+
+
+def get_engine(path: str) -> sa.Engine:
+    engine = sa.create_engine(path)
+    return engine