diff --git a/wrolpi/files/api.py b/wrolpi/files/api.py index 47b86c4a..331be78a 100644 --- a/wrolpi/files/api.py +++ b/wrolpi/files/api.py @@ -305,11 +305,6 @@ async def post_untag_file_group(_, body: schema.TagFileGroupPost): return response.empty(HTTPStatus.NO_CONTENT) -# { -# '/media/directory/sub-dir/the-file-name.suffix': 2, # The chunk number we will receive next. -# } - - @files_bp.post('/upload') async def post_upload(request: Request): """Accepts a multipart/form-data request to upload a single file. diff --git a/wrolpi/files/models.py b/wrolpi/files/models.py index 49e027b2..bb027335 100644 --- a/wrolpi/files/models.py +++ b/wrolpi/files/models.py @@ -3,6 +3,7 @@ import pathlib import shutil import urllib.parse from datetime import datetime +from functools import singledispatchmethod from typing import List, Type, Optional, Iterable from sqlalchemy import Column, String, Computed, BigInteger, Boolean @@ -17,7 +18,7 @@ from wrolpi.downloader import Download from wrolpi.errors import FileGroupIsTagged, UnknownFile from wrolpi.files import indexers from wrolpi.media_path import MediaPathType -from wrolpi.tags import Tag, TagFile +from wrolpi.tags import Tag, TagFile, save_tags_config, sync_tags_directory from wrolpi.vars import PYTEST logger = logger.getChild(__name__) @@ -144,12 +145,56 @@ class FileGroup(ModelHelper, Base): self.viewed = now() @optional_session - def add_tag(self, tag_id_or_name: int | str, session: Session = None) -> TagFile: - return Tag.tag_file_group(self.id, tag_id_or_name, session) + def set_tags(self, tag_names_or_ids: Iterable[str | int], session: Session = None): + """Insert or Delete TagFiles as necessary to match the provided tags for this FileGroup.""" + tag_names_or_ids = list(tag_names_or_ids) + if tag_names_or_ids and isinstance(tag_names_or_ids[0], str): + tag_ids = {Tag.get_id_by_name(i) for i in tag_names_or_ids} + else: + tag_ids = set(tag_names_or_ids) + existing_tag_ids = {i.tag_id for i in self.tag_files} + if new_tag_ids := (tag_ids - existing_tag_ids): + for tag_id in new_tag_ids: + tag_file = TagFile(file_group_id=self.id, tag_id=tag_id) + self.tag_files.append(tag_file) + session.flush(self.tag_files) + + if deleted_tags := (existing_tag_ids - tag_ids): + session.query(TagFile).filter( + TagFile.file_group_id == self.id, + TagFile.tag_id.in_(deleted_tags), + ).delete(synchronize_session=False) + + # Save changes to config. + save_tags_config.activate_switch() + sync_tags_directory.activate_switch() + + return self.tag_files + + @singledispatchmethod @optional_session - def untag(self, tag_id_or_name: int | str, session: Session = None): - Tag.untag_file_group(self.id, tag_id_or_name, session) + def add_tag(self, tag_id: int | str, session: Session = None) -> TagFile: + tags = [i.tag_id for i in self.tag_files] + tags.append(tag_id) + tag_files = self.set_tags(tags, session=session) + return next(i for i in tag_files if i.tag_id == tag_id) + + @add_tag.register + def _(self, tag_name: str, session: Session = None) -> TagFile: + tag_id = Tag.get_id_by_name(tag_name, session=session) + return self.add_tag(tag_id, session=session) + + @singledispatchmethod + @optional_session + def untag(self, tag_id: int | str, session: Session = None): + tags = [i.tag_id for i in self.tag_files if i.tag_id != tag_id] + self.set_tags(tags, session=session) + + @untag.register + def _(self, tag_name: str, session: Session = None) -> TagFile: + tag_id = Tag.get_id_by_name(tag_name, session=session) + return self.untag(tag_id, session=session) def append_files(self, *paths: pathlib.Path): """Add all `paths` to this FileGroup.files.""" diff --git a/wrolpi/files/test/test_lib.py b/wrolpi/files/test/test_lib.py index c94cc5d8..33617a45 100644 --- a/wrolpi/files/test/test_lib.py +++ b/wrolpi/files/test/test_lib.py @@ -16,6 +16,7 @@ from PIL import Image import wrolpi.common from modules.videos import Video from wrolpi.common import timer +from wrolpi.conftest import await_switches from wrolpi.dates import now from wrolpi.errors import InvalidFile, UnknownDirectory, FileGroupIsTagged, NoPrimaryFile from wrolpi.files import lib, indexers diff --git a/wrolpi/files/test/test_models.py b/wrolpi/files/test/test_models.py index e2685e3c..c2fd8fcd 100644 --- a/wrolpi/files/test/test_models.py +++ b/wrolpi/files/test/test_models.py @@ -1,5 +1,8 @@ import pytest +from wrolpi.files import lib +from wrolpi.files.models import FileGroup + @pytest.mark.asyncio async def test_file_group_move(async_client, test_session, test_directory, video_factory): @@ -19,3 +22,23 @@ async def test_file_group_move(async_client, test_session, test_directory, video (test_directory / 'videos/NO CHANNEL/new name.jpg'), (test_directory / 'videos/NO CHANNEL/new name.mp4'), ] + + +@pytest.mark.asyncio +async def test_add_tag_overload(async_client, test_session, make_files_structure, tag_factory): + """FileGroup.add_tag and FileGroup.untag use `singeldispatch`.""" + make_files_structure({'foo.txt': 'foo contents'}) + await lib.refresh_files() + foo: FileGroup = test_session.query(FileGroup).one() + + one, two = await tag_factory(), await tag_factory() + + foo.add_tag(one.name) + assert set(foo.tag_names) == {'one'}, 'Tag should have been found by name' + foo.add_tag(two.id) + assert set(foo.tag_names) == {'one', 'two'}, 'Tag should have been found by ID.' + + foo.untag(one.name) + assert set(foo.tag_names) == {'two'}, 'Tag should have been found by name' + foo.untag(two.id) + assert not foo.tag_names, 'Tag should have been found by ID.' diff --git a/wrolpi/tags.py b/wrolpi/tags.py index 0b2e2a45..7e73abc4 100644 --- a/wrolpi/tags.py +++ b/wrolpi/tags.py @@ -4,7 +4,9 @@ from dataclasses import dataclass, field from datetime import datetime from typing import List, Dict, Tuple, Optional -from sqlalchemy import Column, Integer, String, ForeignKey, BigInteger +import cachetools +from cachetools.keys import hashkey +from sqlalchemy import Column, Integer, String, ForeignKey, BigInteger, event from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import relationship, Session @@ -14,7 +16,7 @@ from wrolpi.common import ModelHelper, Base, logger, ConfigFile, get_media_direc from wrolpi.dates import TZDateTime from wrolpi.db import optional_session, get_db_curs, get_db_session from wrolpi.downloader import save_downloads_config -from wrolpi.errors import UnknownTag, UsedTag, InvalidTag, FileGroupAlreadyTagged, RefreshConflict, NoPrimaryFile +from wrolpi.errors import UnknownTag, UsedTag, InvalidTag, RefreshConflict, NoPrimaryFile from wrolpi.events import Events from wrolpi.switches import register_switch_handler, ActivateSwitchMethod from wrolpi.vars import PYTEST @@ -51,6 +53,22 @@ class TagFile(ModelHelper, Base): .one_or_none() +# Tag.get_id_by_name +def get_id_by_name_key(klass: Base, name: str, session: Session = None): + return hashkey(name) + + +get_id_by_name_cache = cachetools.LRUCache(maxsize=1_000) + + +# Tag.get_name_by_id +def get_name_by_id_key(klass: Base, id_: int, session: Session = None): + return hashkey(id_) + + +get_name_by_id_cache = cachetools.LRUCache(maxsize=1_000) + + class Tag(ModelHelper, Base): __tablename__ = 'tag' id = Column(Integer, primary_key=True) @@ -74,55 +92,32 @@ class Tag(ModelHelper, Base): color=self.color, ) - @staticmethod + @classmethod + @cachetools.cached(cache=get_id_by_name_cache, key=get_id_by_name_key) @optional_session - def get_tag_file_group(file_group_id: int, tag_id_or_name: int | str, session: Session = None) -> Optional[TagFile]: - if isinstance(tag_id_or_name, int): - return TagFile.get_by_primary_keys(file_group_id, tag_id_or_name, session) + def get_id_by_name(cls, name: str, session: Session = None) -> int: + """Returns a Tag's ID, if a Tag matches the provided name. - return TagFile.get_by_tag_name(file_group_id, tag_id_or_name, session) + @raise UnknownTag: If no Tag matches the provided name.""" + if tag := session.query(Tag).filter_by(name=name).one_or_none(): + return tag.id + raise UnknownTag(f'No Tag with name={name}') - @staticmethod + @classmethod + @cachetools.cached(cache=get_name_by_id_cache, key=get_name_by_id_key) @optional_session - def tag_file_group(file_group_id: int, tag_id_or_name: int | str, session: Session = None) -> TagFile: - """Add a TagFile for the provided FileGroup and this Tag. + def get_name_by_id(cls, id_: int, session: Session = None) -> str: + """Returns a Tag's name, if a Tag matches the provided ID. - @warning: Commits the session to keep the config in sync.""" - existing = Tag.get_tag_file_group(file_group_id, tag_id_or_name, session) + @raise UnknownTag: If no Tag matches the provided ID.""" + if tag := session.query(Tag).filter_by(id=id_).one_or_none(): + return tag.id + raise UnknownTag(f'No Tag with id={id_}') - if existing: - raise FileGroupAlreadyTagged('Tag already used') - - tag = Tag.get_by_name(tag_id_or_name, session) if isinstance(tag_id_or_name, str) \ - else Tag.get_by_id(tag_id_or_name, session) - tag_file = TagFile(file_group_id=file_group_id, tag_id=tag.id) - session.add(tag_file) - session.flush([tag_file]) - session.commit() - logger.info(f'Tagged FileGroup {file_group_id} with Tag {tag_id_or_name}') - - # Save changes to config. - save_tags_config.activate_switch() - sync_tags_directory.activate_switch() - return tag_file - - @staticmethod - @optional_session - def untag_file_group(file_group_id: int, tag_id_or_name: int | str, session: Session = None): - """Remove the record of a Tag applied to the FileGroup. - - @warning: Commits the session to keep config in sync.""" - tag_file = Tag.get_tag_file_group(file_group_id, tag_id_or_name, session) - if tag_file: - session.delete(tag_file) - session.commit() - logger.debug(f'Deleted TagFile: {tag_file}') - - # Save changes to config. - save_tags_config.activate_switch() - sync_tags_directory.activate_switch() - else: - logger.warning(f'Could not find TagFile for FileGroup.id={file_group_id}/Tag.id={tag_id_or_name}') + @classmethod + def invalidate_cache(cls): + get_id_by_name_cache.clear() + get_name_by_id_cache.clear() @staticmethod @optional_session @@ -136,7 +131,7 @@ class Tag(ModelHelper, Base): if tag := Tag.get_by_name(name, session): return tag - raise UnknownTag(f'Unknown Tag with name={name}') + raise UnknownTag(f'No Tag with name={name}') @staticmethod @optional_session @@ -150,7 +145,7 @@ class Tag(ModelHelper, Base): if tag := Tag.get_by_id(id_, session): return tag - raise UnknownTag(f'Unknown Tag with id={id_}') + raise UnknownTag(f'No Tag with id={id_}') def has_relations(self) -> bool: """Returns True if this Tag has been used with any FileGroups or Zim Entries.""" @@ -254,6 +249,14 @@ class Tag(ModelHelper, Base): sync_tags_directory.activate_switch() +@event.listens_for(Tag, 'after_insert') +@event.listens_for(Tag, 'after_update') +@event.listens_for(Tag, 'after_delete') +def invalidate_cache(mapper, connection, target): + """Clear Tag cache when Tags are changed.""" + Tag.invalidate_cache() + + @dataclass class TagsConfigValidator: tag_files: list[list] = field(default_factory=list) diff --git a/wrolpi/test/test_tags.py b/wrolpi/test/test_tags.py index bb6c0802..24badef1 100644 --- a/wrolpi/test/test_tags.py +++ b/wrolpi/test/test_tags.py @@ -6,7 +6,7 @@ import yaml from wrolpi import tags from wrolpi.common import is_hardlinked, walk -from wrolpi.errors import FileGroupIsTagged, InvalidTag, UnknownTag, FileGroupAlreadyTagged, UsedTag +from wrolpi.errors import FileGroupIsTagged, InvalidTag, UnknownTag, UsedTag from wrolpi.files import lib as files_lib from wrolpi.files.models import FileGroup from wrolpi.tags import TagFile, Tag @@ -20,19 +20,18 @@ async def test_tags_file_group_json(async_client, test_session, make_files_struc await files_lib.refresh_files() file_group: FileGroup = test_session.query(FileGroup).one() - Tag.tag_file_group(file_group.id, tag_one.id) + file_group.add_tag(tag_one.id) assert file_group.__json__()['tags'] == ['one'] - file_group.add_tag(tag_two.id) + file_group.add_tag(tag_two.name) assert file_group.__json__()['tags'] == ['one', 'two'] - with pytest.raises(FileGroupAlreadyTagged): - # Cannot tag a FileGroup with the same Tag twice. - file_group.add_tag(tag_two.id) - file_group.untag(tag_one.id) assert file_group.__json__()['tags'] == ['two'] + file_group.untag(tag_two.name) + assert file_group.__json__()['tags'] == [] + @pytest.mark.asyncio async def test_tags_model(async_client, test_session, make_files_structure, tag_factory, example_pdf): @@ -66,13 +65,13 @@ async def test_tags_file_group(async_client, test_session, make_files_structure, tag_one = await tag_factory() assert len(tag_one.tag_files) == 0 - Tag.tag_file_group(video_group.id, tag_one.id) + video_group.add_tag(tag_one.id) test_session.commit() assert len(tag_one.tag_files) == 1 assert sorted([i.tag.name for i in video_group.tag_files]) == ['one', ] tag_two = await tag_factory() - Tag.tag_file_group(video_group.id, tag_two.id) + video_group.add_tag(tag_two.id) test_session.commit() assert len(tag_one.tag_files) == 1 assert len(tag_two.tag_files) == 1