rt/wrolpi
1
0
Fork 0
mirror of https://github.com/lrnselfreliance/wrolpi.git synced 2026-10-26 02:48:38 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

Removing optional_session, passing session down from request middleware. Moving session param to the front of functions.
Some checks failed
CodeQL / Analyze (push) Has been cancelled
Details
Jest Tests / jest-tests (push) Has been cancelled
Details

Browse source
This commit is contained in:
Roland 2025-12-10 21:16:10 -07:00
parent 1747442db3
commit 344fe9035f
No known key found for this signature in database
GPG key ID: 1BA620CE6C9A4680
70 changed files with 1035 additions and 1110 deletions
Download patch file Download diff file Expand all files Collapse all files

13
modules/archive/__init__.py
View file

@ -11,7 +11,7 @@ from wrolpi.cmd import SINGLE_FILE_BIN, CHROMIUM
from wrolpi.collections import Collection
from wrolpi.common import logger, register_modeler, register_refresh_cleanup, limit_concurrent, split_lines_by_length, \
slow_logger, get_title_from_html
from wrolpi.db import optional_session, get_db_session
from wrolpi.db import get_db_session
from wrolpi.downloader import Downloader, Download, DownloadResult
from wrolpi.errors import UnrecoverableDownloadError
from wrolpi.files.models import FileGroup
@ -36,8 +36,7 @@ class ArchiveDownloader(Downloader, ABC):
def __repr__(self):
return f'<ArchiveDownloader>'
@optional_session
def already_downloaded(self, *urls: List[str], session: Session = None) -> List:
def already_downloaded(self, session: Session, *urls: List[str]) -> List:
file_groups = list(session.query(FileGroup).filter(FileGroup.url.in_(urls), FileGroup.model == 'archive'))
return file_groups
@ -57,7 +56,7 @@ class ArchiveDownloader(Downloader, ABC):
archive_id = archive.id
with get_db_session() as session:
archive = Archive.find_by_id(archive_id, session)
archive = Archive.find_by_id(session, archive_id)
need_commit = False
if tag_names := download.tag_names:
for name in tag_names:
@ -107,7 +106,7 @@ class ArchiveDownloader(Downloader, ABC):
archive_downloader = ArchiveDownloader()
def model_archive(file_group: FileGroup, session: Session = None) -> Archive:
def model_archive(session: Session, file_group: FileGroup) -> Archive:
"""
Models an Archive from a given FileGroup.
@ -117,8 +116,8 @@ def model_archive(file_group: FileGroup, session: Session = None) -> Archive:
(title, contents) from either the JSON or text readability files.
Args:
session: Database session.
file_group: The FileGroup to model as an Archive.
session: An optional database session for committing changes.
Returns:
An Archive object representing the modeled archive.
@ -289,7 +288,7 @@ async def archive_modeler():
raise
else:
try:
model_archive(file_group, session=session)
model_archive(session, file_group)
# Successfully modeled, mark as indexed
file_group.indexed = True
except InvalidArchive:

31
modules/archive/api.py
View file

@ -5,10 +5,10 @@ from http import HTTPStatus
from sanic import response, Request, Blueprint
from sanic_ext import validate
from sanic_ext.extensions.openapi import openapi
from sqlalchemy.orm import Session
from wrolpi.api_utils import json_response, api_app
from wrolpi.common import logger, wrol_mode_check, api_param_limiter, TRACE_LEVEL
from wrolpi.db import get_db_session
from wrolpi.errors import ValidationError
from wrolpi.events import Events
from wrolpi.schema import JSONErrorResponse
@ -25,8 +25,9 @@ logger = logger.getChild(__name__)
@archive_bp.get('/<archive_id:int>')
@openapi.description('Get an archive')
@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
async def get_archive(_: Request, archive_id: int):
archive = lib.get_archive(archive_id=archive_id)
async def get_archive(request: Request, archive_id: int):
session = request.ctx.session
archive = lib.get_archive(session, archive_id)
archive_file_group = archive.file_group.__json__()
history = [i.file_group.__json__() for i in archive.history]
return json_response({'file_group': archive_file_group, 'history': history})
@ -115,12 +116,12 @@ async def singlefile_upload_switch_handler(url=None):
logger.info(f'Created Archive from upload ({q_size}): {archive}')
Events.send_archive_uploaded(f'Created Archive from upload: {name}', url=archive.location)
url = archive.file_group.url
if url and (download := Download.get_by_url(url)):
if (download.is_failed or download.is_deferred) and download.downloader == ArchiveDownloader.name:
# Download was attempted and failed, user manually archived the URL.
download.complete()
download.location = archive.location
Session.object_session(download).commit()
with get_db_session(commit=True) as session:
if url and (download := Download.get_by_url(session, url)):
if (download.is_failed or download.is_deferred) and download.downloader == ArchiveDownloader.name:
# Download was attempted and failed, user manually archived the URL.
download.complete()
download.location = archive.location
# Call this function again so any new singlefiles can be archived.
singlefile_upload_switch_handler.activate_switch()
@ -154,9 +155,10 @@ async def generate_screenshot_switch_handler(archive_id=None):
try:
await lib.generate_archive_screenshot(archive_id)
# Always send success event since exceptions are raised on failure
archive = lib.get_archive(archive_id=archive_id)
location = archive.location
name = archive.file_group.title or archive.file_group.url
with get_db_session() as session:
archive = lib.get_archive(session, archive_id)
location = archive.location
name = archive.file_group.title or archive.file_group.url
logger.info(f'Generated screenshot for Archive ({q_size}): {archive_id}')
Events.send_screenshot_generated(f'Generated screenshot for: {name}', url=location)
except Exception as e:
@ -196,11 +198,12 @@ async def post_upload_singlefile(request: Request):
@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
@openapi.response(HTTPStatus.BAD_REQUEST, JSONErrorResponse)
@wrol_mode_check
async def post_generate_screenshot(_: Request, archive_id: int):
async def post_generate_screenshot(request: Request, archive_id: int):
"""Queue a screenshot generation request for an Archive."""
session = request.ctx.session
# Verify archive exists
try:
archive = lib.get_archive(archive_id=archive_id)
archive = lib.get_archive(session, archive_id)
except Exception:
return json_response({'error': f'Archive {archive_id} not found'}, status=HTTPStatus.NOT_FOUND)

2
modules/archive/conftest.py
View file

@ -93,7 +93,7 @@ def archive_factory(test_session, archive_directory, make_files_structure, image
archive.validate()
for tag_name in tag_names:
archive.add_tag(tag_name, test_session)
archive.add_tag(test_session, tag_name)
return archive

45
modules/archive/lib.py
View file

@ -24,7 +24,7 @@ from wrolpi.common import get_media_directory, get_relative_to_media_directory,
escape_file_name, aiohttp_post, format_html_string, split_lines_by_length, get_html_soup, get_title_from_html, \
get_wrolpi_config, html_screenshot, ConfigFile
from wrolpi.dates import now, Seconds
from wrolpi.db import get_db_session, get_db_curs, optional_session
from wrolpi.db import get_db_session, get_db_curs
from wrolpi.errors import UnknownArchive, InvalidOrderBy, InvalidDatetime
from wrolpi.events import Events
from wrolpi.files.lib import handle_file_group_search_results
@ -143,7 +143,7 @@ class DomainsConfig(ConfigFile):
# Batch create/update collections
imported_domains = set()
if valid_data_list:
collections = Collection.batch_from_config(valid_data_list, session=session)
collections = Collection.batch_from_config(session, valid_data_list)
imported_domains = {c.name for c in collections}
# Delete domain collections that are no longer in config
@ -453,7 +453,7 @@ async def model_archive_result(url: str, singlefile: str, readability: dict, scr
return archive
def detect_domain_directory(collection: Collection, session: Session) -> Optional[pathlib.Path]:
def detect_domain_directory(session: Session, collection: Collection) -> Optional[pathlib.Path]:
"""
Detect if all archives for a domain collection share a common directory.
@ -546,7 +546,7 @@ def update_domain_directories(session: Session = None) -> int:
updated_count = 0
for collection in collections:
detected_dir = detect_domain_directory(collection, session)
detected_dir = detect_domain_directory(session, collection)
if detected_dir:
collection.directory = detected_dir
session.flush([collection])
@ -598,7 +598,7 @@ def get_or_create_domain_collection(session: Session, url, directory: pathlib.Pa
# Auto-detect directory if not explicitly set and collection doesn't have one
if not directory and not collection.directory:
detected_dir = detect_domain_directory(collection, session)
detected_dir = detect_domain_directory(session, collection)
if detected_dir:
collection.directory = detected_dir
session.flush()
@ -810,10 +810,9 @@ def parse_article_html_metadata(html: Union[bytes, str], assume_utc: bool = True
return metadata
@optional_session
def get_archive(session, archive_id: int) -> Archive:
def get_archive(session: Session, archive_id: int) -> Archive:
"""Get an Archive."""
archive = Archive.find_by_id(archive_id, session=session)
archive = Archive.find_by_id(session, archive_id)
archive.file_group.set_viewed()
return archive
@ -907,17 +906,18 @@ def get_domain(domain_id: int) -> dict:
Raises UnknownArchive if domain not found.
"""
# Get the collection using find_by_id
try:
collection = Collection.find_by_id(domain_id)
except Exception:
# Collection.find_by_id raises UnknownCollection, but we want UnknownArchive
raise UnknownArchive(f"Domain collection with ID {domain_id} not found")
with get_db_session() as session:
try:
collection = Collection.find_by_id(session, domain_id)
except Exception:
# Collection.find_by_id raises UnknownCollection, but we want UnknownArchive
raise UnknownArchive(f"Domain collection with ID {domain_id} not found")
if collection.kind != 'domain':
raise UnknownArchive(f"Collection {domain_id} is not a domain")
if collection.kind != 'domain':
raise UnknownArchive(f"Collection {domain_id} is not a domain")
# Get base domain data from __json__()
domain_data = collection.__json__()
# Get base domain data from __json__()
domain_data = collection.__json__()
# Get archive statistics for this domain
with get_db_curs() as curs:
@ -1045,15 +1045,14 @@ def search_archives(search_str: str, domain: str, limit: int, offset: int, order
return results, total
@optional_session
async def search_domains_by_name(name: str, limit: int = 5, session: Session = None) -> List[dict]:
async def search_domains_by_name(session: Session, name: str, limit: int = 5) -> List[dict]:
"""
Search for domain collections by name.
Args:
session: Database session
name: Search string to match against collection names
limit: Maximum number of results to return
session: Database session
Returns:
List of domain dicts matching the search (in old Domain format for backward compatibility)
@ -1163,7 +1162,7 @@ async def generate_archive_screenshot(archive_id: int) -> pathlib.Path:
RuntimeError: If screenshot generation fails
"""
with get_db_session() as session:
archive = Archive.find_by_id(archive_id, session=session)
archive = Archive.find_by_id(session, archive_id)
if not archive.singlefile_path:
raise ValueError(f'Cannot generate screenshot for Archive {archive_id}: no singlefile')
@ -1175,7 +1174,7 @@ async def generate_archive_screenshot(archive_id: int) -> pathlib.Path:
logger.info(f'Archive {archive_id} already has a screenshot, ensuring it is tracked')
# Ensure the screenshot is tracked in FileGroup.files and FileGroup.data
with get_db_session(commit=True) as tracking_session:
archive = Archive.find_by_id(archive_id, session=tracking_session)
archive = Archive.find_by_id(tracking_session, archive_id)
file_group = archive.file_group
# append_files uses unique_by_predicate, so this is safe even if already tracked
file_group.append_files(archive.screenshot_path)
@ -1213,7 +1212,7 @@ async def generate_archive_screenshot(archive_id: int) -> pathlib.Path:
# Update the Archive to include the new screenshot file
with get_db_session(commit=True) as session:
archive = Archive.find_by_id(archive_id, session=session)
archive = Archive.find_by_id(session, archive_id)
archive.set_screenshot(screenshot_path)
session.flush()

23
modules/archive/models.py
View file

@ -11,7 +11,6 @@ from wrolpi import dates
from wrolpi.collections import Collection
from wrolpi.common import ModelHelper, Base, logger, get_title_from_html, get_wrolpi_config, get_media_directory
from wrolpi.dates import now
from wrolpi.db import optional_session
from wrolpi.errors import UnknownArchive
from wrolpi.files.models import FileGroup
from wrolpi.tags import TagFile
@ -62,33 +61,29 @@ class Archive(Base, ModelHelper):
return None
@staticmethod
@optional_session
def get_by_id(id_: int, session: Session = None) -> Optional['Archive']:
def get_by_id(session: Session, id_: int) -> Optional['Archive']:
"""Attempt to find an Archive with the provided id. Returns None if it cannot be found."""
archive = session.query(Archive).filter_by(id=id_).one_or_none()
return archive
@staticmethod
@optional_session
def find_by_id(id_: int, session: Session = None) -> 'Archive':
def find_by_id(session: Session, id_: int) -> 'Archive':
"""Find an Archive with the provided id, raises an exception when no Archive is found.
@raise UnknownArchive: if the Archive can not be found"""
archive = Archive.get_by_id(id_, session=session)
archive = Archive.get_by_id(session, id_)
if not archive:
raise UnknownArchive(f'Cannot find Archive with id {id_}')
return archive
@staticmethod
@optional_session
def get_by_path(path: pathlib.Path | str, session: Session = None) -> Optional['Archive']:
def get_by_path(session: Session, path: pathlib.Path | str) -> Optional['Archive']:
archive = session.query(Archive).join(FileGroup).filter(FileGroup.primary_path == str(path)).one_or_none()
return archive
@staticmethod
@optional_session
def find_by_path(path: pathlib.Path | str, session: Session = None) -> Optional['Archive']:
archive = Archive.get_by_path(path, session)
def find_by_path(session: Session, path: pathlib.Path | str) -> Optional['Archive']:
archive = Archive.get_by_path(session, path)
if archive:
return archive
raise UnknownArchive(f'Cannot find Archive with path: {path}')
@ -102,10 +97,10 @@ class Archive(Base, ModelHelper):
return False
@staticmethod
def do_model(file_group: FileGroup, session: Session) -> 'Archive':
def do_model(session: Session, file_group: FileGroup) -> 'Archive':
# Local import to avoid circular import within archive module
from modules.archive import model_archive
archive = model_archive(file_group, session)
archive = model_archive(session, file_group)
archive.validate()
file_group.indexed = True
return archive
@ -377,7 +372,7 @@ class Archive(Base, ModelHelper):
# Local import to avoid circular import within archive module
from modules.archive import model_archive
file_group = FileGroup.from_paths(session, *paths)
archive = model_archive(file_group, session)
archive = model_archive(session, file_group)
return archive
def add_tag(self, tag_id_or_name: int | str, session: Session = None) -> TagFile:

4
modules/archive/test/test_collection_id_bug.py
View file

@ -65,7 +65,7 @@ async def test_model_archive_sets_collection_id(async_client, test_session, test
# Model the archive using the production code path
# This should create an Archive with collection_id set
archive = model_archive(file_group, session=test_session)
archive = model_archive(test_session, file_group)
# The bug: collection_id is None, causing NOT NULL constraint violation
assert archive is not None, "Archive should be created"
@ -126,7 +126,7 @@ async def test_model_archive_extracts_url_from_singlefile(async_client, test_ses
file_group = FileGroup.from_paths(test_session, *file_paths)
# Model the archive - it should extract the URL from the singlefile
archive = model_archive(file_group, session=test_session)
archive = model_archive(test_session, file_group)
assert archive is not None
assert archive.file_group.url == url, \

4
modules/archive/test/test_domain_tagging.py
View file

@ -310,7 +310,7 @@ async def test_tag_domain_without_directory_saves_config(test_session, test_dire
test_session.commit()
# Tag using library function (no directory)
await tag_collection(collection.id, tag_name='News', directory=None, session=test_session)
await tag_collection(test_session, collection.id, tag_name='News', directory=None)
test_session.commit()
# Wait for switches to process
@ -340,7 +340,7 @@ async def test_untag_domain_without_directory_saves_config(test_session, test_di
test_session.commit()
# Untag using library function (tag_name=None)
await tag_collection(collection.id, tag_name=None, directory=None, session=test_session)
await tag_collection(test_session, collection.id, tag_name=None, directory=None)
test_session.commit()
# Wait for switches to process

6
modules/archive/test/test_lib.py
View file

@ -883,7 +883,7 @@ async def test_detect_domain_directory_single_archive(async_client, test_directo
assert 'example.com' in str(collection.directory)
# Verify detect_domain_directory also returns the same result
detected = detect_domain_directory(collection, test_session)
detected = detect_domain_directory(test_session, collection)
assert detected is not None
assert str(detected) == 'archive/example.com'
@ -904,7 +904,7 @@ async def test_detect_domain_directory_multiple_archives(async_client, test_dire
assert collection.name == 'test.com'
# Detect directory
detected = detect_domain_directory(collection, test_session)
detected = detect_domain_directory(test_session, collection)
assert detected is not None
assert str(detected) == 'archive/test.com'
@ -919,7 +919,7 @@ def test_detect_domain_directory_no_archives(test_directory, test_session):
test_session.flush()
# Detect directory - should return None
detected = detect_domain_directory(collection, test_session)
detected = detect_domain_directory(test_session, collection)
assert detected is None

31
modules/archive/test/test_models.py
View file

@ -5,6 +5,37 @@ import pytest
from modules.archive import Archive
from wrolpi.common import get_wrolpi_config
from wrolpi.errors import UnknownArchive
def test_archive_get_by_id(test_session, archive_factory):
"""Archive.get_by_id should accept session as first argument, id as second."""
archive = archive_factory()
test_session.commit()
# Session must be first argument (session-first pattern)
found = Archive.get_by_id(test_session, archive.id)
assert found is not None
assert found.id == archive.id
# Non-existent ID should return None
not_found = Archive.get_by_id(test_session, 99999)
assert not_found is None
def test_archive_find_by_id(test_session, archive_factory):
"""Archive.find_by_id should accept session as first argument, id as second."""
archive = archive_factory()
test_session.commit()
# Session must be first argument (session-first pattern)
found = Archive.find_by_id(test_session, archive.id)
assert found is not None
assert found.id == archive.id
# Non-existent ID should raise UnknownArchive
with pytest.raises(UnknownArchive):
Archive.find_by_id(test_session, 99999)
@pytest.mark.asyncio

4
modules/inventory/api.py
View file

@ -28,8 +28,8 @@ def get_brands(_: Request):
@inventory_bp.get('/')
def get_inventories(_: Request):
inventories = inventory.get_inventories()
def get_inventories(request: Request):
inventories = inventory.get_inventories(request.ctx.session)
return json_response(dict(inventories=inventories))

5
modules/inventory/common.py
View file

@ -142,8 +142,9 @@ class InventoriesConfig(ConfigFile):
def dump_config(self, file: pathlib.Path = None, send_events=False, overwrite=False):
inventories = []
for inventory in get_inventories():
inventories.append(inventory.dict())
with get_db_session() as session:
for inventory in get_inventories(session):
inventories.append(inventory.dict())
if not inventories:
raise NoInventories('No Inventories are in the database!')

5
modules/inventory/inventory.py
View file

@ -7,7 +7,7 @@ import psycopg2
from sqlalchemy.orm import Session
from wrolpi.common import logger, Base
from wrolpi.db import get_db_session, get_db_curs, optional_session
from wrolpi.db import get_db_session, get_db_curs
from wrolpi.errors import APIError
from .models import Inventory, Item
@ -85,8 +85,7 @@ DEFAULT_INVENTORIES = [
sort_categories = partial(sorted, key=itemgetter(1, 0))
@optional_session
def get_inventories(session: Session = None) -> List[Inventory]:
def get_inventories(session: Session) -> List[Inventory]:
inventories = session.query(Inventory).filter(
Inventory.deleted_at == None, # noqa
).order_by(

4
modules/inventory/test/test_inventory.py
View file

@ -81,7 +81,7 @@ class TestInventory(PytestCase):
def test_get_inventories(self, test_session):
self.prepare()
inventories = get_inventories()
inventories = get_inventories(test_session)
dates = extract_items(inventories, {'created_at', 'deleted_at', 'viewed_at'})
self.assertItemsTruthyOrFalsey(dates, [{'created_at': True, 'deleted_at': False, 'viewed_at': False}])
self.assertEqual(inventories[0].id, 1)
@ -290,7 +290,7 @@ async def test_inventories_config(await_switches, test_session, test_directory,
import_inventories_config()
inventories = get_inventories()
inventories = get_inventories(test_session)
assert len(inventories) == 1
# ID has increased because we did not reset the sequence when deleting from the table.
assert inventories[0].id == 2

2
modules/map/api.py
View file

@ -41,7 +41,7 @@ async def import_pbfs(_: Request, body: schema.ImportPost):
@map_bp.get('/files')
@openapi.description('Find any map files, get their import status')
def get_files_status(request: Request):
paths = lib.get_import_status()
paths = lib.get_import_status(request.ctx.session)
paths = sorted(paths, key=lambda i: str(i.path))
pending = request.app.shared_ctx.map_importing.get('pending')
if pending:

5
modules/map/lib.py
View file

@ -11,7 +11,7 @@ from wrolpi.api_utils import api_app
from wrolpi.cmd import SUDO_BIN, run_command
from wrolpi.common import get_media_directory, walk, logger, get_wrolpi_config
from wrolpi.dates import timedelta_to_timestamp, seconds_to_timestamp, now
from wrolpi.db import optional_session, get_db_session
from wrolpi.db import get_db_session
from wrolpi.events import Events
from wrolpi.vars import PROJECT_DIR, IS_RPI5
@ -208,8 +208,7 @@ async def run_import_command(*paths: Path):
raise ValueError(f'Importing map file failed with return code {result.return_code}')
@optional_session
def get_import_status(session: Session = None) -> List[MapFile]:
def get_import_status(session: Session) -> List[MapFile]:
paths = get_map_paths()
map_paths = []
for path in paths:

7
modules/videos/api.py
View file

@ -37,12 +37,13 @@ async def statistics(_: Request):
body=schema.ChannelTagInfoRequest,
)
@validate(schema.ChannelTagInfoRequest)
async def channel_tag_info(_: Request, body: schema.ChannelTagInfoRequest):
async def channel_tag_info(request: Request, body: schema.ChannelTagInfoRequest):
from wrolpi.tags import Tag
channel = Channel.find_by_id(body.channel_id) if body.channel_id else None
session = request.ctx.session
channel = Channel.find_by_id(session, body.channel_id) if body.channel_id else None
channel_name = channel.name if channel else None
channel_url = channel.url if channel else None
tag_name = Tag.find_by_name(body.tag_name).name if body.tag_name else None
tag_name = Tag.find_by_name(session, body.tag_name).name if body.tag_name else None
videos_destination = format_videos_destination(channel_name, tag_name, channel_url)
ret = dict(videos_destination=videos_destination)
return json_response(ret)

35
modules/videos/channel/api.py
View file

@ -34,8 +34,9 @@ async def get_channels(_: Request):
@channel_bp.get('/<channel_id:int>')
@openapi.description('Get a Channel')
@openapi.response(HTTPStatus.OK, schema.ChannelResponse)
def channel_get(_: Request, channel_id: int = None):
channel = lib.get_channel(channel_id=channel_id, return_dict=True)
def channel_get(request: Request, channel_id: int = None):
session = request.ctx.session
channel = lib.get_channel(session, channel_id=channel_id, return_dict=True)
channel.pop('info_json')
return json_response({'channel': channel})
@ -49,10 +50,11 @@ def channel_get(_: Request, channel_id: int = None):
@openapi.response(HTTPStatus.BAD_REQUEST, JSONErrorResponse)
@validate(schema.ChannelPostRequest)
@wrol_mode_check
def channel_post(_: Request, body: schema.ChannelPostRequest):
def channel_post(request: Request, body: schema.ChannelPostRequest):
session = request.ctx.session
body.directory = get_media_directory() / body.directory
channel = lib.create_channel(data=body, return_dict=False)
channel = lib.create_channel(session, data=body, return_dict=False)
# Refresh the videos asynchronously
if not PYTEST:
@ -73,9 +75,10 @@ def channel_post(_: Request, body: schema.ChannelPostRequest):
@openapi.response(HTTPStatus.BAD_REQUEST, JSONErrorResponse)
@validate(schema.ChannelPutRequest)
@wrol_mode_check
async def channel_update(_: Request, channel_id: int, body: schema.ChannelPutRequest):
async def channel_update(request: Request, channel_id: int, body: schema.ChannelPutRequest):
session = request.ctx.session
body.directory = pathlib.Path(body.directory) if body.directory else None
channel = await lib.update_channel(data=body, channel_id=channel_id)
channel = await lib.update_channel(session, data=body, channel_id=channel_id)
save_channels_config.activate_switch()
return response.raw('', HTTPStatus.NO_CONTENT,
headers={'Location': f'/api/videos/channels/{channel.id}'})
@ -85,8 +88,9 @@ async def channel_update(_: Request, channel_id: int, body: schema.ChannelPutReq
@openapi.description('Delete a Channel')
@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
@wrol_mode_check
def channel_delete(_: Request, channel_id: int):
channel = lib.delete_channel(channel_id=channel_id)
def channel_delete(request: Request, channel_id: int):
session = request.ctx.session
channel = lib.delete_channel(session, channel_id=channel_id)
Events.send_deleted(f'Deleted Channel: {channel["name"]}')
return response.raw('', HTTPStatus.NO_CONTENT)
@ -95,8 +99,9 @@ def channel_delete(_: Request, channel_id: int):
@openapi.description('Refresh all files in the Channel directory')
@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
@wrol_mode_check
def channel_refresh(_: Request, channel_id: int):
channel: Channel = lib.get_channel(channel_id=channel_id, return_dict=False)
def channel_refresh(request: Request, channel_id: int):
session = request.ctx.session
channel: Channel = lib.get_channel(session, channel_id=channel_id, return_dict=False)
if not channel:
raise UnknownChannel()
@ -115,13 +120,14 @@ def channel_refresh(_: Request, channel_id: int):
@openapi.response(HTTPStatus.BAD_REQUEST, JSONErrorResponse)
@validate(schema.ChannelTagRequest)
@wrol_mode_check
async def channel_tag(_: Request, channel_id: int, body: schema.ChannelTagRequest):
async def channel_tag(request: Request, channel_id: int, body: schema.ChannelTagRequest):
session = request.ctx.session
directory = None
if body.directory:
directory = pathlib.Path(body.directory)
directory = get_media_directory() / directory
await lib.tag_channel(body.tag_name, directory, channel_id)
await lib.tag_channel(session, body.tag_name, directory, channel_id)
return response.raw('', HTTPStatus.NO_CONTENT)
@ -132,8 +138,9 @@ async def channel_tag(_: Request, channel_id: int, body: schema.ChannelTagReques
body=schema.ChannelSearchRequest,
)
@validate(schema.ChannelSearchRequest)
async def channel_search(_: Request, body: schema.ChannelSearchRequest):
channels = await lib.search_channels(body.tag_names)
async def channel_search(request: Request, body: schema.ChannelSearchRequest):
session = request.ctx.session
channels = await lib.search_channels(session, body.tag_names)
ret = dict(
channels=channels,
)

29
modules/videos/channel/lib.py
View file

@ -9,7 +9,7 @@ from wrolpi import flags
from wrolpi.collections import Collection
from wrolpi.common import logger, \
get_media_directory, wrol_mode_check, background_task
from wrolpi.db import get_db_curs, optional_session, get_db_session
from wrolpi.db import get_db_curs, get_db_session
from wrolpi.downloader import save_downloads_config, download_manager, Download
from wrolpi.errors import APIError, ValidationError, RefreshConflict
from wrolpi.tags import Tag
@ -53,7 +53,6 @@ async def get_minimal_channels() -> List[dict]:
COD = Union[dict, Channel]
@optional_session
def get_channel(session: Session, *, channel_id: int = None, source_id: str = None, url: str = None,
directory: str = None, name: str = None, return_dict: bool = True) -> COD:
"""
@ -92,10 +91,9 @@ def get_channel(session: Session, *, channel_id: int = None, source_id: str = No
return channel
@optional_session
async def update_channel(session: Session, *, data: schema.ChannelPutRequest, channel_id: int) -> Channel:
"""Update a Channel's DB record"""
channel = Channel.find_by_id(channel_id, session)
channel = Channel.find_by_id(session, channel_id)
# Verify that the URL/Name/directory aren't taken
check_for_channel_conflicts(
@ -132,7 +130,6 @@ async def update_channel(session: Session, *, data: schema.ChannelPutRequest, ch
return channel
@optional_session
def create_channel(session: Session, data: schema.ChannelPostRequest, return_dict: bool = True) -> Union[Channel, dict]:
"""
Create a new Channel. Check for conflicts with existing Channels.
@ -176,9 +173,8 @@ def create_channel(session: Session, data: schema.ChannelPostRequest, return_dic
return channel.dict() if return_dict else channel
@optional_session(commit=True)
def delete_channel(session: Session, *, channel_id: int):
channel = Channel.find_by_id(channel_id, session=session)
channel = Channel.find_by_id(session, channel_id)
channel_dict = channel.dict()
channel.delete_with_videos()
@ -188,8 +184,7 @@ def delete_channel(session: Session, *, channel_id: int):
return channel_dict
@optional_session
async def search_channels_by_name(name: str, limit: int = 5, session: Session = None,
async def search_channels_by_name(session: Session, name: str, limit: int = 5,
order_by_video_count: bool = False) -> List[Channel]:
name = name or ''
name_no_spaces = ''.join(name.split(' '))
@ -222,8 +217,8 @@ async def search_channels_by_name(name: str, limit: int = 5, session: Session =
async def create_channel_download(channel_id: int, url: str, frequency: int, settings: dict):
"""Create Download for Channel."""
with get_db_session(commit=True) as session:
channel = Channel.find_by_id(channel_id, session=session)
download = channel.get_or_create_download(url, frequency, session, reset_attempts=True)
channel = Channel.find_by_id(session, channel_id)
download = channel.get_or_create_download(session, url, frequency, reset_attempts=True)
download.settings = settings
save_channels_config.activate_switch()
@ -237,8 +232,8 @@ async def update_channel_download(channel_id: int, download_id: int, url: str, f
"""Fetch Channel's Download, update its properties."""
with get_db_session(commit=True) as session:
# Ensure that the Channel exists.
Channel.find_by_id(channel_id, session=session)
download = Download.find_by_id(download_id, session=session)
Channel.find_by_id(session, channel_id)
download = Download.find_by_id(session, download_id)
download.url = url
download.frequency = frequency
download.settings = settings
@ -253,8 +248,7 @@ async def update_channel_download(channel_id: int, download_id: int, url: str, f
@wrol_mode_check
@optional_session
async def tag_channel(tag_name: str | None, directory: pathlib.Path | None, channel_id: int, session: Session = None):
async def tag_channel(session: Session, tag_name: str | None, directory: pathlib.Path | None, channel_id: int):
"""Add a Tag to a Channel, or remove a Tag from a Channel if no `tag_name` is provided.
Move the Channel to the new directory, if provided."""
@ -262,7 +256,7 @@ async def tag_channel(tag_name: str | None, directory: pathlib.Path | None, chan
if directory and flags.refreshing.is_set():
raise RefreshConflict('Refusing to move channel while file refresh is in progress')
channel = Channel.find_by_id(channel_id, session)
channel = Channel.find_by_id(session, channel_id)
# May also clear the tag if `tag_name` is None.
channel.set_tag(tag_name)
@ -287,8 +281,7 @@ async def tag_channel(tag_name: str | None, directory: pathlib.Path | None, chan
save_channels_config.activate_switch()
@optional_session
async def search_channels(tag_names: List[str], session: Session) -> List[Channel]:
async def search_channels(session: Session, tag_names: List[str]) -> List[Channel]:
"""Search Tagged Channels."""
channels = session.query(Channel).join(Collection).join(Tag).filter(Tag.name.in_(tag_names)).all()
return channels

6
modules/videos/channel/test/test_api.py
View file

@ -476,7 +476,7 @@ async def test_tag_channel(async_client, test_session, test_directory, channel_f
tag = await tag_factory('Tag Name')
v1, v2 = video_factory(title='video1', channel_id=channel.id), video_factory(title='video2', channel_id=channel.id)
# Create recurring download which uses the Channel's directory.
download = test_download_manager.recurring_download('https://example.com/1', 60, test_downloader.name,
download = test_download_manager.recurring_download(test_session, 'https://example.com/1', 60, test_downloader.name,
destination=str(test_directory / 'videos/Channel Name'))
test_session.commit()
save_channels_config()
@ -597,7 +597,7 @@ async def test_move_channel(async_client, test_session, test_directory, channel_
)
request, response = await async_client.put(f'/api/videos/channels/{channel.id}', json=body)
assert response.status_code == HTTPStatus.NO_CONTENT, response.content.decode()
channel = Channel.find_by_id(channel.id)
channel = Channel.find_by_id(test_session, channel.id)
assert str(channel.directory) == str(test_directory / 'videos/New Directory')
assert str(vid1.video_path).startswith(str(test_directory / 'videos/New Directory/'))
assert str(vid2.video_path).startswith(str(test_directory / 'videos/New Directory/'))
@ -625,7 +625,7 @@ async def test_move_channel_after_terminal_move(async_client, test_session, test
)
request, response = await async_client.put(f'/api/videos/channels/{channel.id}', json=body)
assert response.status_code == HTTPStatus.NO_CONTENT, response.content.decode()
channel = Channel.find_by_id(channel.id)
channel = Channel.find_by_id(test_session, channel.id)
vid1 = test_session.query(Video).one()
assert str(channel.directory) == str(test_directory / 'videos/New Directory')
assert str(vid1.video_path).startswith(str(test_directory / 'videos/New Directory/'))

24
modules/videos/channel/test/test_lib.py
View file

@ -20,7 +20,7 @@ from modules.videos.models import Channel
def test_get_bad_channel(test_session, test_directory, params):
"""A Channel cannot be found without data."""
with pytest.raises(UnknownChannel):
lib.get_channel(**params)
lib.get_channel(test_session, **params)
def test_get_channel(test_session, test_directory, channel_factory):
@ -38,19 +38,19 @@ def test_get_channel(test_session, test_directory, channel_factory):
test_session.commit()
channel = lib.get_channel(channel_id=channel_id_channel.id, return_dict=False)
channel = lib.get_channel(test_session, channel_id=channel_id_channel.id, return_dict=False)
assert channel == channel_id_channel
channel = lib.get_channel(source_id='the source id', return_dict=False)
channel = lib.get_channel(test_session, source_id='the source id', return_dict=False)
assert channel == source_id_channel
channel = lib.get_channel(url='https://example.com/channel', return_dict=False)
channel = lib.get_channel(test_session, url='https://example.com/channel', return_dict=False)
assert channel == url_channel
# Directory can be relative, or absolute.
channel = lib.get_channel(directory='some directory', return_dict=False)
channel = lib.get_channel(test_session, directory='some directory', return_dict=False)
assert channel == directory_channel
channel = lib.get_channel(directory=f'{test_directory}/some directory', return_dict=False)
channel = lib.get_channel(test_session, directory=f'{test_directory}/some directory', return_dict=False)
assert channel == directory_channel
# Channel is found using the priority of params, incorrect data is ignored.
@ -61,8 +61,8 @@ def test_get_channel(test_session, test_directory, channel_factory):
dict(channel_id=channel_id_channel.id, source_id='the source id'),
]
for p in params:
assert lib.get_channel(**p, return_dict=False) == channel_id_channel
assert lib.get_channel(**p)
assert lib.get_channel(test_session, **p, return_dict=False) == channel_id_channel
assert lib.get_channel(test_session, **p)
@pytest.mark.asyncio
@ -115,7 +115,7 @@ async def test_search_channels_by_name(test_session, channel_factory, video_fact
test_session.commit()
async def assert_search(name: str, channel_names: List[str], order_by_video_count=False):
channels = await lib.search_channels_by_name(name, order_by_video_count=order_by_video_count)
channels = await lib.search_channels_by_name(test_session, name, order_by_video_count=order_by_video_count)
for channel in channels:
channel_name = channel_names.pop(0)
if channel.name != channel_name:
@ -139,7 +139,7 @@ async def test_tag_channel_existing(async_client, test_session, test_directory,
assert str(channel.directory).startswith(f'{test_directory}/videos/one/')
tag = await tag_factory()
await lib.tag_channel(tag.name, channel_directory, channel.id, test_session)
await lib.tag_channel(test_session, tag.name, channel_directory, channel.id)
await await_switches()
# Channel's tag is saved to the config.
@ -161,7 +161,7 @@ async def test_tag_channel_without_directory_saves_config(
tag = await tag_factory(name='Tech')
# Tag without providing directory (directory=None)
await lib.tag_channel(tag.name, None, channel.id, test_session)
await lib.tag_channel(test_session, tag.name, None, channel.id)
test_session.commit()
await await_switches()
@ -189,7 +189,7 @@ async def test_untag_channel_without_directory_saves_config(
assert channel.tag is not None
# Untag without providing directory (tag_name=None, directory=None)
await lib.tag_channel(None, None, channel.id, test_session)
await lib.tag_channel(test_session, None, None, channel.id)
test_session.commit()
await await_switches()

2
modules/videos/channel/test/test_models.py
View file

@ -96,7 +96,7 @@ async def test_channel_download_relationships(test_session, download_channel):
assert test_session.query(Channel).one(), 'Channel should not have been deleted.'
# Create a Download again.
download_channel.get_or_create_download('https://example.com/1', 60, test_session, reset_attempts=True)
download_channel.get_or_create_download(test_session, 'https://example.com/1', 60, reset_attempts=True)
assert test_session.query(Download).count() == 1
# Delete the Channel, and any relationships.
download_channel.delete_with_videos()

14
modules/videos/conftest.py
View file

@ -53,7 +53,7 @@ def channel_factory(test_session, test_directory):
def factory(source_id: str = None, download_frequency: DownloadFrequency = None, url: str = None, name: str = None,
directory: pathlib.Path = None, tag_name: str = None):
name = name or str(uuid4())
tag = Tag.find_by_name(tag_name) if tag_name else None
tag = Tag.find_by_name(test_session, tag_name) if tag_name else None
tag_name = tag.name if tag else None
directory = directory or format_videos_destination(name, tag_name, url or f'https://example.com/{name}')
directory.mkdir(exist_ok=True, parents=True)
@ -79,7 +79,7 @@ def channel_factory(test_session, test_directory):
test_session.add(Directory(path=channel.directory, name=channel.directory.name))
assert channel.id and channel.url
if download_frequency:
download = channel.get_or_create_download(channel.url, download_frequency)
download = channel.get_or_create_download(test_session, channel.url, download_frequency)
assert download.url == channel.url
assert channel.downloads
test_session.commit()
@ -111,7 +111,7 @@ def download_channel(test_session, test_directory, video_download_manager) -> Ch
test_session.add(channel)
test_session.flush([channel])
assert channel and channel.id and channel.url
download = channel.get_or_create_download(channel.url, 60, test_session)
download = channel.get_or_create_download(test_session, channel.url, 60)
assert download.url == channel.url
download.frequency = DownloadFrequency.weekly
test_session.commit()
@ -126,7 +126,7 @@ def simple_video(test_session, test_directory, simple_channel, video_file) -> Vi
video = Video.from_paths(test_session, video_path)
video.channel = simple_channel
test_session.commit()
video = Video.find_by_id(video.id, session=test_session)
video = Video.find_by_id(test_session, video.id)
return video
@ -145,7 +145,7 @@ def video_factory(test_session, test_directory):
path = pathlib.Path(with_video_file) if not isinstance(with_video_file, pathlib.Path) else with_video_file
elif channel_id:
# Put the video in its Channel's directory.
channel = Channel.find_by_id(channel_id)
channel = Channel.find_by_id(test_session, channel_id)
path = (channel.directory or test_directory) / f'{title}.mp4'
path.parent.mkdir(exist_ok=True, parents=True)
else:
@ -182,10 +182,10 @@ def video_factory(test_session, test_directory):
video.channel_id = channel_id
video.source_id = source_id or title
video.file_group.published_datetime = upload_date
video.validate()
video.validate(test_session)
for tag_name in tag_names:
video.add_tag(tag_name)
video.add_tag(test_session, tag_name)
return video

123
modules/videos/downloader.py
View file

@ -22,7 +22,6 @@ from wrolpi.common import logger, get_media_directory, escape_file_name, resolve
trim_file_name, cached_multiprocessing_result, get_absolute_media_path
from wrolpi.dates import now
from wrolpi.db import get_db_session
from wrolpi.db import optional_session
from wrolpi.downloader import Downloader, Download, DownloadResult
from wrolpi.errors import UnrecoverableDownloadError
from wrolpi.files.lib import glob_shared_stem, split_path_stem_and_suffix
@ -147,7 +146,7 @@ async def fetch_video_duration(url: str) -> int:
otherwise, fetches the duration using yt-dlp."""
with get_db_session() as session:
# Join with Video to get duration; a URL may have other files (such as an Archive).
video = Video.get_by_url(url, session=session)
video = Video.get_by_url(session, url)
if video and (duration := video.file_group.length):
logger.debug(f'Using already-known duration {duration} for {url} from FileGroup')
return duration
@ -199,7 +198,7 @@ class ChannelDownloader(Downloader, ABC):
if not name:
raise ValueError(f'Could not find name')
channel_source_id = info.get('channel_id') or info.get('id')
channel = get_or_create_channel(channel_source_id, download.url, name, channel_tag_name)
channel = get_or_create_channel(session, channel_source_id, download.url, name, channel_tag_name)
channel.dict() # get all attributes while we have the session.
location = f'/videos/channel/{channel.id}/video' if channel and channel.id else None
@ -225,7 +224,7 @@ class ChannelDownloader(Downloader, ABC):
else:
logger.debug('Not updating channel because this is a playlist')
downloads = self.get_missing_videos(download)
downloads = self.get_missing_videos(session, download)
return DownloadResult(
success=True,
location=location,
@ -261,7 +260,7 @@ class ChannelDownloader(Downloader, ABC):
session.refresh(channel)
@staticmethod
def get_missing_videos(download: Download) -> List[str]:
def get_missing_videos(session: Session, download: Download) -> List[str]:
"""
Return all URLs of Videos in the `info_json` which need to be downloaded.
"""
@ -323,7 +322,7 @@ class ChannelDownloader(Downloader, ABC):
downloads = [normalize_video_url(i) for i in downloads]
# Only download those that have not yet been downloaded.
already_downloaded = [i.url for i in video_downloader.already_downloaded(*downloads)]
already_downloaded = [i.url for i in video_downloader.already_downloaded(session, *downloads)]
downloads = [i for i in downloads if i not in already_downloaded]
return downloads
@ -343,65 +342,64 @@ class VideoDownloader(Downloader, ABC):
def __repr__(self):
return f'<VideoDownloader>'
@optional_session
def already_downloaded(self, *urls: str, session: Session = None) -> List:
def already_downloaded(self, session: Session, *urls: str) -> List:
# We only consider a video record with a video file as "downloaded".
file_groups = list(session.query(FileGroup).filter(FileGroup.url.in_(urls), FileGroup.model == 'video'))
return file_groups
async def do_download(self, download: Download) -> DownloadResult:
if download.attempts >= 10:
raise UnrecoverableDownloadError('Max download attempts reached')
url = normalize_video_url(download.url)
# Copy settings into Download (they may be from the ChannelDownloader)
settings = download.settings or dict()
download.destination = download.destination or settings.get('destination')
tag_names = download.tag_names = download.tag_names or settings.get('tag_names')
# Video may have been downloaded previously, get its location for error reporting.
with get_db_session() as session:
video_ = Video.get_by_url(url, session)
if download.attempts >= 10:
raise UnrecoverableDownloadError('Max download attempts reached')
url = normalize_video_url(download.url)
# Copy settings into Download (they may be from the ChannelDownloader)
settings = download.settings or dict()
download.destination = download.destination or settings.get('destination')
tag_names = download.tag_names = download.tag_names or settings.get('tag_names')
# Video may have been downloaded previously, get its location for error reporting.
video_ = Video.get_by_url(session, url)
location = video_.location if video_ else None
try:
download.info_json = download.info_json or extract_info(url)
except yt_dlp.utils.DownloadError as e:
# Video may be private.
try:
raise RuntimeError('Failed to extract_info') from e
except Exception as e:
return DownloadResult(
success=False,
location=location,
error='\n'.join(traceback.format_exception(e)),
)
download.info_json = download.info_json or extract_info(url)
except yt_dlp.utils.DownloadError as e:
# Video may be private.
try:
raise RuntimeError('Failed to extract_info') from e
except Exception as e:
return DownloadResult(
success=False,
location=location,
error='\n'.join(traceback.format_exception(e)),
)
if not download.info_json:
raise ValueError(f'Cannot download video with no info_json.')
if not download.info_json:
raise ValueError(f'Cannot download video with no info_json.')
channel, channel_directory, channel_id, destination, settings = await self._get_channel(download)
channel, channel_directory, channel_id, destination, settings = await self._get_channel(session, download)
if destination:
# Download to the directory specified in the settings.
out_dir = destination
logger.debug(f'Downloading {url} to destination {destination} from settings')
elif channel:
out_dir = channel_directory
logger.debug(f'Downloading {url} to channel directory: {channel_directory}')
else:
# Download to the default directory if this video has no channel.
out_dir = get_no_channel_directory()
logger.debug(f'Downloading {url} to default directory')
if destination:
# Download to the directory specified in the settings.
out_dir = destination
logger.debug(f'Downloading {url} to destination {destination} from settings')
elif channel:
out_dir = channel_directory
logger.debug(f'Downloading {url} to channel directory: {channel_directory}')
else:
# Download to the default directory if this video has no channel.
out_dir = get_no_channel_directory()
logger.debug(f'Downloading {url} to default directory')
# Make output directory. (Maybe string from settings)
out_dir = out_dir if isinstance(out_dir, pathlib.Path) else pathlib.Path(out_dir)
out_dir = get_absolute_media_path(out_dir)
out_dir.mkdir(exist_ok=True, parents=True)
# Make output directory. (Maybe string from settings)
out_dir = out_dir if isinstance(out_dir, pathlib.Path) else pathlib.Path(out_dir)
out_dir = get_absolute_media_path(out_dir)
out_dir.mkdir(exist_ok=True, parents=True)
if tag_names and logger.isEnabledFor(logging.DEBUG):
logger.debug(f'Downloading {url} with {tag_names=}')
if tag_names and logger.isEnabledFor(logging.DEBUG):
logger.debug(f'Downloading {url} with {tag_names=}')
config = get_videos_downloader_config()
@ -526,12 +524,12 @@ class VideoDownloader(Downloader, ABC):
with get_db_session(commit=True) as session:
# Second session is started because SQLAlchemy will forget what we have done.
video = Video.get_by_id(video_id, session)
video = Video.get_by_id(session, video_id)
if tag_names:
existing_names = video.file_group.tag_names
for name in tag_names:
if name not in existing_names:
video.add_tag(name)
video.add_tag(session, name)
location = video.location
await video.get_ffprobe_json()
@ -589,7 +587,7 @@ class VideoDownloader(Downloader, ABC):
return result
@staticmethod
async def _get_channel(download: Download) \
async def _get_channel(session: Session, download: Download) \
-> Tuple[Optional[Channel], Optional[pathlib.Path], Optional[int], Optional[pathlib.Path], dict]:
settings = download.settings or dict()
channel_tag_name = i[0] if (i := settings.get('channel_tag_name')) else None
@ -603,7 +601,8 @@ class VideoDownloader(Downloader, ABC):
if channel_name or channel_source_id or channel_url:
# Try to find the channel via info_json from yt-dlp.
try:
channel = get_or_create_channel(source_id=channel_source_id, url=channel_url, name=channel_name,
channel = get_or_create_channel(session, source_id=channel_source_id, url=channel_url,
name=channel_name,
tag_name=channel_tag_name)
if channel:
found_channel = 'yt_dlp'
@ -615,7 +614,7 @@ class VideoDownloader(Downloader, ABC):
if not channel and destination:
# Destination may find Channel if not already found.
try:
channel = get_channel(directory=destination, return_dict=False)
channel = get_channel(session, directory=destination, return_dict=False)
found_channel = 'download_settings_directory'
logger.debug(f'Found a channel that shares the destination directory {channel=}')
except UnknownChannel:
@ -630,10 +629,9 @@ class VideoDownloader(Downloader, ABC):
try:
# Find Channel by collection_id or URL
if collection_id:
with get_db_session() as session:
channel = session.query(Channel).filter_by(collection_id=collection_id).one_or_none()
channel = session.query(Channel).filter_by(collection_id=collection_id).one_or_none()
if not channel and channel_url:
channel = get_channel(url=channel_url, return_dict=False)
channel = get_channel(session, url=channel_url, return_dict=False)
logger.debug(f'Found channel with collection_id or channel url {channel=}')
except UnknownChannel:
# We do not need a Channel if we have a destination directory.
@ -811,14 +809,15 @@ channel_downloader = ChannelDownloader()
video_downloader = VideoDownloader()
def get_or_create_channel(source_id: str = None, url: str = None, name: str = None, tag_name: str = None) -> Channel:
def get_or_create_channel(session: Session, source_id: str = None, url: str = None, name: str = None,
tag_name: str = None) -> Channel:
"""
Attempt to find a Channel using the provided params. The params are in order of reliability.
Creates a new Channel if one cannot be found.
"""
try:
channel = get_channel(source_id=source_id, url=url, name=name, return_dict=False)
channel = get_channel(session, source_id=source_id, url=url, name=name, return_dict=False)
# Get all properties while we have the session.
channel.dict()
return channel
@ -839,7 +838,7 @@ def get_or_create_channel(source_id: str = None, url: str = None, name: str = No
directory=str(channel_directory.relative_to(get_media_directory())),
tag_name=tag_name,
)
channel = create_channel(data=data, return_dict=False)
channel = create_channel(session, data=data, return_dict=False)
# Create the directory now that the channel is approved.
channel_directory.mkdir(exist_ok=True)
# Get all properties while we have the session.

9
modules/videos/lib.py
View file

@ -84,7 +84,7 @@ def extract_video_info_json(video: Video) -> VideoInfoJSON:
EXTRACT_SUBTITLES = False
def validate_video(video: Video, channel_generate_poster: bool):
def validate_video(session: Session, video: Video, channel_generate_poster: bool):
"""
Validate a single video.
@ -116,6 +116,7 @@ def validate_video(video: Video, channel_generate_poster: bool):
from modules.videos.channel.lib import get_channel
try:
if channel := get_channel(
session,
source_id=video_info_json.channel_source_id,
url=video_info_json.channel_url,
directory=video.video_path.parent,
@ -292,7 +293,7 @@ class ChannelsConfig(ConfigFile):
# Batch create/update Collections
from wrolpi.collections import Collection
collections = Collection.batch_from_config(channel_data_list, session=session)
collections = Collection.batch_from_config(session, channel_data_list)
# Pre-fetch existing Channels by collection_id in ONE query
collection_ids = [c.id for c in collections]
@ -624,7 +625,7 @@ async def fetch_channel_source_id(channel_id: int):
try:
with get_db_session() as session:
from modules.videos.channel.lib import get_channel
channel = Channel.find_by_id(channel_id)
channel = Channel.find_by_id(session, channel_id)
channel.source_id = channel.source_id or get_channel_source_id(channel.url)
if not channel.source_id:
channel_import_logger.warning(f'Unable to fetch source_id for {channel.url}')
@ -682,7 +683,7 @@ def link_channel_and_downloads(session: Session, channel_: Type[Base] = Channel,
# Associate any Download which shares a Channel's URL.
for download in downloads:
channel = channel_.get_by_url(download.url, session)
channel = channel_.get_by_url(session, download.url)
if channel and not download.collection_id:
download.collection_id = channel.collection_id
need_commit = True

86
modules/videos/models.py
View file

@ -9,7 +9,7 @@ from sqlalchemy.orm.collections import InstrumentedList
from modules.videos.errors import UnknownVideo, UnknownChannel
from wrolpi.captions import read_captions
from wrolpi.common import Base, ModelHelper, logger, get_media_directory, background_task, replace_file
from wrolpi.db import get_db_curs, get_db_session, optional_session
from wrolpi.db import get_db_curs, get_db_session
from wrolpi.downloader import Download
from wrolpi.files.lib import refresh_files, split_path_stem_and_suffix
from wrolpi.files.models import FileGroup
@ -99,14 +99,14 @@ class Video(ModelHelper, Base):
return file_group.mimetype.startswith('video/')
@staticmethod
def do_model(file_group: FileGroup, session: Session) -> 'Video':
def do_model(session: Session, file_group: FileGroup) -> 'Video':
# Check if a Video already exists for this FileGroup
video = session.query(Video).filter_by(file_group_id=file_group.id).one_or_none()
if not video:
# Create new Video if it doesn't exist
video = Video(file_group_id=file_group.id, file_group=file_group)
session.add(video)
video.validate()
video.validate(session)
file_group.indexed = True
video.flush()
return video
@ -208,12 +208,12 @@ class Video(ModelHelper, Base):
break
# Fetch the videos by id, if they exist.
previous_video = Video.find_by_id(previous_id, session) if previous_id else None
next_video = Video.find_by_id(next_id, session) if next_id else None
previous_video = Video.find_by_id(session, previous_id) if previous_id else None
next_video = Video.find_by_id(session, next_id) if next_id else None
return previous_video, next_video
def validate(self, session: Session = None):
def validate(self, session: Session):
"""Perform a validation of this video and it's files."""
if not self.file_group.primary_path:
# Can't validate if there is no video file.
@ -222,7 +222,7 @@ class Video(ModelHelper, Base):
logger.debug(f'Validating {self}')
from .lib import validate_video
try:
validate_video(self, self.channel.generate_posters if self.channel else False)
validate_video(session, self, self.channel.generate_posters if self.channel else False)
except Exception as e:
logger.error(f'Failed to validate video {self}', exc_info=e)
if PYTEST:
@ -239,7 +239,7 @@ class Video(ModelHelper, Base):
# If this Video is in a Channel's directory, then it is part of that Channel.
session: Session = session or Session.object_session(self)
if session and (channel := Channel.get_by_path(self.video_path.parent, session)):
if session and (channel := Channel.get_by_path(session, self.video_path.parent)):
self.channel = channel
logger.debug(f'{self} has Channel {channel}')
@ -263,7 +263,7 @@ class Video(ModelHelper, Base):
session.add(video)
session.flush([video, file_group])
video.validate()
video.validate(session)
session.flush([video, ])
return video
@ -322,46 +322,40 @@ class Video(ModelHelper, Base):
return caption_text
@staticmethod
@optional_session
def get_by_path(path, session: Session) -> Optional['Video']:
def get_by_path(session: Session, path) -> Optional['Video']:
video = session.query(Video) \
.join(FileGroup, FileGroup.id == Video.file_group_id) \
.filter(FileGroup.primary_path == path).one_or_none()
return video
@staticmethod
@optional_session
def get_by_url(url: str, session: Session) -> Optional['Video']:
def get_by_url(session: Session, url: str) -> Optional['Video']:
video = session.query(Video) \
.join(FileGroup, FileGroup.id == Video.file_group_id) \
.filter(FileGroup.url == url).one_or_none()
return video
@staticmethod
@optional_session
def get_by_id(id_: int, session: Session = None) -> Optional['Video']:
def get_by_id(session: Session, id_: int) -> Optional['Video']:
"""Attempt to find a Video with the provided id. Returns None if it cannot be found."""
video = session.query(Video).filter(Video.id == id_).one_or_none()
return video
@staticmethod
@optional_session
def find_by_id(id_: int, session: Session = None) -> 'Video':
def find_by_id(session: Session, id_: int) -> 'Video':
"""Find a Video with the provided id, raises an exception if it cannot be found.
@raise UnknownVideo: if the Video can not be found"""
video = Video.get_by_id(id_, session)
video = Video.get_by_id(session, id_)
if not video:
raise UnknownVideo(f'Cannot find Video with id {id_}')
return video
def add_tag(self, tag_id_or_name: int | str, session: Session = None) -> TagFile:
session = session or Session.object_session(self)
return self.file_group.add_tag(tag_id_or_name, session=session)
def add_tag(self, session: Session, tag_id_or_name: int | str) -> TagFile:
return self.file_group.add_tag(session, tag_id_or_name)
def untag(self, tag_id_or_name: int | str, session: Session = None):
session = session or Session.object_session(self)
self.file_group.untag(tag_id_or_name, session)
def untag(self, session: Session, tag_id_or_name: int | str):
self.file_group.untag(session, tag_id_or_name)
async def get_ffprobe_json(self) -> dict:
"""Return the ffprobe json object if previously stored.
@ -468,7 +462,7 @@ class Video(ModelHelper, Base):
largest_video = duplicate_videos.pop(largest_video_idx)
# Delete the video at the destination, if it is not the largest.
existing_video = Video.get_by_path(video_path, session=session)
existing_video = Video.get_by_path(session, video_path)
if existing_video and largest_video != existing_video:
delete_video(existing_video)
for idx, video in enumerate(duplicate_videos):
@ -480,7 +474,7 @@ class Video(ModelHelper, Base):
largest_video.file_group.move(video_path)
for tag_name in tag_names:
if tag_name not in largest_video.file_group.tag_names:
largest_video.add_tag(tag_name, session)
largest_video.add_tag(session, tag_name)
logger.warning(f'Deleting duplicate videos: {duplicate_videos}')
@ -660,7 +654,7 @@ class Channel(ModelHelper, Base):
downloads = data.pop('downloads', [])
if tag_name := data.pop('tag_name', None):
self.tag = Tag.find_by_name(tag_name, session)
self.tag = Tag.find_by_name(session, tag_name)
for key, value in data.items():
setattr(self, key, value)
@ -677,7 +671,7 @@ class Channel(ModelHelper, Base):
if isinstance(download, dict):
url = download['url']
# Use download frequency from downloads config before the channels config.
if existing_download := Download.get_by_url(url, session=session):
if existing_download := Download.get_by_url(session, url):
frequency = existing_download.frequency
else:
frequency = download.get('frequency')
@ -691,7 +685,7 @@ class Channel(ModelHelper, Base):
logger.error(f'Refusing to create Download for Channel without frequency: {url}')
continue
download = self.get_or_create_download(url, frequency, session=session, reset_attempts=True)
download = self.get_or_create_download(session, url, frequency, reset_attempts=True)
download.collection_id = self.collection_id
session.add(download)
@ -754,7 +748,7 @@ class Channel(ModelHelper, Base):
logger.error(f'Refusing to create Download for Channel without frequency: {url}')
continue
download_obj = self.get_or_create_download(url, frequency, session=session, reset_attempts=True)
download_obj = self.get_or_create_download(session, url, frequency, reset_attempts=True)
download_obj.collection_id = self.collection_id
session.add(download_obj)
@ -776,29 +770,25 @@ class Channel(ModelHelper, Base):
return config
@classmethod
def from_config(cls, data: dict, session: Session = None) -> 'Channel':
def from_config(cls, session: Session, data: dict) -> 'Channel':
"""
Create or update a Channel from config data. This also creates/updates the Collection.
Args:
data: Config dict containing channel metadata (name, directory, url, source_id, etc.)
session: Database session
data: Config dict containing channel metadata (name, directory, url, source_id, etc.)
Returns:
The created or updated Channel
"""
from wrolpi.collections import Collection
from wrolpi.db import get_db_session
if session is None:
session = get_db_session()
# Ensure this is treated as a channel collection
data = data.copy()
data['kind'] = 'channel'
# Create or update the Collection first
collection = Collection.from_config(data, session=session)
collection = Collection.from_config(session, data)
# Extract Channel-specific fields
url = data.get('url')
@ -823,7 +813,7 @@ class Channel(ModelHelper, Base):
return channel
@staticmethod
def get_by_path(path: pathlib.Path, session: Session) -> Optional['Channel']:
def get_by_path(session: Session, path: pathlib.Path) -> Optional['Channel']:
if not path:
raise RuntimeError('Must provide path to get Channel')
path = pathlib.Path(path) if isinstance(path, str) else path
@ -883,9 +873,9 @@ class Channel(ModelHelper, Base):
# Clear tag.
self.tag = None
elif isinstance(tag_id_or_name, int):
self.tag = Tag.find_by_id(tag_id_or_name, session)
self.tag = Tag.find_by_id(session, tag_id_or_name)
elif isinstance(tag_id_or_name, str):
self.tag = Tag.find_by_name(tag_id_or_name, session)
self.tag = Tag.find_by_name(session, tag_id_or_name)
# Copy the id in case some code needs the id before flush.
self.tag_id = self.tag.id if self.tag else None
return self.tag
@ -916,19 +906,17 @@ class Channel(ModelHelper, Base):
background_task(_())
@staticmethod
@optional_session
def get_by_id(id_: int, session: Session = None) -> Optional['Channel']:
def get_by_id(session: Session, id_: int) -> Optional['Channel']:
"""Attempt to find a Channel with the provided id. Returns None if it cannot be found."""
channel = session.query(Channel).filter_by(id=id_).one_or_none()
return channel
@staticmethod
@optional_session
def find_by_id(id_: int, session: Session = None) -> 'Channel':
def find_by_id(session: Session, id_: int) -> 'Channel':
"""Find a Channel with the provided id, raises an exception when no Channel is found.
@raise UnknownChannel: if the channel can not be found"""
if channel := Channel.get_by_id(id_, session=session):
if channel := Channel.get_by_id(session, id_):
return channel
raise UnknownChannel(f'Cannot find channel with id {id_}')
@ -941,8 +929,7 @@ class Channel(ModelHelper, Base):
elif len(matching_entries) > 1:
raise RuntimeError(f'More than one info_json entry matches {video_source_id}')
@optional_session
def get_or_create_download(self, url: str, frequency: int, session: Session = None,
def get_or_create_download(self, session: Session, url: str, frequency: int,
reset_attempts: bool = False) -> Download:
"""Get a Download record, if it does not exist, create it. Create a Download if necessary
which goes into this Channel's directory.
@ -952,7 +939,7 @@ class Channel(ModelHelper, Base):
from modules.videos.downloader import ChannelDownloader, VideoDownloader
return self.collection.get_or_create_download(
url, frequency, session=session,
session, url, frequency,
reset_attempts=reset_attempts,
downloader_name=ChannelDownloader.name,
sub_downloader_name=VideoDownloader.name
@ -966,8 +953,7 @@ class Channel(ModelHelper, Base):
return f'https://www.youtube.com/feeds/videos.xml?channel_id={self.source_id}'
@staticmethod
@optional_session
def get_by_url(url: str, session: Session = None) -> Optional['Channel']:
def get_by_url(session: Session, url: str) -> Optional['Channel']:
if not url:
raise RuntimeError('Must provide URL to get Channel')
channel = session.query(Channel).filter_by(url=url).one_or_none()

16
modules/videos/test/test_api.py
View file

@ -106,7 +106,7 @@ async def test_refresh_videos(async_client, test_session, test_directory, simple
video1_video_path = str(video1.video_path)
video1.poster_path.unlink()
await async_client.post('/api/files/refresh')
video1 = Video.get_by_path(video1_video_path, test_session)
video1 = Video.get_by_path(test_session, video1_video_path)
assert not video1.poster_path
@ -141,9 +141,9 @@ async def test_channels_with_videos(test_session, async_client, test_directory,
assert {i[0] for i in test_session.query(Collection.name).join(Channel).filter(Collection.kind == 'channel')} == {
'channel1', 'channel2'}, 'Channels were changed.'
vid1 = Video.get_by_path(vid1_path)
vid2 = Video.get_by_path(vid2_path)
vid3 = Video.get_by_path(vid3_path)
vid1 = Video.get_by_path(test_session, vid1_path)
vid2 = Video.get_by_path(test_session, vid2_path)
vid3 = Video.get_by_path(test_session, vid3_path)
assert vid1.channel == channel1
assert vid2.channel == channel2
@ -234,13 +234,13 @@ async def test_search_videos(test_session, video_factory, assert_video_search, s
vid3: Video = video_factory(upload_date='2022-09-18', with_video_file=True, title='vid3')
tag1, tag2 = await tag_factory(), await tag_factory()
vid1.add_tag(tag1.id)
vid1.add_tag(test_session, tag1.id)
# vid2 has two tags
vid2.add_tag(tag1.id)
vid2.add_tag(tag2.id)
vid2.add_tag(test_session, tag1.id)
vid2.add_tag(test_session, tag2.id)
vid3.add_tag(tag2.id)
vid3.add_tag(test_session, tag2.id)
test_session.commit()
assert test_session.query(Video).count() == 3

12
modules/videos/test/test_common.py
View file

@ -122,8 +122,10 @@ async def test_update_view_count(test_session, channel_factory, video_factory):
# An outdated view count will be overwritten.
vid = test_session.query(Video).filter_by(id=1).one()
vid.view_count = 8
test_session.commit() # Commit so raw SQL in update_view_counts_and_censored doesn't deadlock
check_view_counts({'vid1': 8, 'vid2': 11, 'vid3': 13, 'vid4': 14})
await update_view_counts_and_censored(channel1.id)
test_session.expire_all() # Expire cached values to see raw SQL changes
check_view_counts({'vid1': 10, 'vid2': 11, 'vid3': 13, 'vid4': 14})
@ -192,15 +194,15 @@ async def test_import_channel_downloads(await_switches, test_session, channel_fa
assert next_download == str(download.next_download)
# Creating Download that matches Channel2's URL means they are related. Delete it and it should be re-created.
channel2.get_or_create_download(channel2.url, 60, test_session)
channel2.get_or_create_download(test_session, channel2.url, 60)
save_channels_config()
Download.find_by_url(channel2.url).delete(add_to_skip_list=False)
Download.find_by_url(test_session, channel2.url).delete(add_to_skip_list=False)
test_session.commit()
await await_switches()
# Missing Download is re-created on import.
import_channels_config()
channel2 = Channel.get_by_id(channel2.id)
channel2 = Channel.get_by_id(test_session, channel2.id)
downloads = test_session.query(Download).all()
assert len(downloads) == 1, downloads
assert len(channel1.downloads) == 1
@ -208,8 +210,8 @@ async def test_import_channel_downloads(await_switches, test_session, channel_fa
assert downloads[0].downloader == 'video_channel'
# Add a Download to Channel2 which does not match Channel.url.
channel2.get_or_create_download('https://example.org', 60, session=test_session)
channel2.get_or_create_download('https://example.com/channel2', 60, session=test_session)
channel2.get_or_create_download(test_session, 'https://example.org', 60)
channel2.get_or_create_download(test_session, 'https://example.com/channel2', 60)
test_session.commit()
save_channels_config()
await await_switches()

39
modules/videos/test/test_downloader.py
View file

@ -98,7 +98,7 @@ async def test_download_no_channel(test_session, video_download_manager, test_di
mock_video_extract_info.return_value = info_json
with mock.patch('modules.videos.downloader.VideoDownloader.prepare_filename') as mock_prepare_filename:
mock_prepare_filename.return_value = (video_path, {'id': 'foo'})
video_download_manager.create_download(url, video_downloader.name)
video_download_manager.create_download(test_session, url, video_downloader.name)
await video_download_manager.wait_for_all_downloads()
mock_video_process_runner.assert_called_once()
@ -119,7 +119,7 @@ async def test_download_video_tags(test_session, video_download_manager, video_f
with mock.patch('modules.videos.downloader.VideoDownloader.prepare_filename') as mock_prepare_filename:
mock_prepare_filename.return_value = (video_path, {'id': 'foo'})
settings = dict(tag_names=[tag1.name, tag2.name])
video_download_manager.create_download(url, video_downloader.name, settings=settings)
video_download_manager.create_download(test_session, url, video_downloader.name, settings=settings)
await video_download_manager.wait_for_all_downloads()
mock_video_process_runner.assert_called_once()
@ -141,7 +141,7 @@ async def test_download_channel(test_session, test_directory, simple_channel, vi
mock_video_extract_info.return_value = example_channel_json
with mock.patch('modules.videos.downloader.get_channel') as mock_get_channel:
mock_get_channel.return_value = simple_channel
simple_channel.get_or_create_download(url, 60, test_session)
simple_channel.get_or_create_download(test_session, url, 60)
# Download channel.
await video_download_manager.wait_for_all_downloads()
# Download videos.
@ -166,7 +166,7 @@ async def test_download_channel(test_session, test_directory, simple_channel, vi
test_session.commit()
with mock.patch('modules.videos.downloader.get_channel') as mock_get_channel:
mock_get_channel.return_value = simple_channel
download = simple_channel.get_or_create_download(url, 60, test_session)
download = simple_channel.get_or_create_download(test_session, url, 60)
download.frequency = 100
download.settings = {'title_include': '2'}
await video_download_manager.wait_for_all_downloads()
@ -184,7 +184,7 @@ async def test_download_channel(test_session, test_directory, simple_channel, vi
test_session.commit()
with mock.patch('modules.videos.downloader.get_channel') as mock_get_channel:
mock_get_channel.return_value = simple_channel
download = simple_channel.get_or_create_download(url, 60, test_session)
download = simple_channel.get_or_create_download(test_session, url, 60)
download.frequency = 100
download.settings = {'title_exclude': '2'}
await video_download_manager.wait_for_all_downloads()
@ -233,30 +233,30 @@ async def test_get_or_create_channel(async_client, test_session, test_directory,
(dict(name='qux'), c4),
]
for kwargs, expected in tests:
channel = get_or_create_channel(**kwargs)
channel = get_or_create_channel(test_session, **kwargs)
assert expected.id == channel.id, f'Expected {expected} for {kwargs} but got {channel}'
# A new channel is created. It will not be automatically downloaded.
channel = get_or_create_channel(source_id='quux', name='quux', url='https://example.org')
channel = get_or_create_channel(test_session, source_id='quux', name='quux', url='https://example.org')
assert channel.id == 5
assert channel.source_id == 'quux'
assert channel.name == 'quux'
assert channel.url == 'https://example.org'
# New channel can be retrieved.
assert get_or_create_channel(source_id='quux') == channel
assert get_or_create_channel(test_session, source_id='quux') == channel
# A Channel can be created with a tag name.
channel = get_or_create_channel(name='One Channel', tag_name=one.name)
channel = get_or_create_channel(test_session, name='One Channel', tag_name=one.name)
await await_switches()
assert channel.directory == test_directory / f'videos/one/One Channel'
assert channel.tag_name == one.name and channel.tag == one
# Getting the same Channel will not change the tag.
channel = get_or_create_channel(name='One Channel', tag_name=two.name)
channel = get_or_create_channel(test_session, name='One Channel', tag_name=two.name)
assert channel.directory == test_directory / f'videos/one/One Channel'
assert channel.tag_name == one.name and channel.tag == one
channel = get_or_create_channel(name='One Channel')
channel = get_or_create_channel(test_session, name='One Channel')
assert channel.directory == test_directory / f'videos/one/One Channel'
assert channel.directory.is_dir()
assert channel.tag_name == one.name and channel.tag == one
@ -275,7 +275,7 @@ def test_bad_downloader(test_session, video_download_manager):
Attempting to use an unknown downloader should raise an error.
"""
with pytest.raises(InvalidDownload):
video_download_manager.create_download('https://example.com', downloader_name='bad downloader')
video_download_manager.create_download(test_session, 'https://example.com', downloader_name='bad downloader')
@pytest.mark.asyncio
@ -303,7 +303,7 @@ async def test_video_download(test_session, test_directory, mock_video_extract_i
mock_video_extract_info.return_value = example_video_json
mock_prepare_filename.return_value = (video_path, {'id': 'foo'})
video_download_manager.create_download(url, video_downloader.name)
video_download_manager.create_download(test_session, url, video_downloader.name)
await video_download_manager.wait_for_all_downloads()
mock_video_process_runner.assert_called_once()
@ -338,7 +338,7 @@ async def test_download_result(test_session, test_directory, video_download_mana
mock_video_extract_info.return_value = example_video_json
with mock.patch('modules.videos.downloader.VideoDownloader.prepare_filename') as mock_prepare_filename:
mock_prepare_filename.return_value = [video_file, {'id': 'foo'}]
video_download_manager.create_download('https://example.com', video_downloader.name)
video_download_manager.create_download(test_session, 'https://example.com', video_downloader.name)
await video_download_manager.wait_for_all_downloads()
# Sleep to allow background tasks to finish.
await asyncio.sleep(1)
@ -362,7 +362,7 @@ async def test_download_destination(test_session, test_directory, video_download
# This result will be ignored during the test, we just need some value so the download does not fail.
mock_video_prepare_filename.return_value = str(test_directory / 'test video.mp4')
video_download_manager.create_download('https://example.com/1', downloader_name=VideoDownloader.name)
video_download_manager.create_download(test_session, 'https://example.com/1', downloader_name=VideoDownloader.name)
await video_download_manager.wait_for_all_downloads()
# Output directory matches the channel directory.
assert mock_video_prepare_filename.call_args_list[0].kwargs['ydl'].params['outtmpl']['default'] \
@ -370,7 +370,7 @@ async def test_download_destination(test_session, test_directory, video_download
mock_video_prepare_filename.reset_mock()
video_download_manager.create_download('https://example.com/2', downloader_name=VideoDownloader.name,
video_download_manager.create_download(test_session, 'https://example.com/2', downloader_name=VideoDownloader.name,
destination=f'{test_directory}/custom')
await video_download_manager.wait_for_all_downloads()
# Output directory matches the custom directory specified.
@ -446,7 +446,8 @@ async def test_download_playlist(test_session, test_directory, mock_video_extrac
"""All videos in a playlist can be downloaded for it's Channel."""
download = Download(url='https://example.com/playlist-url')
test_session.add(download)
channel = get_or_create_channel(example_playlist_json['channel_id'], download.url, example_channel_json['uploader'])
channel = get_or_create_channel(test_session, source_id=example_playlist_json['channel_id'], url=download.url,
name=example_channel_json['uploader'])
video = Video.from_paths(test_session, video_file)
video.channel_id = channel.id
video.source_id = 'video 1 id'
@ -577,7 +578,7 @@ async def test_video_download_cookies(test_session, test_directory, mock_video_e
mock_video_extract_info.return_value = example_video_json
mock_prepare_filename.return_value = (video_path, {'id': 'foo'})
video_download_manager.create_download(url, video_downloader.name, settings=settings)
video_download_manager.create_download(test_session, url, video_downloader.name, settings=settings)
await video_download_manager.wait_for_all_downloads()
download, cmd, out_dir = mock_video_process_runner.call_args[0]
@ -604,7 +605,7 @@ async def test_video_download_always_use_cookies(test_session, test_directory, m
mock_video_extract_info.return_value = example_video_json
mock_prepare_filename.return_value = (video_path, {'id': 'foo'})
video_download_manager.create_download(url, video_downloader.name)
video_download_manager.create_download(test_session, url, video_downloader.name)
await video_download_manager.wait_for_all_downloads()
download, cmd, out_dir = mock_video_process_runner.call_args[0]

10
modules/videos/test/test_lib.py
View file

@ -73,7 +73,7 @@ async def test_video_factory(test_session, video_factory, channel_factory):
assert video.source_id == 'some id'
def test_validate_video(test_directory, video_factory, image_bytes_factory, video_file):
def test_validate_video(test_session, test_directory, video_factory, image_bytes_factory, video_file):
"""A video poster will be generated only if the channel permits."""
video_file = video_file.rename(test_directory / 'Channel Name_20050607_1234567890_The Title.mp4')
vid1 = video_factory(with_video_file=video_file, with_info_json=True)
@ -81,7 +81,7 @@ def test_validate_video(test_directory, video_factory, image_bytes_factory, vide
vid1.source_id = None
assert not vid1.poster_path
validate_video(vid1, True)
validate_video(test_session, vid1, True)
assert vid1.video_path == video_file
assert vid1.poster_path, 'Poster was not created'
assert vid1.poster_path.is_file(), 'Poster path does not exist'
@ -108,7 +108,7 @@ def test_validate_video(test_directory, video_factory, image_bytes_factory, vide
'view_count': 406
}
vid1.replace_info_json(info_json)
validate_video(vid1, False)
validate_video(test_session, vid1, False)
assert vid1.file_group.published_datetime and vid1.file_group.published_datetime.year == 2024
assert vid1.file_group.title == 'The full title'
@ -116,7 +116,7 @@ def test_validate_video(test_directory, video_factory, image_bytes_factory, vide
vid2 = video_factory(with_video_file=True, with_poster_ext='.png')
assert vid2.poster_path and vid2.poster_path.suffix == '.png', 'Poster was not initialized'
assert {i.suffix for i in vid2.file_group.my_paths()} == {'.png', '.mp4'}
validate_video(vid2, True)
validate_video(test_session, vid2, True)
assert vid2.poster_path.is_file(), 'New poster was not generated'
assert vid2.poster_path.suffix == '.jpg' and vid2.poster_path.stat().st_size > 0
assert {i.suffix for i in vid2.file_group.my_paths()} == {'.jpg', '.png', '.mp4'}
@ -194,7 +194,7 @@ def test_link_channel_and_downloads(test_session, channel_factory, test_download
assert not any(i.collection_id for i in test_download_manager.get_downloads(test_session))
# `link_channel_and_downloads` links Downloads to Collections.
lib.link_channel_and_downloads(session=test_session)
lib.link_channel_and_downloads(test_session)
assert test_session.query(Download).count() == 3
assert all(i.collection_id for i in test_download_manager.get_recurring_downloads(test_session))
assert not any(i.collection_id for i in test_download_manager.get_once_downloads(test_session))

6
modules/videos/test/test_video.py
View file

@ -34,7 +34,7 @@ async def test_delete_video_with_tag(async_client, test_session, video_factory,
"""You cannot delete a video with a tag."""
video = video_factory(with_video_file=True, with_poster_ext='png')
tag = await tag_factory()
video.add_tag(tag.name)
video.add_tag(test_session, tag.name)
test_session.commit()
with pytest.raises(FileGroupIsTagged):
@ -66,8 +66,8 @@ async def test_video_channel_refresh(async_client, test_session, test_directory,
assert test_session.query(Video).count() == 2
# The video is associated with `simple_channel` because its file is in the Channel's directory.
video1: Video = Video.get_by_path(video1_path, test_session)
video2: Video = Video.get_by_path(video2_path, test_session)
video1: Video = Video.get_by_path(test_session, video1_path)
video2: Video = Video.get_by_path(test_session, video2_path)
assert video1.channel == channel
assert not video2.channel
assert video1.__json__()['video']['channel']

5
modules/videos/video/api.py
View file

@ -78,13 +78,14 @@ async def search_videos(_: Request, body: schema.VideoSearchRequest):
@openapi.response(HTTPStatus.NO_CONTENT)
@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
@wrol_mode_check
async def video_delete(_: Request, video_ids: str):
async def video_delete(request: Request, video_ids: str):
try:
video_ids = [int(i) for i in str(video_ids).split(',')]
except Exception:
raise ValidationError('Unable to parse video ids')
lib.delete_videos(*video_ids)
session = request.ctx.session
lib.delete_videos(session, *video_ids)
save_channels_config.activate_switch()
Events.send_deleted(f'Deleted {len(video_ids)} videos')

9
modules/videos/video/lib.py
View file

@ -10,7 +10,7 @@ from sqlalchemy.orm import Session
from modules.videos.models import Video, Channel
from wrolpi.common import logger, limit_concurrent, wrol_mode_check
from wrolpi.dates import now
from wrolpi.db import get_db_session, optional_session
from wrolpi.db import get_db_session
from wrolpi.downloader import download_manager
from wrolpi.files.lib import handle_file_group_search_results
from wrolpi.files.models import FileGroup
@ -26,7 +26,7 @@ def get_video_for_app(video_id: int) -> Tuple[dict, Optional[dict], Optional[dic
Get a Video, with it's prev/next videos. Mark the Video as viewed.
"""
with get_db_session(commit=True) as session:
video = Video.find_by_id(video_id, session=session)
video = Video.find_by_id(session, video_id)
video.file_group.set_viewed()
previous_video, next_video = video.get_surrounding_videos()
@ -42,7 +42,7 @@ def get_video(video_id: int) -> Video:
Get a Video, with it's prev/next videos. Mark the Video as viewed.
"""
with get_db_session() as session:
video = Video.find_by_id(video_id, session=session)
video = Video.find_by_id(session, video_id)
return video
@ -159,8 +159,7 @@ def search_videos(
return results, total
@optional_session
def delete_videos(*video_ids: int, session: Session = None):
def delete_videos(session: Session, *video_ids: int):
videos = list(session.query(Video).filter(Video.id.in_(video_ids)))
if not videos:
raise UnknownVideo('Could not find videos to delete')

67
modules/zim/api.py
View file

@ -8,7 +8,6 @@ from sanic_ext.extensions.openapi import openapi
from wrolpi import lang
from wrolpi.api_utils import json_response
from wrolpi.common import logger
from wrolpi.db import get_db_session
from wrolpi.downloader import download_manager
from wrolpi.events import Events
from . import lib, schema
@ -24,12 +23,12 @@ logger = logger.getChild(__name__)
summary='List all known Zim files',
response=schema.GetZimsResponse,
)
async def get_zims(_: Request):
with get_db_session() as session:
zims = lib.get_zims(session=session)
resp = {
'zims': [i.__json__() for i in zims],
}
async def get_zims(request: Request):
session = request.ctx.session
zims = lib.get_zims(session)
resp = {
'zims': [i.__json__() for i in zims],
}
return json_response(resp)
@ -37,9 +36,11 @@ async def get_zims(_: Request):
@openapi.definition(
summary='Delete all Zim files',
)
async def delete_zims(_: Request, zim_ids: str):
async def delete_zims(request: Request, zim_ids: str):
session = request.ctx.session
zim_ids = [int(i) for i in zim_ids.split(',')]
lib.delete_zims(zim_ids)
lib.delete_zims(session, zim_ids)
return response.empty()
@ -49,9 +50,10 @@ async def delete_zims(_: Request, zim_ids: str):
body=schema.ZimAutoSearchRequest,
)
@validate(schema.ZimAutoSearchRequest)
async def post_set_zim_auto_search(_: Request, zim_id: int, body: schema.ZimAutoSearchRequest):
async def post_set_zim_auto_search(request: Request, zim_id: int, body: schema.ZimAutoSearchRequest):
session = request.ctx.session
zim_id = int(zim_id)
lib.set_zim_auto_search(zim_id, body.auto_search)
lib.set_zim_auto_search(session, zim_id, body.auto_search)
return response.empty()
@ -62,8 +64,9 @@ async def post_set_zim_auto_search(_: Request, zim_id: int, body: schema.ZimAuto
response=schema.ZimSearchResponse,
)
@validate(schema.ZimSearchRequest)
async def search_zim(_: Request, zim_id: int, body: schema.ZimSearchRequest):
headlines = lib.headline_zim(body.search_str, zim_id, tag_names=body.tag_names, offset=body.offset,
async def search_zim(request: Request, zim_id: int, body: schema.ZimSearchRequest):
session = request.ctx.session
headlines = lib.headline_zim(session, body.search_str, zim_id, tag_names=body.tag_names, offset=body.offset,
limit=body.limit)
return json_response({'zim': headlines})
@ -92,8 +95,10 @@ async def get_zim_entry(_: Request, zim_id: int, zim_path: str):
body=schema.TagZimEntry,
)
@validate(schema.TagZimEntry)
async def post_zim_tag(_: Request, body: schema.TagZimEntry):
await lib.add_tag(body.tag_name, body.zim_id, body.zim_entry)
async def post_zim_tag(request: Request, body: schema.TagZimEntry):
session = request.ctx.session
await lib.add_tag(session, body.tag_name, body.zim_id, body.zim_entry)
return response.empty(HTTPStatus.CREATED)
@ -103,8 +108,10 @@ async def post_zim_tag(_: Request, body: schema.TagZimEntry):
body=schema.TagZimEntry,
)
@validate(schema.TagZimEntry)
async def post_zim_untag(_: Request, body: schema.TagZimEntry):
await lib.untag(body.tag_name, body.zim_id, body.zim_entry)
async def post_zim_untag(request: Request, body: schema.TagZimEntry):
session = request.ctx.session
await lib.untag(session, body.tag_name, body.zim_id, body.zim_entry)
return response.empty(HTTPStatus.NO_CONTENT)
@ -113,10 +120,10 @@ async def post_zim_untag(_: Request, body: schema.TagZimEntry):
summary='Retrieve Zim subscriptions',
response=schema.ZimSubscriptions,
)
async def get_zim_subscriptions(_: Request):
async def get_zim_subscriptions(request: Request):
catalog = lib.get_kiwix_catalog()
with get_db_session() as session:
subscriptions = {i: j.__json__() for i, j in lib.get_kiwix_subscriptions(session).items()}
session = request.ctx.session
subscriptions = {i: j.__json__() for i, j in lib.get_kiwix_subscriptions(session).items()}
resp = dict(
catalog=catalog,
iso_639_codes=lang.ISO_639_CODES,
@ -131,8 +138,10 @@ async def get_zim_subscriptions(_: Request):
body=schema.ZimSubscribeRequest,
)
@validate(schema.ZimSubscribeRequest)
async def post_zim_subscribe(_: Request, body: schema.ZimSubscribeRequest):
await lib.subscribe(body.name, body.language)
async def post_zim_subscribe(request: Request, body: schema.ZimSubscribeRequest):
session = request.ctx.session
await lib.subscribe(session, body.name, body.language)
if download_manager.disabled or download_manager.stopped:
# Downloads are disabled, warn the user.
Events.send_downloads_disabled('Download created. But, downloads are disabled.')
@ -144,8 +153,10 @@ async def post_zim_subscribe(_: Request, body: schema.ZimSubscribeRequest):
@openapi.definition(
summary='Unsubscribe to a particular Kiwix Zim',
)
async def delete_zim_subscription(_: Request, subscription_id: int):
await lib.unsubscribe(subscription_id)
async def delete_zim_subscription(request: Request, subscription_id: int):
session = request.ctx.session
await lib.unsubscribe(session, subscription_id)
Events.send_deleted(f'Zim subscription deleted')
return response.empty(HTTPStatus.NO_CONTENT)
@ -175,16 +186,18 @@ async def delete_outdated_zims(_: Request):
@zim_bp.post('/search_estimates')
@validate(json=schema.SearchEstimateRequest)
async def post_search_estimates(_: Request, body: schema.SearchEstimateRequest):
async def post_search_estimates(request: Request, body: schema.SearchEstimateRequest):
"""Get an estimated count of FileGroups/Zims which may or may not have been tagged."""
if not body.search_str and not body.tag_names:
return response.empty(HTTPStatus.BAD_REQUEST)
session = request.ctx.session
if body.tag_names:
# Get actual count of entries tagged with the tag names.
zims_estimates = list()
for zim, count in Zims.entries_with_tags(body.tag_names).items():
for zim, count in Zims.entries_with_tags(session, body.tag_names).items():
d = dict(
estimate=count,
**zim.__json__(),
@ -193,7 +206,7 @@ async def post_search_estimates(_: Request, body: schema.SearchEstimateRequest):
else:
# Get estimates using libzim.
zims_estimates = list()
for zim, estimate in Zims.estimate(body.search_str).items():
for zim, estimate in Zims.estimate(session, body.search_str).items():
d = dict(
estimate=estimate,
**zim.__json__(),

54
modules/zim/lib.py
View file

@ -18,7 +18,7 @@ from wrolpi import flags
from wrolpi.cmd import run_command
from wrolpi.common import register_modeler, logger, extract_html_text, extract_headlines, get_media_directory, walk, \
register_refresh_cleanup, background_task, get_wrolpi_config, unique_by_predicate
from wrolpi.db import get_db_session, optional_session, get_db_curs
from wrolpi.db import get_db_session, get_db_curs
from wrolpi.downloader import DownloadFrequency
from wrolpi.files.lib import refresh_files, split_file_name_words
from wrolpi.files.models import FileGroup
@ -148,15 +148,14 @@ def get_unique_paths(*paths: str) -> List[str]:
return list(unique_paths)
@optional_session
def headline_zim(search_str: str, zim_id: int, tag_names: List[str] = None, offset: int = 0, limit: int = 10,
session: Session = None) -> Dict:
def headline_zim(session: Session, search_str: str, zim_id: int, tag_names: List[str] = None, offset: int = 0,
limit: int = 10) -> Dict:
zim_id = int(zim_id)
zim = get_zim(zim_id, session=session)
zim = get_zim(session, zim_id)
if tag_names:
# TODO what is the count goes over the limit?
entries = zim.entries_with_tags(tag_names, offset=offset, limit=1_000, session=session)
entries = zim.entries_with_tags(session, tag_names, offset=offset, limit=1_000)
estimate = 0
elif search_str:
# Get suggested results (this searches titles).
@ -232,8 +231,7 @@ def get_estimates(search_str: str) -> List[int]:
return estimates
@optional_session
def get_zim(zim_id: int, session: Session = None) -> Zim:
def get_zim(session: Session, zim_id: int) -> Zim:
"""Return the Zim record of the provided path.
@raise UnknownZim: When no record exists with the id.
@ -251,40 +249,36 @@ def get_zim(zim_id: int, session: Session = None) -> Zim:
@functools.lru_cache(maxsize=1_000)
def get_entry(path: str, zim_id: int) -> Entry:
with get_db_session() as session:
zim = get_zim(zim_id, session=session)
zim = get_zim(session, zim_id)
entry = zim.find_entry(path)
return entry
@optional_session
def get_zims(ids: List[int] = None, session: Session = None) -> List[Zim]:
def get_zims(session: Session, ids: List[int] = None) -> List[Zim]:
if ids:
zims = session.query(Zim).filter(Zim.id.in_(ids)).all()
else:
zims = Zims.get_all(session=session)
zims = Zims.get_all(session)
return zims
@optional_session(commit=True)
def delete_zims(ids: List[int], session: Session = None):
zims = get_zims(ids, session=session)
def delete_zims(session: Session, ids: List[int]):
zims = get_zims(session, ids)
for zim in zims:
zim.delete()
@optional_session
async def add_tag(tag_name: str, zim_id: int, zim_entry: str, session: Session = None) -> TagZimEntry:
async def add_tag(session: Session, tag_name: str, zim_id: int, zim_entry: str) -> TagZimEntry:
zim = session.query(Zim).filter_by(id=zim_id).one()
tag_zim_entry = zim.tag_entry(tag_name, zim_entry)
tag_zim_entry = zim.tag_entry(session, tag_name, zim_entry)
session.commit()
return tag_zim_entry
@optional_session
async def untag(tag_name: str, zim_id: int, zim_entry: str, session: Session = None):
async def untag(session: Session, tag_name: str, zim_id: int, zim_entry: str):
zim: Zim = session.query(Zim).filter_by(id=zim_id).one()
zim.untag_entry(tag_name, zim_entry)
zim.untag_entry(session, tag_name, zim_entry)
session.commit()
@ -292,17 +286,15 @@ def get_kiwix_catalog():
return KIWIX_CATALOG.copy()
@optional_session
def get_kiwix_subscriptions(session: Session = None) -> Dict[str, ZimSubscription]:
def get_kiwix_subscriptions(session: Session) -> Dict[str, ZimSubscription]:
subscriptions = session.query(ZimSubscription)
subscriptions_by_name = {i.name: i for i in subscriptions}
return subscriptions_by_name
@optional_session(commit=True)
async def subscribe(name: str, language: str, session: Session = None,
async def subscribe(session: Session, name: str, language: str,
frequency: int = DownloadFrequency.days180) -> ZimSubscription:
subscription = ZimSubscription.get_or_create(name, session=session)
subscription = ZimSubscription.get_or_create(session, name)
subscription.language = language
if name == 'Wikipedia (with images)':
@ -378,15 +370,14 @@ async def subscribe(name: str, language: str, session: Session = None,
if language not in kiwix.KIWIX_CATALOG_BY_NAME[name]['languages']:
raise ValueError(f'Language {repr(str(language))} is invalid for {name}')
subscription.change_download(url, frequency, session=session)
subscription.change_download(session, url, frequency)
session.add(subscription)
session.flush([subscription])
session.commit()
return subscription
@optional_session
async def unsubscribe(subscription_id: int, session: Session = None):
async def unsubscribe(session: Session, subscription_id: int):
subscription: ZimSubscription = session.query(ZimSubscription).filter_by(id=subscription_id).one_or_none()
if subscription:
session.delete(subscription)
@ -596,8 +587,7 @@ async def restart_kiwix():
return result.return_code
@optional_session
def set_zim_auto_search(zim_id: int, auto_search: bool, session: Session = None):
zim = Zim.find_by_id(zim_id, session)
def set_zim_auto_search(session: Session, zim_id: int, auto_search: bool):
zim = Zim.find_by_id(session, zim_id)
zim.auto_search = auto_search
session.commit()

83
modules/zim/models.py
View file

@ -14,7 +14,7 @@ from modules.zim.errors import UnknownZimEntry, UnknownZimTagEntry, UnknownZim
from wrolpi import dates, tags
from wrolpi.common import Base, logger, get_relative_to_media_directory, ModelHelper
from wrolpi.dates import TZDateTime
from wrolpi.db import optional_session, get_db_curs
from wrolpi.db import get_db_curs
from wrolpi.downloader import Download, download_manager
from wrolpi.files.models import FileGroup
from wrolpi.media_path import MediaPathType
@ -100,31 +100,27 @@ class Zim(Base, ModelHelper):
return zim
@staticmethod
@optional_session
def get_by_path(path: pathlib.Path, session: Session = None) -> Optional['Zim']:
def get_by_path(session: Session, path: pathlib.Path) -> Optional['Zim']:
zim = session.query(Zim).filter_by(path=path).one_or_none()
return zim
@staticmethod
@optional_session
def find_by_path(path: pathlib.Path, session: Session = None) -> 'Zim':
zim = Zim.get_by_path(path, session=session)
def find_by_path(session: Session, path: pathlib.Path) -> 'Zim':
zim = Zim.get_by_path(session, path)
if not zim:
raise UnknownZim(f'Cannot find zim at path')
return zim
@staticmethod
@optional_session
def get_by_id(zim_id: int, session: Session = None) -> Optional['Zim']:
zim = session.query(Zim).filter_by(id=zim_id).one_or_none()
def get_by_id(session: Session, id_: int) -> Optional['Zim']:
zim = session.query(Zim).filter_by(id=id_).one_or_none()
return zim
@staticmethod
@optional_session
def find_by_id(zim_id: int, session: Session = None) -> 'Zim':
zim = Zim.get_by_id(zim_id, session=session)
def find_by_id(session: Session, id_: int) -> 'Zim':
zim = Zim.get_by_id(session, id_)
if not zim:
raise UnknownZim(f'Cannot find zim with ID {zim_id}')
raise UnknownZim(f'Cannot find zim with ID {id_}')
return zim
def get_zim(self) -> Archive:
@ -149,7 +145,7 @@ class Zim(Base, ModelHelper):
return False
@staticmethod
def do_model(file_group: FileGroup, session: Session) -> 'Zim':
def do_model(session: Session, file_group: FileGroup) -> 'Zim':
from modules.zim.lib import model_zim
zim = model_zim(file_group, session)
return zim
@ -226,26 +222,23 @@ class Zim(Base, ModelHelper):
raise UnknownZimEntry(f'Cannot find entry at {path=}')
def tag_entry(self, tag_id_or_name: int | str, zim_entry: str) -> 'TagZimEntry':
def tag_entry(self, session: Session, tag_id_or_name: int | str, zim_entry: str) -> 'TagZimEntry':
"""Create a TagZimEntry for this Zim at the provided entry (path) for the provided Tag."""
session = Session.object_session(self)
# Ensure the entry exists.
self.find_entry(zim_entry)
tag = Tag.get_by_name(tag_id_or_name, session) if isinstance(tag_id_or_name, str) \
tag = Tag.get_by_name(session, tag_id_or_name) if isinstance(tag_id_or_name, str) \
else Tag.get_by_id(tag_id_or_name, session)
tag_zim_entry = TagZimEntry(tag=tag, zim=self, zim_entry=zim_entry)
session.add(tag_zim_entry)
tags.save_tags_config.activate_switch()
return tag_zim_entry
def untag_entry(self, tag_id_or_name: int | str, zim_entry: str):
def untag_entry(self, session: Session, tag_id_or_name: int | str, zim_entry: str):
"""Removes any TagZimEntry of the Zim entry at the provided path for the provided Tag.
@raise UnknownZimTagEntry: No entry exists to remove."""
session = Session.object_session(self)
tag = Tag.get_by_name(tag_id_or_name, session) if isinstance(tag_id_or_name, str) \
tag = Tag.get_by_name(session, tag_id_or_name) if isinstance(tag_id_or_name, str) \
else Tag.get_by_id(tag_id_or_name, session)
tag_zim_entry = TagZimEntry.get_by_primary_keys(tag.id, self.id, zim_entry, session)
if tag_zim_entry:
@ -274,23 +267,23 @@ class Zim(Base, ModelHelper):
return entries
@optional_session
def entries_with_tags(*args, session: Session = None, **kwargs) -> List[Entry]:
def entries_with_tags(self, session: Session, tag_names: List[str], limit: int = 10, offset: int = 0) -> List[
Entry]:
"""Return Zim Entries tagged with the provided names."""
limit = int(kwargs.pop('limit', 10))
offset = int(kwargs.pop('offset', 0))
args = list(args)
zim: Zim | None = args.pop(0) if isinstance(args[0], Zim) else None
zim_id = zim.id if zim else None
tag_names: List[str] = args.pop(0)
if limit < 1:
raise RuntimeError('Limit must be greater than zero')
if args:
raise RuntimeError('Extra unknown arguments')
tags_sub_select, params = tag_names_to_zim_sub_select(tag_names, zim_id=zim_id)
tags_sub_select, params = tag_names_to_zim_sub_select(tag_names, zim_id=self.id)
entries = Zim._entries_with_tags_process(limit, offset, params, session, tags_sub_select)
return entries
@staticmethod
def all_entries_with_tags(session: Session, tag_names: List[str], limit: int = 10, offset: int = 0) -> List[Entry]:
"""Return Zim Entries tagged with the provided names from ALL Zims."""
if limit < 1:
raise RuntimeError('Limit must be greater than zero')
tags_sub_select, params = tag_names_to_zim_sub_select(tag_names, zim_id=None)
entries = Zim._entries_with_tags_process(limit, offset, params, session, tags_sub_select)
return entries
@ -321,8 +314,7 @@ class Zims:
"""Convenience class which searches all known Zim files."""
@staticmethod
@optional_session
def get_all(session: Session = None) -> List[Zim]:
def get_all(session: Session) -> List[Zim]:
"""Returns all Zim records."""
zims = list()
for zim in session.query(Zim).order_by(Zim.path):
@ -335,9 +327,8 @@ class Zims:
return zims
@classmethod
@optional_session
def estimate(cls, search_str: str, session: Session = None) -> OrderedDictType[Zim, int]:
zims = cls.get_all(session=session)
def estimate(cls, session: Session, search_str: str) -> OrderedDictType[Zim, int]:
zims = cls.get_all(session)
results = OrderedDict()
for zim in zims:
if zim.auto_search:
@ -345,12 +336,11 @@ class Zims:
return results
@classmethod
@optional_session
def entries_with_tags(cls, tag_names: List[str], session: Session = None) -> OrderedDictType[Zim, int]:
zims = cls.get_all(session=session)
def entries_with_tags(cls, session: Session, tag_names: List[str]) -> OrderedDictType[Zim, int]:
zims = cls.get_all(session)
results = OrderedDict()
for zim in zims:
results[zim] = len(zim.entries_with_tags(tag_names, session=session))
results[zim] = len(zim.entries_with_tags(session, tag_names))
return results
@ -408,10 +398,10 @@ class ZimSubscription(Base):
)
return d
def change_download(self, url: str, frequency: int, session: Session):
def change_download(self, session: Session, url: str, frequency: int):
# A Zim file is found using the KiwixCatalogDownloader, then handled by the KiwixZimDownloader.
old_url = self.download.url if self.download else url
download = download_manager.get_or_create_download(old_url, session=session, reset_attempts=True)
download = download_manager.get_or_create_download(session, old_url, reset_attempts=True)
download.url = url
download.frequency = frequency
download.downloader = 'kiwix_catalog'
@ -423,8 +413,7 @@ class ZimSubscription(Base):
return download
@staticmethod
@optional_session
def get_or_create(name: str, session: Session = None) -> 'ZimSubscription':
def get_or_create(session: Session, name: str) -> 'ZimSubscription':
subscription = session.query(ZimSubscription).filter_by(name=name).one_or_none()
if subscription:
return subscription

18
modules/zim/test/test_api.py
View file

@ -12,9 +12,9 @@ from wrolpi.downloader import Download
async def test_zim_search(async_client, test_session, test_zim, tag_factory, assert_zim_search):
"""Entries can be tagged and retrieved using their containing Zim."""
tag1, tag2 = await tag_factory(), await tag_factory()
test_zim.tag_entry(tag1.name, 'one')
test_zim.tag_entry(tag2.name, 'one')
test_zim.tag_entry(tag1.name, 'two')
test_zim.tag_entry(test_session, tag1.name, 'one')
test_zim.tag_entry(test_session, tag2.name, 'one')
test_zim.tag_entry(test_session, tag1.name, 'two')
test_session.commit()
# Test searching with only `search_str`.
@ -83,18 +83,18 @@ async def test_zim_search(async_client, test_session, test_zim, tag_factory, ass
async def test_entries_by_tag(async_client, test_session, test_zim, tag_factory):
"""Entries can be tagged and retrieved using their containing Zim."""
tag1, tag2 = await tag_factory(), await tag_factory()
test_zim.tag_entry(tag1.name, 'one')
test_zim.tag_entry(tag2.name, 'one')
test_zim.tag_entry(tag1.name, 'two')
test_zim.tag_entry(test_session, tag1.name, 'one')
test_zim.tag_entry(test_session, tag2.name, 'one')
test_zim.tag_entry(test_session, tag1.name, 'two')
test_session.commit()
result = test_zim.entries_with_tags([tag1.name, tag2.name])
result = test_zim.entries_with_tags(test_session, [tag1.name, tag2.name])
assert [(i.path, i.title) for i in result] == [('one', 'One'), ]
result = test_zim.entries_with_tags([tag1.name, ])
result = test_zim.entries_with_tags(test_session, [tag1.name, ])
assert [(i.path, i.title) for i in result] == [('one', 'One'), ('two', 'Two')]
result = test_zim.entries_with_tags([tag2.name, ])
result = test_zim.entries_with_tags(test_session, [tag2.name, ])
assert [(i.path, i.title) for i in result] == [('one', 'One'), ]

50
modules/zim/test/test_lib.py
View file

@ -20,10 +20,10 @@ async def test_get_zim(async_client, test_session, zim_path_factory):
zim_path_factory()
await files_lib.refresh_files()
assert lib.get_zim(1)
assert lib.get_zim(test_session, 1)
with pytest.raises(UnknownZim):
assert lib.get_zim(2)
assert lib.get_zim(test_session, 2)
@pytest.mark.asyncio
@ -54,15 +54,15 @@ async def test_zim_get_entry(async_client, test_session, zim_path_factory):
async def test_zim_bad_entry(async_client, test_session, test_zim, tag_factory):
tag1 = await tag_factory()
with pytest.raises(UnknownZimEntry):
test_zim.tag_entry(tag1.name, 'path does not exist')
test_zim.tag_entry(test_session, tag1.name, 'path does not exist')
@pytest.mark.asyncio
async def test_zim_get_entries_tags(async_client, test_session, test_zim, tag_factory):
tag1, tag2 = await tag_factory('tag1'), await tag_factory('tag2')
test_zim.tag_entry(tag1.name, 'one')
test_zim.tag_entry(tag2.name, 'one')
test_zim.tag_entry(tag1.name, 'two')
test_zim.tag_entry(test_session, tag1.name, 'one')
test_zim.tag_entry(test_session, tag2.name, 'one')
test_zim.tag_entry(test_session, tag1.name, 'two')
test_session.commit()
assert lib.get_entries_tags(['one', 'two', 'home'], 1) \
@ -72,7 +72,7 @@ async def test_zim_get_entries_tags(async_client, test_session, test_zim, tag_fa
@pytest.mark.asyncio
async def test_tag_zim_entry_model(async_client, test_session, zim_factory, tag_factory):
tag = await tag_factory()
tze = zim_factory('zim1').tag_entry(tag.name, 'one')
tze = zim_factory('zim1').tag_entry(test_session, tag.name, 'one')
# Test TagZimEntry.__repr__
assert repr(tze) == "<TagZimEntry tag='one' zim=zim1 zim_entry=one>"
@ -83,21 +83,21 @@ async def test_zim_entries_with_tags(async_client, test_session, zim_factory, ta
tag1, tag2 = await tag_factory(), await tag_factory()
zim1, zim2 = zim_factory('zim1'), zim_factory('zim2')
# Tag Zim1 with both tags, but Zim2 shares one of the tags and is returned when no Zim is passed.
zim1.tag_entry(tag1.name, 'one')
zim1.tag_entry(tag2.name, 'two')
zim2.tag_entry(tag2.name, 'home')
zim1.tag_entry(test_session, tag1.name, 'one')
zim1.tag_entry(test_session, tag2.name, 'two')
zim2.tag_entry(test_session, tag2.name, 'home')
test_session.commit()
# Get entries from zim1 only.
entries = zim1.entries_with_tags([tag1.name, ])
entries = zim1.entries_with_tags(test_session, [tag1.name, ])
assert len(entries) == 1
assert entries[0].path == 'one'
entries = zim1.entries_with_tags([tag2.name, ])
entries = zim1.entries_with_tags(test_session, [tag2.name, ])
assert len(entries) == 1
assert entries[0].path == 'two'
# Get entries from any Zim.
entries = lib.Zim.entries_with_tags([tag2.name, ])
entries = lib.Zim.all_entries_with_tags(test_session, [tag2.name, ])
assert {i.path for i in entries} == {'two', 'home'}, 'Tag2 was applied to `two` and `home`'
@ -135,9 +135,9 @@ async def test_zim_tags_config(async_client, test_session, test_directory, test_
# Tag three Zim entries.
tag1, tag2 = await tag_factory('Tag1'), await tag_factory('Tag2')
test_zim.tag_entry(tag1.name, 'one')
test_zim.tag_entry(tag2.name, 'one')
test_zim.tag_entry(tag1.name, 'two')
test_zim.tag_entry(test_session, tag1.name, 'one')
test_zim.tag_entry(test_session, tag2.name, 'one')
test_zim.tag_entry(test_session, tag1.name, 'two')
test_session.commit()
await await_switches()
# Tags are created in empty database.
@ -151,7 +151,7 @@ async def test_zim_tags_config(async_client, test_session, test_directory, test_
]
# Remove a tag, the config should change.
test_zim.untag_entry(tag1.name, 'one')
test_zim.untag_entry(test_session, tag1.name, 'one')
test_session.commit()
await await_switches()
@ -186,7 +186,7 @@ def test_zim_all_entries(test_session, test_zim):
@pytest.mark.asyncio
async def test_zim_download(test_session, kiwix_download_zim, test_directory, test_zim_bytes, flags_lock,
await_switches):
await lib.subscribe('Wikipedia (with images)', 'es')
await lib.subscribe(test_session, 'Wikipedia (with images)', 'es')
# Downloading the catalog should lead to a new Zim file being downloaded.
await kiwix_download_zim(expected_url='https://download.kiwix.org/zim/wikipedia/wikipedia_es_all_maxi_2023-06.zim')
@ -264,12 +264,12 @@ async def test_zim_tag_migration(await_switches, test_session, test_directory, z
zim3 = zim_factory('wikipedia_en_all_maxi_2020-03.zim') # The latest Zim.
tag1, tag2 = await tag_factory('tag1'), await tag_factory('tag2')
# Tag outdated Zim entries.
zim1.tag_entry(tag1.name, 'home')
zim1.tag_entry(tag2.name, 'one')
zim1.tag_entry(test_session, tag1.name, 'home')
zim1.tag_entry(test_session, tag2.name, 'one')
# This should not be migrated.
zim2.tag_entry(tag2.name, 'two')
zim2.tag_entry(test_session, tag2.name, 'two')
# Tag the latest Zim as well.
zim3.tag_entry(tag2.name, 'one')
zim3.tag_entry(test_session, tag2.name, 'one')
test_session.commit()
assert test_session.query(TagZimEntry).count() == 4
await await_switches()
@ -358,7 +358,7 @@ def test_zim_download_url_to_name(url, expected):
@pytest.mark.asyncio
async def test_zim_subscription_download_import(async_client, test_session):
# Subscription creates a ZimSubscription and Download
await lib.subscribe('Wikisource', 'en', session=test_session)
await lib.subscribe(test_session, 'Wikisource', 'en')
# Add a once-download. This should not be associated with a ZimSubscription.
test_session.add(
Download(url='https://download.kiwix.org/zim/wikibooks/wikibooks_en_all_maxi_2021-03.zim', status='complete'))
@ -376,7 +376,7 @@ async def test_zim_subscription_download_import(async_client, test_session):
for _ in range(1): # Import twice.
# Importing the config restores the Download and ZimSubscription.
await import_downloads_config(session=test_session)
await import_downloads_config(test_session)
recurring, once = test_session.query(Download).order_by(Download.frequency)
assert once.url == 'https://download.kiwix.org/zim/wikibooks/wikibooks_en_all_maxi_2021-03.zim'
assert test_session.query(Download).count() == 2
@ -404,7 +404,7 @@ async def test_zim_modeler(async_client, test_session, zim_factory):
assert fg.model == 'zim'
def test_get_custom_zims_directory(test_directory, test_wrolpi_config):
def test_get_custom_zims_directory(async_client, test_directory, test_wrolpi_config):
"""Custom directory can be used for zims directory."""
assert lib.get_zim_directory() == (test_directory / 'zims')

30
wrolpi/api_utils.py
View file

@ -119,6 +119,36 @@ def json_error_handler(request: Request, exception: Exception):
api_app.error_handler.add(Exception, json_error_handler)
@api_app.middleware('request')
async def inject_session(request: Request):
"""Inject a database session into the request context."""
from wrolpi.db import get_db_context
engine, session = get_db_context()
request.ctx.session = session
request.ctx._db_engine = engine
@api_app.middleware('response')
async def cleanup_session(request: Request, response_: HTTPResponse):
"""Cleanup session after request completes."""
if hasattr(request.ctx, 'session'):
session = request.ctx.session
try:
if response_.status < 400:
session.commit()
else:
session.rollback()
except Exception:
session.rollback()
raise
finally:
# Don't close session during tests - the test_session fixture manages it.
if not PYTEST:
session.close()
return response_
PERPETUAL_WORKERS = list()

31
wrolpi/collections/api.py
View file

@ -41,8 +41,9 @@ async def get_collections_endpoint(request: Request):
- metadata: UI metadata (columns, fields, routes, messages) if kind is specified
"""
kind = request.args.get('kind')
session = request.ctx.session
collections = lib.get_collections(kind=kind)
collections = lib.get_collections(session, kind=kind)
response_data = {
'collections': collections,
@ -56,7 +57,7 @@ async def get_collections_endpoint(request: Request):
@openapi.summary('Get a single collection by ID')
@openapi.response(HTTPStatus.OK, description="Collection details with statistics")
@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
async def get_collection_endpoint(_: Request, collection_id: int):
async def get_collection_endpoint(request: Request, collection_id: int):
"""
Get details for a single collection including type-specific statistics.
@ -65,8 +66,9 @@ async def get_collection_endpoint(_: Request, collection_id: int):
- Channel collections: includes video_count and size
- Other types: includes item_count and total_size
"""
session = request.ctx.session
try:
collection_data = lib.get_collection_with_stats(collection_id)
collection_data = lib.get_collection_with_stats(session, collection_id)
return json_response({'collection': collection_data})
except UnknownCollection as e:
return json_response({'error': str(e)}, status=HTTPStatus.NOT_FOUND)
@ -82,7 +84,7 @@ async def get_collection_endpoint(_: Request, collection_id: int):
@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
@openapi.response(HTTPStatus.BAD_REQUEST, JSONErrorResponse)
@wrol_mode_check
async def put_collection_endpoint(_: Request, collection_id: int, body: schema.CollectionUpdateRequest):
async def put_collection_endpoint(request: Request, collection_id: int, body: schema.CollectionUpdateRequest):
"""
Update collection properties (directory, tag, description).
@ -93,8 +95,10 @@ async def put_collection_endpoint(_: Request, collection_id: int, body: schema.C
Note: To clear a field, pass an empty string. To leave unchanged, omit the field.
"""
session = request.ctx.session
try:
collection = lib.update_collection(
session,
collection_id=collection_id,
directory=body.directory,
tag_name=body.tag_name,
@ -102,7 +106,7 @@ async def put_collection_endpoint(_: Request, collection_id: int, body: schema.C
)
# Return updated collection data
collection_data = lib.get_collection_with_stats(collection_id)
collection_data = lib.get_collection_with_stats(session, collection_id)
return json_response({'collection': collection_data})
except UnknownCollection as e:
@ -146,7 +150,7 @@ async def refresh_collection_endpoint(_: Request, collection_id: int):
@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
@openapi.response(HTTPStatus.BAD_REQUEST, JSONErrorResponse)
@wrol_mode_check
async def tag_collection_endpoint(_: Request, collection_id: int, body: schema.CollectionTagRequest):
async def tag_collection_endpoint(request: Request, collection_id: int, body: schema.CollectionTagRequest):
"""
Tag a collection and optionally move its files to a new directory.
@ -157,8 +161,10 @@ async def tag_collection_endpoint(_: Request, collection_id: int, body: schema.C
If no directory is specified, the collection must already have one.
"""
session = request.ctx.session
try:
result = await lib.tag_collection(
session,
collection_id=collection_id,
tag_name=body.tag_name,
directory=body.directory
@ -179,7 +185,7 @@ async def tag_collection_endpoint(_: Request, collection_id: int, body: schema.C
)
@openapi.response(HTTPStatus.OK, schema.CollectionTagInfoResponse, description="Tag info retrieved successfully")
@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
async def get_tag_info_endpoint(_: Request, collection_id: int, body: schema.CollectionTagInfoRequest):
async def get_tag_info_endpoint(request: Request, collection_id: int, body: schema.CollectionTagInfoRequest):
"""
Get information about tagging a collection with a specific tag.
@ -189,8 +195,10 @@ async def get_tag_info_endpoint(_: Request, collection_id: int, body: schema.Col
This is useful for showing users the suggested directory before they commit to tagging.
"""
session = request.ctx.session
try:
tag_info = lib.get_tag_info(
session,
collection_id=collection_id,
tag_name=body.tag_name
)
@ -207,7 +215,7 @@ async def get_tag_info_endpoint(_: Request, collection_id: int, body: schema.Col
@openapi.response(HTTPStatus.NO_CONTENT, description="Collection deleted successfully")
@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
@wrol_mode_check
def delete_collection_endpoint(_: Request, collection_id: int):
async def delete_collection_endpoint(request: Request, collection_id: int):
"""
Delete a collection and orphan its child items.
@ -217,8 +225,9 @@ def delete_collection_endpoint(_: Request, collection_id: int):
- Archives remain in the database but are no longer associated with a domain
- Triggers domain config save
"""
session = request.ctx.session
try:
collection = lib.delete_collection(collection_id=collection_id)
collection = lib.delete_collection(session, collection_id=collection_id)
from wrolpi.events import Events
Events.send_deleted(f'Deleted {collection["kind"]} collection: {collection["name"]}')
return response.raw('', HTTPStatus.NO_CONTENT)
@ -234,7 +243,7 @@ def delete_collection_endpoint(_: Request, collection_id: int):
validate=True,
)
@openapi.response(HTTPStatus.OK, description="Search results")
async def search_collections_endpoint(_: Request, body: schema.CollectionSearchRequest):
async def search_collections_endpoint(request: Request, body: schema.CollectionSearchRequest):
"""
Search collections by kind, tags, and name.
@ -245,7 +254,9 @@ async def search_collections_endpoint(_: Request, body: schema.CollectionSearchR
All filters are optional and can be combined.
"""
session = request.ctx.session
collections = lib.search_collections(
session,
kind=body.kind,
tag_names=body.tag_names if body.tag_names else None,
search_str=body.search_str

42
wrolpi/collections/lib.py
View file

@ -10,8 +10,10 @@ from typing import List, Optional, Dict
from sqlalchemy import func
from sqlalchemy.orm import Session
from wrolpi import flags
from wrolpi.common import logger, get_relative_to_media_directory
from wrolpi.db import get_db_session, optional_session
from wrolpi.db import get_db_session
from wrolpi.errors import RefreshConflict
from wrolpi.errors import ValidationError
from wrolpi.events import Events
from wrolpi.tags import Tag
@ -32,8 +34,7 @@ __all__ = [
]
@optional_session
def get_collections(kind: Optional[str] = None, session: Session = None) -> List[dict]:
def get_collections(session: Session, kind: Optional[str] = None) -> List[dict]:
"""
Get all collections, optionally filtered by kind.
@ -41,8 +42,8 @@ def get_collections(kind: Optional[str] = None, session: Session = None) -> List
avoiding the N+1 query problem that occurs with per-collection queries.
Args:
kind: Optional collection kind to filter by (e.g., 'domain', 'channel')
session: Database session
kind: Optional collection kind to filter by (e.g., 'domain', 'channel')
Returns:
List of collection dicts with statistics for each collection type
@ -135,14 +136,13 @@ def get_collections(kind: Optional[str] = None, session: Session = None) -> List
return result
@optional_session
def get_collection_with_stats(collection_id: int, session: Session = None) -> dict:
def get_collection_with_stats(session: Session, collection_id: int) -> dict:
"""
Get a single collection with type-specific statistics.
Args:
collection_id: The collection ID
session: Database session
collection_id: The collection ID
Returns:
Collection dict with statistics
@ -209,23 +209,22 @@ def get_collection_with_stats(collection_id: int, session: Session = None) -> di
return data
@optional_session
def update_collection(
session: Session,
collection_id: int,
directory: Optional[str] = None,
tag_name: Optional[str] = None,
description: Optional[str] = None,
session: Session = None
) -> Collection:
"""
Update a collection's properties.
Args:
session: Database session
collection_id: The collection ID
directory: New directory (relative or absolute path), or None to clear
tag_name: New tag name, empty string to clear, or None to leave unchanged
description: New description, or None to leave unchanged
session: Database session
Returns:
Updated Collection object
@ -319,21 +318,20 @@ def refresh_collection(collection_id: int, send_events: bool = True) -> None:
Events.send_directory_refresh(f'Refreshing: {relative_dir}')
@optional_session(commit=True)
async def tag_collection(
session: Session,
collection_id: int,
tag_name: Optional[str] = None,
directory: Optional[str] = None,
session: Session = None
) -> Dict:
"""
Tag a collection and optionally move files to a new directory, or remove a tag if no tag_name is provided.
Args:
session: Database session
collection_id: The collection ID
tag_name: Tag name to apply, or None to remove the tag
directory: Optional new directory for the collection
session: Database session
Returns:
Dict with tag information and suggested directory
@ -343,9 +341,6 @@ async def tag_collection(
ValidationError: If tagging requirements not met
RefreshConflict: If a file refresh is in progress and directory change is requested
"""
from wrolpi import flags
from wrolpi.errors import RefreshConflict
collection = session.query(Collection).filter_by(id=collection_id).one_or_none()
if not collection:
@ -465,11 +460,10 @@ async def tag_collection(
}
@optional_session
def get_tag_info(
session: Session,
collection_id: int,
tag_name: Optional[str],
session: Session = None
) -> Dict:
"""
Get information about tagging a collection with a specific tag.
@ -478,9 +472,9 @@ def get_tag_info(
For collections without a directory, returns suggested_directory=None.
Args:
session: Database session
collection_id: The collection ID
tag_name: Tag name to check
session: Database session
Returns:
Dict with suggested_directory, conflict flag, and optional conflict_message
@ -534,10 +528,9 @@ def get_tag_info(
}
@optional_session
def delete_collection(
session: Session,
collection_id: int,
session: Session = None
) -> Dict:
"""
Delete a collection and orphan its child items.
@ -548,8 +541,8 @@ def delete_collection(
- Triggers domain config save
Args:
collection_id: The collection ID to delete
session: Database session
collection_id: The collection ID to delete
Returns:
Dict with collection information
@ -588,21 +581,20 @@ def delete_collection(
return collection_dict
@optional_session
def search_collections(
session: Session,
kind: Optional[str] = None,
tag_names: Optional[List[str]] = None,
search_str: Optional[str] = None,
session: Session = None
) -> List[dict]:
"""
Search collections by kind, tags, and search string.
Args:
session: Database session
kind: Optional collection kind filter
tag_names: Optional list of tag names to filter by
search_str: Optional search string for collection names
session: Database session
Returns:
List of matching collection dicts

80
wrolpi/collections/models.py
View file

@ -8,7 +8,6 @@ from sqlalchemy.orm.collections import InstrumentedList
from wrolpi import flags
from wrolpi.common import Base, ModelHelper, logger, get_media_directory, get_relative_to_media_directory, \
unique_by_predicate
from wrolpi.db import optional_session
from wrolpi.downloader import Download, save_downloads_config
from wrolpi.errors import ValidationError
from wrolpi.events import Events
@ -180,17 +179,16 @@ class Collection(ModelHelper, Base):
return config
@optional_session
def get_or_create_download(self, url: str, frequency: int, session: Session = None,
def get_or_create_download(self, session: Session, url: str, frequency: int,
reset_attempts: bool = False,
downloader_name: str = None,
sub_downloader_name: str = None) -> 'Download':
"""Get or create a Download for this Collection.
Args:
session: Database session
url: The URL to download
frequency: Seconds between re-downloading
session: Database session
reset_attempts: Reset download attempts if True
downloader_name: Name of the downloader to use
sub_downloader_name: Name of the sub-downloader for child downloads
@ -210,10 +208,10 @@ class Collection(ModelHelper, Base):
if not frequency:
raise InvalidDownload('Download for Collection must have a frequency')
download = Download.get_by_url(url, session=session)
download = Download.get_by_url(session, url)
if not download:
download = download_manager.recurring_download(
url, frequency, downloader_name, session=session,
session, url, frequency, downloader_name,
sub_downloader_name=sub_downloader_name,
destination=str(self.directory) if self.directory else None,
reset_attempts=reset_attempts,
@ -224,8 +222,7 @@ class Collection(ModelHelper, Base):
return download
@staticmethod
@optional_session
def from_config(data: dict, session: Session = None) -> 'Collection':
def from_config(session: Session, data: dict) -> 'Collection':
"""
Create or update a Collection from config data.
@ -233,8 +230,8 @@ class Collection(ModelHelper, Base):
If the collection has no directory, items are managed manually (not from config).
Args:
data: Config dict containing collection metadata
session: Database session
data: Config dict containing collection metadata
Returns:
The created or updated Collection
@ -268,7 +265,7 @@ class Collection(ModelHelper, Base):
# - If directory is provided, resolve by unique directory
# - Otherwise, fall back to first match by (name, kind)
if directory:
collection = Collection.get_by_path(directory, session)
collection = Collection.get_by_path(session, directory)
else:
collection = session.query(Collection).filter_by(name=name, kind=kind).first()
@ -280,7 +277,7 @@ class Collection(ModelHelper, Base):
collection.kind = kind
if tag_name:
tag = Tag.get_by_name(tag_name, session)
tag = Tag.get_by_name(session, tag_name)
if tag:
collection.tag = tag
else:
@ -299,7 +296,7 @@ class Collection(ModelHelper, Base):
)
if tag_name:
tag = Tag.get_by_name(tag_name, session)
tag = Tag.get_by_name(session, tag_name)
if tag:
collection.tag = tag
else:
@ -312,15 +309,14 @@ class Collection(ModelHelper, Base):
# If directory-restricted, populate with FileGroups from that directory
if collection.directory and collection.directory.is_dir():
logger.info(f'Populating collection {repr(name)} from directory {collection.directory}')
collection.populate_from_directory(session=session)
collection.populate_from_directory(session)
elif collection.directory and not collection.directory.is_dir():
logger.warning(f'Collection directory does not exist: {collection.directory}')
return collection
@staticmethod
@optional_session
def batch_from_config(data_list: List[dict], session: Session = None) -> List['Collection']:
def batch_from_config(session: Session, data_list: List[dict]) -> List['Collection']:
"""
Create or update multiple Collections from config data in batch.
More efficient than calling from_config() in a loop because it pre-fetches
@ -330,8 +326,8 @@ class Collection(ModelHelper, Base):
use claim_videos_for_channels() after batch_from_config() to assign videos to channels.
Args:
data_list: List of config dicts containing collection metadata
session: Database session
data_list: List of config dicts containing collection metadata
Returns:
List of created or updated Collections
@ -460,12 +456,11 @@ class Collection(ModelHelper, Base):
return collections
def populate_from_directory(self, session: Session = None):
def populate_from_directory(self, session: Session):
"""
Populate this collection with all FileGroups in the collection's directory.
Only works for directory-restricted collections.
"""
session = session or Session.object_session(self)
if not self.is_directory_restricted:
logger.warning(f'Cannot populate unrestricted collection from directory: {self}')
@ -490,7 +485,7 @@ class Collection(ModelHelper, Base):
# Add new FileGroups
new_file_groups = [fg for fg in file_groups if fg.id not in existing_fg_ids]
if new_file_groups:
self.add_file_groups(new_file_groups, session=session)
self.add_file_groups(session, new_file_groups)
logger.info(f'Added {len(new_file_groups)} FileGroups to collection {repr(self.name)}')
def validate_file_group(self, file_group: FileGroup) -> bool:
@ -560,7 +555,7 @@ class Collection(ModelHelper, Base):
position = (max_position or 0) + 1
else:
# Insert at specific position - need to shift existing items
self._shift_positions(position, shift_by=1, session=session)
self._shift_positions(session, position, shift_by=1)
item = CollectionItem(
collection_id=self.id,
@ -572,19 +567,18 @@ class Collection(ModelHelper, Base):
return item
def add_file_groups(self, file_groups: List[FileGroup], session: Session = None) -> List['CollectionItem']:
def add_file_groups(self, session: Session, file_groups: List[FileGroup]) -> List['CollectionItem']:
"""
Add multiple FileGroups to this Collection in batch.
More efficient than calling add_file_group in a loop.
Args:
file_groups: List of FileGroups to add
session: Database session
file_groups: List of FileGroups to add
Returns:
List of created CollectionItems
"""
session = session or Session.object_session(self)
# Validate all file groups first
for fg in file_groups:
@ -625,9 +619,8 @@ class Collection(ModelHelper, Base):
logger.debug(f'Added {len(items)} FileGroups to Collection {self.id}')
return items
def remove_file_group(self, file_group_id: int, session: Session = None):
def remove_file_group(self, session: Session, file_group_id: int):
"""Remove a FileGroup from this Collection."""
session = session or Session.object_session(self)
item = session.query(CollectionItem).filter_by(
collection_id=self.id,
@ -640,18 +633,17 @@ class Collection(ModelHelper, Base):
session.flush()
# Shift remaining items down
self._shift_positions(position + 1, shift_by=-1, session=session)
self._shift_positions(session, position + 1, shift_by=-1)
def remove_file_groups(self, file_group_ids: List[int], session: Session = None):
def remove_file_groups(self, session: Session, file_group_ids: List[int]):
"""
Remove multiple FileGroups from this Collection in batch.
More efficient than calling remove_file_group in a loop.
Args:
file_group_ids: List of FileGroup IDs to remove
session: Database session
file_group_ids: List of FileGroup IDs to remove
"""
session = session or Session.object_session(self)
if not file_group_ids:
return
@ -707,7 +699,7 @@ class Collection(ModelHelper, Base):
item.position = new_position
session.flush()
def _shift_positions(self, from_position: int, shift_by: int, session: Session):
def _shift_positions(self, session: Session, from_position: int, shift_by: int):
"""Helper to shift positions of items."""
session.query(CollectionItem).filter(
CollectionItem.collection_id == self.id,
@ -715,10 +707,8 @@ class Collection(ModelHelper, Base):
).update({'position': CollectionItem.position + shift_by})
session.flush()
def get_items(self, limit: Optional[int] = None, offset: int = 0,
session: Session = None) -> List['CollectionItem']:
def get_items(self, session: Session, limit: Optional[int] = None, offset: int = 0) -> List['CollectionItem']:
"""Get items in this collection, ordered by position."""
session = session or Session.object_session(self)
query = session.query(CollectionItem).filter_by(
collection_id=self.id
@ -732,25 +722,22 @@ class Collection(ModelHelper, Base):
return query.all()
@staticmethod
@optional_session
def get_by_id(id_: int, session: Session = None) -> Optional['Collection']:
def get_by_id(session: Session, id_: int) -> Optional['Collection']:
"""Attempt to find a Collection with the provided id. Returns None if not found."""
return session.query(Collection).filter_by(id=id_).one_or_none()
@staticmethod
@optional_session
def find_by_id(id_: int, session: Session = None) -> 'Collection':
def find_by_id(session: Session, id_: int) -> 'Collection':
"""Find a Collection with the provided id, raises an exception if not found.
@raise UnknownCollection: if the collection cannot be found
"""
if collection := Collection.get_by_id(id_, session=session):
if collection := Collection.get_by_id(session, id_):
return collection
raise UnknownCollection(f'Cannot find Collection with id {id_}')
@staticmethod
@optional_session
def get_by_path(path: pathlib.Path, session: Session = None) -> Optional['Collection']:
def get_by_path(session: Session, path: pathlib.Path) -> Optional['Collection']:
"""Find a Collection by its directory path. Returns None if not found."""
if not path:
return None
@ -759,13 +746,12 @@ class Collection(ModelHelper, Base):
return session.query(Collection).filter_by(directory=path).one_or_none()
@staticmethod
@optional_session
def find_by_path(path: pathlib.Path, session: Session = None) -> 'Collection':
def find_by_path(session: Session, path: pathlib.Path) -> 'Collection':
"""Find a Collection by its directory path, raises an exception if not found.
@raise UnknownCollection: if the collection cannot be found
"""
if collection := Collection.get_by_path(path, session=session):
if collection := Collection.get_by_path(session, path):
return collection
raise UnknownCollection(f'Cannot find Collection with directory {path}')
@ -798,9 +784,9 @@ class Collection(ModelHelper, Base):
if tag_id_or_name is None:
self.tag = None
elif isinstance(tag_id_or_name, int):
self.tag = Tag.find_by_id(tag_id_or_name, session)
self.tag = Tag.find_by_id(session, tag_id_or_name)
elif isinstance(tag_id_or_name, str):
self.tag = Tag.find_by_name(tag_id_or_name, session)
self.tag = Tag.find_by_name(session, tag_id_or_name)
self.tag_id = self.tag.id if self.tag else None
return self.tag
@ -873,7 +859,7 @@ class Collection(ModelHelper, Base):
def change_download_destinations(from_directory: pathlib.Path, to_directory: pathlib.Path):
"""Update download destinations from one directory to another."""
downloads = list(self.downloads)
downloads.extend(Download.get_all_by_destination(from_directory, session=session))
downloads.extend(Download.get_all_by_destination(session, from_directory))
downloads = unique_by_predicate(downloads, lambda i: i.id)
for download in downloads:
download.destination = to_directory
@ -898,7 +884,7 @@ class Collection(ModelHelper, Base):
if send_events:
Events.send_file_move_completed(f'Collection {repr(self.name)} was moved (directory missing)')
else:
await move_files(directory, *list(old_directory.iterdir()), session=session)
await move_files(session, directory, *list(old_directory.iterdir()))
if send_events:
Events.send_file_move_completed(f'Collection {repr(self.name)} was moved')
except Exception as e:

3
wrolpi/collections/test/test_api.py
View file

@ -343,6 +343,9 @@ class TestCollectionDeletion:
# Check response
assert response.status_code == HTTPStatus.NO_CONTENT
# Expire all cached objects to see changes made by the API's separate session
test_session.expire_all()
# Verify collection is deleted
assert test_session.query(Collection).filter_by(id=collection_id).count() == 0

14
wrolpi/collections/test/test_collection_lifecycle.py
View file

@ -49,11 +49,11 @@ async def test_collection_lifecycle_end_to_end(async_client, test_session: Sessi
'directory': str(rel_dir),
'kind': 'channel',
}
coll = Collection.from_config(cfg, session=test_session)
coll = Collection.from_config(test_session, cfg)
test_session.commit()
# After creation, items should be auto-populated from directory
items = coll.get_items(session=test_session)
items = coll.get_items(test_session)
fg_ids = {i.file_group_id for i in items}
assert {v1.file_group_id, v2.file_group_id, v3.file_group_id}.issubset(fg_ids)
@ -74,18 +74,18 @@ async def test_collection_lifecycle_end_to_end(async_client, test_session: Sessi
test_session.commit()
# Populate from directory to pick up new file group
coll.populate_from_directory(session=test_session)
coll.populate_from_directory(test_session)
test_session.commit()
items = coll.get_items(session=test_session)
items = coll.get_items(test_session)
fg_ids = {i.file_group_id for i in items}
assert v4.file_group_id in fg_ids
# 6) Remove one video from the collection only (do not delete the video)
coll.remove_file_group(v2.file_group_id, session=test_session)
coll.remove_file_group(test_session, v2.file_group_id)
test_session.commit()
items = coll.get_items(session=test_session)
items = coll.get_items(test_session)
fg_ids = {i.file_group_id for i in items}
assert v2.file_group_id not in fg_ids, 'Removed video should not be visible in collection fetch methods'
# Files should remain on disk
@ -100,7 +100,7 @@ async def test_collection_lifecycle_end_to_end(async_client, test_session: Sessi
assert not v3.file_group.primary_path.exists()
# The collection should no longer reference it (FK cascade or absence after populate)
items = coll.get_items(session=test_session)
items = coll.get_items(test_session)
fg_ids = {i.file_group_id for i in items}
assert v3.file_group_id not in fg_ids

12
wrolpi/collections/test/test_collection_tagging.py
View file

@ -38,11 +38,11 @@ async def test_tagging_directory_collection_moves_files_and_updates_config(
'directory': str(rel_dir),
'kind': 'channel',
}
coll = Collection.from_config(cfg, session=test_session)
coll = Collection.from_config(test_session, cfg)
test_session.commit()
# Verify it populated
assert {i.file_group_id for i in coll.get_items(session=test_session)} == {v1.file_group_id, v2.file_group_id}
assert {i.file_group_id for i in coll.get_items(test_session)} == {v1.file_group_id, v2.file_group_id}
# Dump config and verify (config stores absolute paths)
collections_config.dump_config()
@ -63,7 +63,7 @@ async def test_tagging_directory_collection_moves_files_and_updates_config(
dest_dir.mkdir(parents=True, exist_ok=True)
# Perform the move using the model helper
await coll.move_collection(dest_dir, session=test_session)
await coll.move_collection(dest_dir, test_session)
test_session.commit()
# Validate the directory changed and files moved
@ -105,9 +105,9 @@ async def test_tagging_unrestricted_collection_updates_config_only(
test_session.commit()
# Create an unrestricted collection and add both files
coll = Collection.from_config({'name': 'Loose Stuff', 'kind': 'channel'}, session=test_session)
coll = Collection.from_config(test_session, {'name': 'Loose Stuff', 'kind': 'channel'})
# Unrestricted, so add manually
coll.add_file_groups([fg1, fg2], session=test_session)
coll.add_file_groups(test_session, [fg1, fg2])
test_session.commit()
# Dump config and verify directory is absent/None
@ -201,10 +201,10 @@ async def test_tag_collection_removes_tag_and_updates_directory(test_session, te
# Remove the tag AND update the directory
result = await tag_collection(
test_session,
collection_id=collection.id,
tag_name=None, # Remove tag
directory=str(untagged_dir), # New directory
session=test_session
)
test_session.commit()
await await_switches()

30
wrolpi/collections/test/test_domain_collections.py
View file

@ -39,7 +39,7 @@ class TestDomainValidation:
'name': 'example.com',
'kind': 'domain',
}
collection = Collection.from_config(config, session=test_session)
collection = Collection.from_config(test_session, config)
assert collection.name == 'example.com'
assert collection.kind == 'domain'
assert collection.directory is None
@ -50,7 +50,7 @@ class TestDomainValidation:
'kind': 'domain',
}
with pytest.raises(ValueError, match='Invalid domain name'):
Collection.from_config(invalid_config, session=test_session)
Collection.from_config(test_session, invalid_config)
# Domain starting with dot should fail
invalid_config2 = {
@ -58,7 +58,7 @@ class TestDomainValidation:
'kind': 'domain',
}
with pytest.raises(ValueError, match='Invalid domain name'):
Collection.from_config(invalid_config2, session=test_session)
Collection.from_config(test_session, invalid_config2)
# Domain ending with dot should fail
invalid_config3 = {
@ -66,7 +66,7 @@ class TestDomainValidation:
'kind': 'domain',
}
with pytest.raises(ValueError, match='Invalid domain name'):
Collection.from_config(invalid_config3, session=test_session)
Collection.from_config(test_session, invalid_config3)
def test_from_config_allows_any_name_for_channel_kind(self, test_session: Session):
"""Test that channel kind does not enforce domain validation."""
@ -75,7 +75,7 @@ class TestDomainValidation:
'name': 'My Channel Without Dots',
'kind': 'channel',
}
collection = Collection.from_config(config, session=test_session)
collection = Collection.from_config(test_session, config)
assert collection.name == 'My Channel Without Dots'
assert collection.kind == 'channel'
@ -86,7 +86,7 @@ class TestDomainValidation:
'kind': 'domain',
'description': 'Archives from example.org',
}
collection = Collection.from_config(config, session=test_session)
collection = Collection.from_config(test_session, config)
test_session.commit()
assert collection.name == 'example.org'
@ -101,7 +101,7 @@ class TestDomainValidation:
'name': 'blog.example.com',
'kind': 'domain',
}
collection = Collection.from_config(config, session=test_session)
collection = Collection.from_config(test_session, config)
test_session.commit()
assert collection.name == 'blog.example.com'
@ -115,7 +115,7 @@ class TestDomainValidation:
'kind': 'domain',
'description': 'Initial description',
}
collection = Collection.from_config(config, session=test_session)
collection = Collection.from_config(test_session, config)
test_session.commit()
collection_id = collection.id
@ -125,7 +125,7 @@ class TestDomainValidation:
'kind': 'domain',
'description': 'Updated description',
}
updated = Collection.from_config(updated_config, session=test_session)
updated = Collection.from_config(test_session, updated_config)
assert updated.id == collection_id
assert updated.description == 'Updated description'
@ -136,7 +136,7 @@ class TestDomainValidation:
'kind': 'domain',
'description': 'Test domain',
}
collection = Collection.from_config(config, session=test_session)
collection = Collection.from_config(test_session, config)
test_session.commit()
exported = collection.to_config()
@ -157,8 +157,8 @@ class TestDomainValidation:
'kind': 'channel',
}
domain_coll = Collection.from_config(domain_config, session=test_session)
channel_coll = Collection.from_config(channel_config, session=test_session)
domain_coll = Collection.from_config(test_session, domain_config)
channel_coll = Collection.from_config(test_session, channel_config)
test_session.commit()
# Both should exist with different IDs
@ -190,7 +190,7 @@ class TestDirectoryValidation:
# Should raise ValidationError
with pytest.raises(ValidationError, match="must be under media directory"):
Collection.from_config(config, test_session)
Collection.from_config(test_session, config)
def test_from_config_accepts_relative_path(self, test_session, test_directory):
"""Test that from_config converts relative paths to absolute under media directory."""
@ -202,7 +202,7 @@ class TestDirectoryValidation:
'directory': 'archive/example.com' # Relative path
}
collection = Collection.from_config(config, test_session)
collection = Collection.from_config(test_session, config)
test_session.flush()
# Should be converted to absolute path under media directory
@ -223,7 +223,7 @@ class TestDirectoryValidation:
'directory': str(valid_path)
}
collection = Collection.from_config(config, test_session)
collection = Collection.from_config(test_session, config)
test_session.flush()
# Should accept the path

1
wrolpi/collections/test/test_domain_moving.py
View file

@ -368,6 +368,7 @@ async def test_tag_domain_comprehensive(
# Create recurring download which uses the domain's directory
download = test_download_manager.recurring_download(
test_session,
'https://example.com/feed',
60,
test_downloader.name,

14
wrolpi/collections/test/test_domains_config.py
View file

@ -51,9 +51,9 @@ class TestDomainsConfig:
test_directory: pathlib.Path, async_client):
"""Test that dumping domains.yaml only exports domain collections."""
# Create some domain and channel collections
domain1 = Collection.from_config({'name': 'example.com', 'kind': 'domain'}, session=test_session)
domain2 = Collection.from_config({'name': 'test.org', 'kind': 'domain'}, session=test_session)
channel = Collection.from_config({'name': 'My Channel', 'kind': 'channel'}, session=test_session)
domain1 = Collection.from_config(test_session, {'name': 'example.com', 'kind': 'domain'})
domain2 = Collection.from_config(test_session, {'name': 'test.org', 'kind': 'domain'})
channel = Collection.from_config(test_session, {'name': 'My Channel', 'kind': 'channel'})
test_session.commit()
# Dump to config
@ -126,8 +126,8 @@ class TestDomainsConfig:
async_client):
"""Test that domains removed from config are deleted from database."""
# Create two domain collections
domain1 = Collection.from_config({'name': 'example.com', 'kind': 'domain'}, session=test_session)
domain2 = Collection.from_config({'name': 'test.org', 'kind': 'domain'}, session=test_session)
domain1 = Collection.from_config(test_session, {'name': 'example.com', 'kind': 'domain'})
domain2 = Collection.from_config(test_session, {'name': 'test.org', 'kind': 'domain'})
test_session.commit()
# Import config with only one domain
@ -152,11 +152,11 @@ class TestDomainsConfig:
async_client):
"""Test that updating a domain in config updates the existing collection."""
# Create initial domain
domain = Collection.from_config({
domain = Collection.from_config(test_session, {
'name': 'example.com',
'kind': 'domain',
'description': 'Original description',
}, session=test_session)
})
test_session.commit()
domain_id = domain.id

36
wrolpi/collections/test/test_models.py Normal file
View file

@ -0,0 +1,36 @@
import pytest
from wrolpi.collections import Collection
from wrolpi.errors import UnknownCollection
def test_collection_get_by_id(test_session, test_directory):
"""Collection.get_by_id should accept session as first argument, id as second."""
collection = Collection(name='test collection', directory=test_directory / 'test_collection')
test_session.add(collection)
test_session.flush()
# Session must be first argument (session-first pattern)
found = Collection.get_by_id(test_session, collection.id)
assert found is not None
assert found.id == collection.id
# Non-existent ID should return None
not_found = Collection.get_by_id(test_session, 99999)
assert not_found is None
def test_collection_find_by_id(test_session, test_directory):
"""Collection.find_by_id should accept session as first argument, id as second."""
collection = Collection(name='test collection', directory=test_directory / 'test_collection')
test_session.add(collection)
test_session.flush()
# Session must be first argument (session-first pattern)
found = Collection.find_by_id(test_session, collection.id)
assert found is not None
assert found.id == collection.id
# Non-existent ID should raise UnknownCollection
with pytest.raises(UnknownCollection):
Collection.find_by_id(test_session, 99999)

10
wrolpi/common.py
View file

@ -324,20 +324,20 @@ class ModelHelper:
return d
@classmethod
def upsert(cls, file, session: Session) -> Base:
def upsert(cls, session: Session, file) -> Base:
raise NotImplementedError('This model has not defined this method.')
@staticmethod
def get_by_path(path, session) -> Base:
def get_by_path(session: Session, path) -> Base:
raise NotImplementedError('This model has not defined this method.')
@staticmethod
def get_by_id(id_: int, session: Session = None) -> Optional[Base]:
def get_by_id(session: Session, id_: int) -> Optional[Base]:
"""Attempts to get a model instance by its id. Returns None if no instance can be found."""
raise NotImplementedError('This model has not defined this method.')
@staticmethod
def find_by_id(id_: int, session: Session = None) -> Base:
def find_by_id(session: Session, id_: int) -> Base:
"""Get a model instance by its id, raise an error if no instance is found."""
raise NotImplementedError('This model has not defined this method.')
@ -346,7 +346,7 @@ class ModelHelper:
raise NotImplementedError('This model has not defined this method.')
@staticmethod
def do_model(file_group, session: Session):
def do_model(session: Session, file_group):
raise NotImplementedError('This model has not defined this method.')
def flush(self, session: Session = None):

10
wrolpi/conftest.py
View file

@ -148,8 +148,11 @@ async def background_task_listener(request, response):
@pytest.fixture
async def async_client(test_directory) -> SanicASGITestClient:
"""Get an Async Sanic Test Client with all default routes attached."""
async def async_client(test_directory, test_session) -> SanicASGITestClient:
"""Get an Async Sanic Test Client with all default routes attached.
Depends on test_session to ensure the database mock is in place for session middleware.
"""
api_app.signalize()
attach_shared_contexts(api_app)
initialize_configs_contexts(api_app)
@ -414,6 +417,7 @@ def make_files_structure(test_directory) -> Callable[[Union[List[Union[pathlib.P
def create_files(paths: Union[List, Dict], file_groups: bool = False, session: Session = None) \
-> List[pathlib.Path]:
# Note: session is optional here for backwards compatibility - file_groups=True requires session
files = []
@iterify(list)
@ -534,7 +538,7 @@ def tag_factory(test_session) -> Callable[[Optional[str]], Awaitable[Tag]]:
async def factory(name: str = None) -> Tag:
if not name:
name = names.pop(0)
tag = await upsert_tag(name, f'#{str(count) * 6}', session=test_session)
tag = await upsert_tag(test_session, name, f'#{str(count) * 6}')
return tag
return factory

118
wrolpi/db.py
View file

@ -1,8 +1,6 @@
import inspect
import types
from contextlib import contextmanager
from functools import wraps
from typing import ContextManager, Tuple, List, Union, Type, Generator
from typing import Tuple, List, Union, Type, Generator, Any
import psycopg2
import psycopg2.extensions
@ -12,7 +10,7 @@ from psycopg2._psycopg import cursor
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import NullPool
from wrolpi.common import logger, Base, partition
from wrolpi.common import logger, Base
from wrolpi.vars import DB_NAME, DB_USER, DB_PASSWORD, DB_HOST, DB_PORT, DOCKERIZED, PYTEST
logger = logger.getChild(__name__)
@ -81,7 +79,7 @@ def get_db_context() -> Tuple[sqlalchemy.engine.Engine, Session]:
@contextmanager
def get_db_session(commit: bool = False) -> ContextManager[Session]:
def get_db_session(commit: bool = False) -> Generator[Session, Any, None]:
"""
Context manager that creates a DB session. This will automatically rollback changes, unless `commit` is True.
"""
@ -95,24 +93,38 @@ def get_db_session(commit: bool = False) -> ContextManager[Session]:
raise
finally:
# Rollback only if a transaction hasn't been committed.
if session.transaction.is_active:
# In tests, the test_session fixture manages the session lifecycle,
# so we should not rollback here - that would undo other test operations.
if not PYTEST and session.transaction.is_active:
session.rollback()
@contextmanager
def get_db_conn(isolation_level=psycopg2.extensions.ISOLATION_LEVEL_DEFAULT):
local_engine, session = get_db_context()
connection = local_engine.raw_connection()
connection.set_isolation_level(isolation_level)
try:
yield connection, session
except sqlalchemy.exc.DatabaseError:
session.rollback()
raise
finally:
# Rollback only if a transaction hasn't been committed.
if session.transaction.is_active:
connection.rollback()
if PYTEST:
# During tests, use the session's connection to avoid lock issues.
# The test_session is mocked to be shared, so getting a raw_connection would
# create a new connection that blocks waiting for locks held by test_session.
connection = session.connection().connection
try:
yield connection, session
except sqlalchemy.exc.DatabaseError:
session.rollback()
raise
# Don't rollback in finally during tests - test_session manages this
else:
connection = local_engine.raw_connection()
connection.set_isolation_level(isolation_level)
try:
yield connection, session
except sqlalchemy.exc.DatabaseError:
session.rollback()
raise
finally:
# Rollback only if a transaction hasn't been committed.
if session.transaction.is_active:
connection.rollback()
@contextmanager
@ -124,82 +136,20 @@ def get_db_curs(commit: bool = False, isolation_level=psycopg2.extensions.ISOLAT
curs = connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
try:
yield curs
if commit:
if commit and not PYTEST:
# Don't commit during tests - test_session manages the transaction
connection.commit()
except sqlalchemy.exc.DatabaseError:
session.rollback()
raise
finally:
# Rollback only if a transaction hasn't been committed.
if session.transaction.is_active:
# Don't rollback during tests - test_session manages the transaction
if not PYTEST and session.transaction.is_active:
connection.rollback()
def optional_session(commit: Union[callable, bool] = False) -> callable:
"""
Wraps a function, if a Session is passed it will be used. Otherwise, a new session will be
created and passed to the function.
"""
def find_session(*f_args, session: Session = None, **f_kwargs):
session = f_kwargs.get('session', session)
if not session:
# Find the session in the args.
session, new_args = partition(lambda i: isinstance(i, Session), f_args)
if session:
session = session[0]
f_args = tuple(new_args)
return session, f_args, f_kwargs
def call_func(func, session, args, kwargs):
if session:
result = func(*args, session=session, **kwargs)
if commit:
session.commit()
return result
else:
with get_db_session(commit=commit) as session:
result = func(*args, session=session, **kwargs)
return result
async def call_func_async(func, session, args, kwargs):
"""Async version of call_func that properly awaits the coroutine before committing."""
if session:
result = await func(*args, session=session, **kwargs)
if commit:
session.commit()
return result
else:
with get_db_session(commit=commit) as session:
result = await func(*args, session=session, **kwargs)
return result
def wrapper(*w_args, **w_kwargs):
if len(w_args) == 1 and len(w_kwargs) == 0 and callable(w_args[0]):
func = w_args[0]
if inspect.iscoroutinefunction(func):
@wraps(func)
async def wrapped(*args, session: Session = None, **kwargs):
session, args, kwargs = find_session(*args, session=session, **kwargs)
return await call_func_async(func, session, args, kwargs)
else:
@wraps(func)
def wrapped(*args, session: Session = None, **kwargs):
session, args, kwargs = find_session(*args, session=session, **kwargs)
return call_func(func, session, args, kwargs)
return wrapped
else:
# No params were passed to this wrapper, `commit` is actually the function we are wrapping.
session_, w_args, w_kwargs = find_session(*w_args, **w_kwargs)
return call_func(commit, session_, w_args, w_kwargs)
return wrapper
@optional_session
def get_ranked_models(ranked_primary_keys: List, model: Type[Base], session: Session = None) -> List[Base]:
def get_ranked_models(ranked_primary_keys: List, model: Type[Base], session: Session) -> List[Base]:
"""Get all objects whose primary keys are in the `ranked_primary_keys`, preserve their order."""
pkey = sqlalchemy.inspect(model).primary_key[0]
pkey_name = pkey.name

133
wrolpi/downloader.py
View file

@ -32,7 +32,7 @@ from wrolpi.common import Base, ModelHelper, logger, wrol_mode_check, zig_zag, C
wrol_mode_enabled, background_task, get_absolute_media_path, timer, aiohttp_get, \
get_download_info, trim_file_name, get_wrolpi_config, TRACE_LEVEL
from wrolpi.dates import TZDateTime, now, Seconds
from wrolpi.db import get_db_session, get_db_curs, optional_session
from wrolpi.db import get_db_session, get_db_curs
from wrolpi.errors import InvalidDownload, UnrecoverableDownloadError, UnknownDownload, ValidationError, DownloadError
from wrolpi.events import Events
from wrolpi.media_path import MediaPathType
@ -252,33 +252,29 @@ class Download(ModelHelper, Base): # noqa
save_downloads_config.activate_switch()
@staticmethod
def get_by_id(id_: int, session: Session = None) -> Optional['Download']:
def get_by_id(session: Session, id_: int) -> Optional['Download']:
download = session.query(Download).filter(Download.id == id_).one_or_none()
return download
@classmethod
@optional_session
def find_by_id(cls, id_: int, session: Session = None) -> 'Download':
if download := cls.get_by_id(id_, session):
def find_by_id(cls, session: Session, id_) -> 'Download':
if download := cls.get_by_id(session, id_):
return download
raise UnknownDownload(f'Cannot find Download with id: {id_}')
@staticmethod
@optional_session
def get_by_url(url: str, session: Session = None) -> Optional['Download']:
def get_by_url(session: Session, url: str) -> Optional['Download']:
download = session.query(Download).filter(Download.url == url).one_or_none()
return download
@classmethod
@optional_session
def find_by_url(cls, url: str, session: Session = None) -> 'Download':
if download := cls.get_by_url(url, session):
def find_by_url(cls, session: Session, url: str) -> 'Download':
if download := cls.get_by_url(session, url):
return download
raise UnknownDownload(f'Cannot find Download with URL: {url}')
@classmethod
@optional_session
def get_all_by_destination(cls, destination: str | pathlib.Path, session: Session = None) -> List['Download']:
def get_all_by_destination(cls, session: Session, destination: str | pathlib.Path) -> List['Download']:
destination = str(destination)
downloads = session.query(Download).filter_by(destination=destination).all()
return downloads
@ -310,8 +306,7 @@ class Downloader:
async def do_download(self, download: Download) -> DownloadResult:
raise NotImplementedError()
@optional_session
def already_downloaded(self, *urls: List[str], session: Session = None):
def already_downloaded(self, session: Session, *urls: List[str]):
raise NotImplementedError()
async def process_runner(self, download: Download, cmd: Tuple[str | pathlib.Path, ...], cwd: pathlib.Path,
@ -564,12 +559,12 @@ class DownloadManager:
return downloader
raise InvalidDownload(f'Cannot find downloader with name {name}')
def get_or_create_download(self, url: str, session: Session, reset_attempts: bool = False) -> Download:
def get_or_create_download(self, session: Session, url: str, reset_attempts: bool = False) -> Download:
"""Get a Download by its URL, if it cannot be found create one."""
if not url:
raise ValueError('Download must have a URL')
if download := Download.get_by_url(url, session=session):
if download := Download.get_by_url(session, url):
return download
if self.is_skipped(url):
@ -582,8 +577,7 @@ class DownloadManager:
session.add(download)
return download
@optional_session
def create_downloads(self, urls: List[str], downloader_name: str, session: Session = None,
def create_downloads(self, session: Session, urls: List[str], downloader_name: str,
reset_attempts: bool = False, sub_downloader_name: str = None,
destination: str | pathlib.Path = None, tag_names: List[str] = None, settings: dict = None) \
-> List[Download]:
@ -619,7 +613,7 @@ class DownloadManager:
self.log_warning(f'Skipping {url} because it is in the download_manager.yaml skip list.')
continue
download = self.get_or_create_download(url, session, reset_attempts=reset_attempts)
download = self.get_or_create_download(session, url, reset_attempts=reset_attempts)
# Download may have failed, try again.
download.renew(reset_attempts=reset_attempts)
download.destination = destination or None
@ -649,18 +643,16 @@ class DownloadManager:
return downloads
@optional_session
def create_download(self, url: str, downloader_name: str, session: Session = None, reset_attempts: bool = False,
def create_download(self, session: Session, url: str, downloader_name: str, reset_attempts: bool = False,
sub_downloader_name: str = None, destination: str | pathlib.Path = None,
tag_names: List[str] = None, settings: Dict = None) -> Download:
"""Schedule a URL for download. If the URL failed previously, it may be retried."""
downloads = self.create_downloads([url], session=session, downloader_name=downloader_name,
downloads = self.create_downloads(session, [url], downloader_name=downloader_name,
reset_attempts=reset_attempts, sub_downloader_name=sub_downloader_name,
destination=destination, tag_names=tag_names, settings=settings)
return downloads[0]
@optional_session
def recurring_download(self, url: str, frequency: int, downloader_name: str, session: Session = None,
def recurring_download(self, session: Session, url: str, frequency: int, downloader_name: str,
sub_downloader_name: str = None, reset_attempts: bool = False,
destination: str | pathlib.Path = None, tag_names: List[str] = None,
settings: Dict = None, collection_id: int = None) -> Download:
@ -672,7 +664,7 @@ class DownloadManager:
if downloader_name == ScrapeHTMLDownloader.name and frequency:
raise InvalidDownload(f'Cannot schedule recurring download for {downloader_name=}')
download, = self.create_downloads([url, ], session=session, downloader_name=downloader_name,
download, = self.create_downloads(session, [url, ], downloader_name=downloader_name,
reset_attempts=reset_attempts, sub_downloader_name=sub_downloader_name,
destination=destination, tag_names=tag_names, settings=settings)
download.frequency = frequency
@ -681,20 +673,18 @@ class DownloadManager:
# Only recurring Downloads can be Collection Downloads - look up via Channel.
from modules.videos.models import Channel
if channel := Channel.get_by_url(url=download.url, session=session):
if channel := Channel.get_by_url(session, url=download.url):
download.collection_id = channel.collection_id
session.commit()
return download
@optional_session
def update_download(self, id_: int, url: str, downloader: str,
def update_download(self, session: Session, id_: int, url: str, downloader: str,
destination: str | pathlib.Path = None, tag_names: List[str] = None,
sub_downloader: str | None = None, frequency: int = None,
settings: Dict = None, collection_id: int = None,
session: Session = None) -> Download:
download = Download.find_by_id(id_, session=session)
settings: Dict = None, collection_id: int = None) -> Download:
download = Download.find_by_id(session, id_)
if collection_id and not frequency:
raise InvalidDownload(f'A once-download cannot be associated with a Collection')
download.url = url
@ -714,8 +704,7 @@ class DownloadManager:
return download
@wrol_mode_check
@optional_session
async def dispatch_downloads(self, session: Session = None):
async def dispatch_downloads(self):
"""Dispatch Sanic signals to start downloads. This only starts as many downloads as the
SIMULTANEOUS_DOWNLOAD_DOMAINS variable."""
if not self.can_download:
@ -728,23 +717,25 @@ class DownloadManager:
return
# Find download whose domain isn't already being downloaded.
new_downloads = list(session.query(Download).filter(
Download.status == 'new',
Download.domain not in self.processing_domains,
).order_by(
Download.frequency.is_(None),
Download.frequency,
Download.id))
count = 0
for download in new_downloads:
domain = download.domain
if domain not in self.processing_domains and len(self.processing_domains) < SIMULTANEOUS_DOWNLOAD_DOMAINS:
self._add_processing_domain(domain)
context = dict(download_id=download.id, download_url=download.url)
await api_app.dispatch('wrolpi.download.download', context=context)
count += 1
if count:
self.log_debug(f'Added {count} downloads to queue.')
with get_db_session(commit=True) as session:
new_downloads = list(session.query(Download).filter(
Download.status == 'new',
Download.domain not in self.processing_domains,
).order_by(
Download.frequency.is_(None),
Download.frequency,
Download.id))
count = 0
for download in new_downloads:
domain = download.domain
if domain not in self.processing_domains and len(
self.processing_domains) < SIMULTANEOUS_DOWNLOAD_DOMAINS:
self._add_processing_domain(domain)
context = dict(download_id=download.id, download_url=download.url)
await api_app.dispatch('wrolpi.download.download', context=context)
count += 1
if count:
self.log_debug(f'Added {count} downloads to queue.')
async def do_downloads(self):
"""Schedule any downloads that are new.
@ -809,7 +800,6 @@ class DownloadManager:
stmt = "UPDATE download SET status='new' WHERE status='pending' OR status='deferred'"
curs.execute(stmt)
@optional_session
def get_new_downloads(self, session: Session) -> Generator[Download, None, None]:
"""
Get all "new" downloads. This method fetches the first download each iteration, so it will fetch downloads
@ -826,8 +816,7 @@ class DownloadManager:
last = download
yield download
@optional_session
def get_recurring_downloads(self, session: Session = None, limit: int = None):
def get_recurring_downloads(self, session: Session, limit: int = None):
"""Get all Downloads that will be downloaded in the future."""
query = session.query(Download).filter(
Download.frequency != None # noqa
@ -840,8 +829,7 @@ class DownloadManager:
downloads = query.all()
return downloads
@optional_session
def get_once_downloads(self, session: Session = None, limit: int = None):
def get_once_downloads(self, session: Session, limit: int = None):
"""Get all Downloads that will not reoccur."""
query = session.query(Download).filter(
Download.frequency == None # noqa
@ -863,8 +851,7 @@ class DownloadManager:
renewed_count = 0
for download in recurring:
# A new download may not have a `next_download`, create it if necessary.
download.next_download = download.next_download or self.calculate_next_download(download,
session=session)
download.next_download = download.next_download or self.calculate_next_download(session, download)
if download.next_download < now_ and download.status not in (
DownloadStatus.new, DownloadStatus.pending):
download.renew()
@ -879,19 +866,17 @@ class DownloadManager:
downloads = list(session.query(Download).all())
return downloads
@optional_session
def delete_download(self, download_id: int, session: Session = None):
def delete_download(self, session: Session, download_id: int):
"""Delete a Download. Returns True if a Download was deleted, otherwise return False."""
if download := Download.get_by_id(download_id, session=session):
if download := Download.get_by_id(session, download_id):
# This saves the config twice.
download.delete()
return True
return False
@optional_session
def restart_download(self, download_id: int, session: Session = None) -> Download:
def restart_download(self, session: Session, download_id: int) -> Download:
"""Renews a download and resets its download attempts."""
download = Download.find_by_id(download_id, session=session)
download = Download.find_by_id(session, download_id)
download.renew(reset_attempts=True)
session.commit()
@ -905,7 +890,7 @@ class DownloadManager:
api_app.shared_ctx.download_manager_data.update(download_manager_data)
with get_db_session(commit=True) as session:
if download := Download.get_by_id(download_id, session=session):
if download := Download.get_by_id(session, download_id):
download.error = 'User stopped this download'
download.fail()
@ -1045,8 +1030,7 @@ class DownloadManager:
return summary
@staticmethod
@optional_session
def calculate_next_download(download: Download, session: Session) -> Optional[datetime]:
def calculate_next_download(session: Session, download: Download) -> Optional[datetime]:
"""
If the download is "deferred", download soon. But, slowly increase the time between attempts.
@ -1099,7 +1083,6 @@ class DownloadManager:
stmt = stmt.where(Download.status == status)
return stmt
@optional_session
def delete_completed(self, session: Session) -> List[int]:
"""Delete any completed download records."""
stmt = self._delete_downloads_q(once=True, status=DownloadStatus.complete)
@ -1107,7 +1090,6 @@ class DownloadManager:
session.commit()
return deleted_ids
@optional_session
def delete_failed(self, session: Session):
"""Delete any failed download records."""
stmt = self._delete_downloads_q(once=True, status=DownloadStatus.failed, returning=Download.url)
@ -1118,7 +1100,6 @@ class DownloadManager:
session.commit()
@optional_session
def delete_once(self, session: Session):
"""Delete all once-download records."""
stmt = self._delete_downloads_q(once=True)
@ -1165,7 +1146,7 @@ async def signal_download_download(download_id: int, download_url: str):
with get_db_session(commit=True) as session:
# Mark the download as started in new session so the change is committed.
download = Download.find_by_id(download_id, session=session)
download = Download.find_by_id(session, download_id)
download.started()
download_domain = download.domain
@ -1205,12 +1186,12 @@ async def signal_download_download(download_id: int, download_url: str):
download.location = result.location or download.location or None
# Clear any old errors if the download succeeded.
download.error = result.error if result.error else None
download.next_download = download_manager.calculate_next_download(download, session)
download.next_download = download_manager.calculate_next_download(session, download)
urls = download.filter_excluded(result.downloads) if result.downloads else None
if urls:
worker_logger.info(f'Adding {len(result.downloads)} downloads from result of {download.url}')
download_manager.create_downloads(urls, session, downloader_name=download.sub_downloader,
download_manager.create_downloads(session, urls, downloader_name=download.sub_downloader,
settings=result.settings)
if try_again is False and not download.frequency:
@ -1405,7 +1386,7 @@ class DownloadManagerConfig(ConfigFile):
subscription = session.query(ZimSubscription).filter_by(name=name, language=language).one_or_none()
if not subscription:
subscription = ZimSubscription(name=name, language=language)
subscription.change_download(download.url, download.frequency, session=session)
subscription.change_download(session, download.url, download.frequency)
session.add(subscription)
need_commit = True
@ -1455,7 +1436,6 @@ def save_downloads_config():
save_downloads_config: ActivateSwitchMethod
@optional_session
async def import_downloads_config(session: Session):
"""Upsert all Downloads in the Download Manager Config into the DB.
@ -1512,7 +1492,8 @@ class RSSDownloader(Downloader):
logger.warning(f'RSS entry {idx} did not have a link!')
# Only download new URLs.
urls = [i for i in urls if i not in [i.url for i in sub_downloader.already_downloaded(*urls)]]
with get_db_session() as session:
urls = [i for i in urls if i not in [i.url for i in sub_downloader.already_downloaded(session, *urls)]]
# Remove skipped URLs before duration checks. (Typically done after this by the download worker).
urls = [i for i in urls if not download_manager.is_skipped(i)]
@ -1602,7 +1583,7 @@ class RSSDownloader(Downloader):
from modules.videos.models import Channel
channel = Channel.get_by_source_id(session, f'UC{yt_channel_id}')
if channel:
download_ = Download.get_by_id(download_id, session=session)
download_ = Download.get_by_id(session, download_id)
download_.collection = channel.collection
download_.collection_id = channel.collection_id
download_.location = download_.location or channel.location

42
wrolpi/files/api.py
View file

@ -30,10 +30,11 @@ logger = logger.getChild(__name__)
body=schema.FilesRequest,
)
@validate(schema.FilesRequest)
async def get_files(_: Request, body: schema.FilesRequest):
async def get_files(request: Request, body: schema.FilesRequest):
directories = body.directories or []
session = request.ctx.session
files = lib.list_directories_contents(directories)
files = lib.list_directories_contents(session, directories)
return json_response({'files': files})
@ -43,13 +44,14 @@ async def get_files(_: Request, body: schema.FilesRequest):
body=schema.FileRequest,
)
@validate(schema.FileRequest)
async def get_file(_: Request, body: schema.FileRequest):
async def get_file(request: Request, body: schema.FileRequest):
session = request.ctx.session
try:
file = lib.get_file_dict(body.file)
file = lib.get_file_dict(session, body.file)
except FileNotFoundError:
raise InvalidFile()
background_task(lib.set_file_viewed(get_media_directory() / body.file))
background_task(lib.set_file_viewed(session, get_media_directory() / body.file))
return json_response({'file': file})
@ -135,7 +137,7 @@ def post_directories(_, body: schema.DirectoriesRequest):
body=schema.DirectoriesSearchRequest,
)
@validate(schema.DirectoriesSearchRequest)
async def post_search_directories(_, body: schema.DirectoriesSearchRequest):
async def post_search_directories(request, body: schema.DirectoriesSearchRequest):
if len(body.path) <= 1:
# Need more data before searching.
return json_response({
@ -145,6 +147,7 @@ async def post_search_directories(_, body: schema.DirectoriesSearchRequest):
'domain_directories': [],
})
session = request.ctx.session
path = get_media_directory() / body.path
if logger.isEnabledFor(TRACE_LEVEL):
logger.trace(f'post_search_directories: {path=}')
@ -158,7 +161,7 @@ async def post_search_directories(_, body: schema.DirectoriesSearchRequest):
# Search Channels by name.
from modules.videos.channel.lib import search_channels_by_name
channels = await search_channels_by_name(name=body.path)
channels = await search_channels_by_name(session, name=body.path)
channel_directories = [dict(path=i.directory, name=i.name) for i in channels]
channel_paths = [i['path'] for i in channel_directories]
if logger.isEnabledFor(TRACE_LEVEL):
@ -166,7 +169,7 @@ async def post_search_directories(_, body: schema.DirectoriesSearchRequest):
# Search Domains by name.
from modules.archive.lib import search_domains_by_name
domains = await search_domains_by_name(name=body.path)
domains = await search_domains_by_name(session, name=body.path)
# search_domains_by_name returns dicts with 'directory' and 'domain' keys
domain_directories = [dict(path=i['directory'], domain=i['domain']) for i in domains]
domain_paths = [i['path'] for i in domain_directories]
@ -175,7 +178,7 @@ async def post_search_directories(_, body: schema.DirectoriesSearchRequest):
# Get all Directory that match but do not contain the above directories.
excluded = [str(i) for i in channel_paths + domain_paths + matching_directories]
directories = await lib.search_directories_by_name(name=body.path, excluded=excluded)
directories = await lib.search_directories_by_name(session, name=body.path, excluded=excluded)
# Return only the top 20 directories.
directories = [i.__json__() for i in directories]
@ -251,7 +254,7 @@ async def post_create_directory(_: Request, body: schema.Directory):
body=schema.Move,
)
@validate(schema.Move)
async def post_move(_: Request, body: schema.Move):
async def post_move(request: Request, body: schema.Move):
destination = get_media_directory() / body.destination
if not destination.is_dir():
raise UnknownDirectory(f'Cannot move files into {destination} because it does not exist')
@ -263,7 +266,7 @@ async def post_move(_: Request, body: schema.Move):
raise FileConflict(f'Cannot move mounted directory!')
try:
await lib.move(destination, *sources)
await lib.move(request.ctx.session, destination, *sources)
Events.send_file_move_completed(f'Moving files to {body.destination} succeeded')
except Exception as e:
Events.send_file_move_failed(f'Moving files to {body.destination} failed!')
@ -278,7 +281,7 @@ async def post_move(_: Request, body: schema.Move):
body=schema.Rename,
)
@validate(schema.Rename)
async def post_rename(_: Request, body: schema.Rename):
async def post_rename(request: Request, body: schema.Rename):
path = get_media_directory() / body.path
new_path = path.with_name(body.new_name)
if not path.exists():
@ -287,7 +290,7 @@ async def post_rename(_: Request, body: schema.Rename):
raise FileConflict(f'File already exists: {new_path}')
try:
await lib.rename(path, body.new_name, send_events=True)
await lib.rename(request.ctx.session, path, body.new_name, send_events=True)
except Exception as e:
logger.error(f'Failed to rename {path} to {new_path}', exc_info=e)
raise FileConflict(f'Failed to rename {path} to {new_path}') from e
@ -302,7 +305,7 @@ async def post_rename(_: Request, body: schema.Rename):
@files_bp.post('/tag')
@validate(schema.TagFileGroupPost)
async def post_tag_file_group(_, body: schema.TagFileGroupPost):
async def post_tag_file_group(request: Request, body: schema.TagFileGroupPost):
if not body.tag_id and not body.tag_name:
return json_response(
dict(error='tag_id and tag_name cannot both be empty'),
@ -314,14 +317,16 @@ async def post_tag_file_group(_, body: schema.TagFileGroupPost):
HTTPStatus.BAD_REQUEST,
)
await lib.add_file_group_tag(body.file_group_id, body.file_group_primary_path, body.tag_name, body.tag_id)
await lib.add_file_group_tag(request.ctx.session, body.file_group_id, body.file_group_primary_path, body.tag_name,
body.tag_id)
return response.empty(HTTPStatus.CREATED)
@files_bp.post('/untag')
@validate(schema.TagFileGroupPost)
async def post_untag_file_group(_, body: schema.TagFileGroupPost):
await lib.remove_file_group_tag(body.file_group_id, body.file_group_primary_path, body.tag_name, body.tag_id)
async def post_untag_file_group(request: Request, body: schema.TagFileGroupPost):
await lib.remove_file_group_tag(request.ctx.session, body.file_group_id, body.file_group_primary_path,
body.tag_name, body.tag_id)
return response.empty(HTTPStatus.NO_CONTENT)
@ -382,8 +387,9 @@ async def post_upload(request: Request):
if not isinstance(tag_names, list):
raise FileUploadFailed('tag_names must be a list')
tag_names = [str(i) for i in tag_names]
session = request.ctx.session
for tag_name in tag_names:
if not Tag.get_by_name(tag_name):
if not Tag.get_by_name(session, tag_name):
raise FileUploadFailed(f'Tag does not exist: {tag_name}')
filename = pathlib.Path(filename.lstrip('/'))

2
wrolpi/files/ebooks.py
View file

@ -177,7 +177,7 @@ class EBook(ModelHelper, Base):
return False
@staticmethod
def do_model(file_group: FileGroup, session: Session) -> 'EBook':
def do_model(session: Session, file_group: FileGroup) -> 'EBook':
ebook = model_ebook(file_group, session)
file_group.indexed = True
return ebook

112
wrolpi/files/lib.py
View file

@ -28,7 +28,7 @@ from wrolpi.common import get_media_directory, wrol_mode_check, logger, limit_co
get_wrolpi_config, \
timer, chunks, unique_by_predicate, get_paths_in_media_directory, TRACE_LEVEL, get_relative_to_media_directory
from wrolpi.dates import now, from_timestamp, months_selector_to_where, date_range_to_where
from wrolpi.db import get_db_session, get_db_curs, mogrify, optional_session
from wrolpi.db import get_db_session, get_db_curs, mogrify
from wrolpi.downloader import download_manager, Download
from wrolpi.errors import InvalidFile, UnknownDirectory, UnknownFile, UnknownTag, FileConflict, FileGroupIsTagged, \
NoPrimaryFile, InvalidDirectory, IgnoredDirectoryError
@ -55,8 +55,7 @@ __all__ = ['list_directories_contents', 'delete', 'split_path_stem_and_suffix',
'move', 'rename', 'delete_directory', 'handle_file_group_search_results', 'get_file_location_href']
@optional_session
def get_file_tag_names(file: pathlib.Path, session: Session = None) -> List[str]:
def get_file_tag_names(session: Session, file: pathlib.Path) -> List[str]:
"""Returns all Tag names for the provided file path."""
tags = session.query(Tag) \
.join(TagFile, Tag.id == TagFile.tag_id) \
@ -66,7 +65,7 @@ def get_file_tag_names(file: pathlib.Path, session: Session = None) -> List[str]
return names
def _get_file_dict(file: pathlib.Path) -> Dict:
def _get_file_dict(session: Session, file: pathlib.Path) -> Dict:
media_directory = get_media_directory()
try:
size = file.stat().st_size
@ -80,20 +79,19 @@ def _get_file_dict(file: pathlib.Path) -> Dict:
path=file.relative_to(media_directory),
size=size,
mimetype=mimetype,
tags=get_file_tag_names(file),
tags=get_file_tag_names(session, file),
)
def get_file_dict(file: str) -> Dict:
def get_file_dict(session: Session, file: str) -> Dict:
media_directory = get_media_directory()
return _get_file_dict(media_directory / file)
return _get_file_dict(session, media_directory / file)
@optional_session
async def set_file_viewed(file: pathlib.Path, session: Session = None):
async def set_file_viewed(session: Session, file: pathlib.Path):
"""Change FileGroup.viewed to the current datetime."""
try:
fg = FileGroup.find_by_path(file, session)
fg = FileGroup.find_by_path(session, file)
except UnknownFile:
fg = FileGroup.from_paths(session, file)
fg.do_model(session)
@ -116,7 +114,7 @@ def _get_directory_dict(directory: pathlib.Path,
)
def _get_recursive_directory_dict(directory: pathlib.Path, directories: List[pathlib.Path]) -> Dict:
def _get_recursive_directory_dict(session: Session, directory: pathlib.Path, directories: List[pathlib.Path]) -> Dict:
directories_cache = str(sorted(directories))
d = _get_directory_dict(directory, directories_cache)
@ -124,11 +122,11 @@ def _get_recursive_directory_dict(directory: pathlib.Path, directories: List[pat
children = dict()
for path in directory.iterdir():
if path.is_dir() and path in directories:
children[f'{path.name}/'] = _get_recursive_directory_dict(path, directories)
children[f'{path.name}/'] = _get_recursive_directory_dict(session, path, directories)
elif path.is_dir():
children[f'{path.name}/'] = _get_directory_dict(path, directories_cache)
else:
children[path.name] = _get_file_dict(path)
children[path.name] = _get_file_dict(session, path)
d['children'] = children
return d
@ -136,7 +134,7 @@ def _get_recursive_directory_dict(directory: pathlib.Path, directories: List[pat
HIDDEN_DIRECTORIES = ('lost+found',)
def list_directories_contents(directories_: List[str]) -> Dict:
def list_directories_contents(session: Session, directories_: List[str]) -> Dict:
"""List all files down to (and within) the directories provided.
This should only be used by the frontend for it's FileBrowser."""
@ -157,9 +155,9 @@ def list_directories_contents(directories_: List[str]) -> Dict:
# Never show ignored directories.
continue
if path.is_dir():
paths[f'{path.name}/'] = _get_recursive_directory_dict(path, directories)
paths[f'{path.name}/'] = _get_recursive_directory_dict(session, path, directories)
else:
paths[path.name] = _get_file_dict(path)
paths[path.name] = _get_file_dict(session, path)
return paths
@ -772,8 +770,7 @@ def upsert_directories(parent_directories, directories):
curs.execute(stmt, dict(directory=f'{directory.absolute()}%', idempotency=idempotency))
@optional_session
async def search_directories_by_name(name: str, excluded: List[str] = None, limit: int = 20, session: Session = None) \
async def search_directories_by_name(session: Session, name: str, excluded: List[str] = None, limit: int = 20) \
-> List[Directory]:
"""Find the Directories whose names contain the `name` string."""
excluded = excluded or []
@ -1050,8 +1047,8 @@ def group_files_by_stem(files: List[pathlib.Path], pre_sorted: bool = False) ->
yield group
async def _get_tag_and_file_group(file_group_id: int, file_group_primary_path: str, tag_name: str, tag_id: int,
session: Session):
async def _get_tag_and_file_group(session: Session, file_group_id: int, file_group_primary_path: str, tag_name: str,
tag_id: int):
if file_group_id:
file_group: FileGroup = session.query(FileGroup).filter_by(id=file_group_id).one_or_none()
if not file_group:
@ -1059,7 +1056,7 @@ async def _get_tag_and_file_group(file_group_id: int, file_group_primary_path: s
elif file_group_primary_path:
path = get_media_directory() / file_group_primary_path
file_group = FileGroup.get_by_path(path, session=session)
file_group = FileGroup.get_by_path(session, path)
# Create the FileGroup, if possible.
if not file_group and path.is_file():
file_group = FileGroup.from_paths(session, path)
@ -1074,7 +1071,7 @@ async def _get_tag_and_file_group(file_group_id: int, file_group_primary_path: s
if not tag:
raise UnknownTag(f'Cannot find Tag with id {tag_id}')
elif tag_name:
tag: Tag = Tag.get_by_name(tag_name, session)
tag: Tag = Tag.get_by_name(session, tag_name)
if not tag:
raise UnknownTag(f'Cannot find Tag with name {tag_name}')
else:
@ -1083,19 +1080,17 @@ async def _get_tag_and_file_group(file_group_id: int, file_group_primary_path: s
return file_group, tag
@optional_session(commit=True)
async def add_file_group_tag(file_group_id: int, file_group_primary_path: str, tag_name: str, tag_id: int,
session: Session = None) -> TagFile:
file_group, tag = await _get_tag_and_file_group(file_group_id, file_group_primary_path, tag_name, tag_id, session)
tag_file = file_group.add_tag(tag.id, session)
async def add_file_group_tag(session: Session, file_group_id: int, file_group_primary_path: str, tag_name: str,
tag_id: int) -> TagFile:
file_group, tag = await _get_tag_and_file_group(session, file_group_id, file_group_primary_path, tag_name, tag_id)
tag_file = file_group.add_tag(session, tag.id)
return tag_file
@optional_session(commit=True)
async def remove_file_group_tag(file_group_id: int, file_group_primary_path: str, tag_name: str, tag_id: int,
session: Session = None):
file_group, tag = await _get_tag_and_file_group(file_group_id, file_group_primary_path, tag_name, tag_id, session)
file_group.untag(tag.id, session)
async def remove_file_group_tag(session: Session, file_group_id: int, file_group_primary_path: str, tag_name: str,
tag_id: int):
file_group, tag = await _get_tag_and_file_group(session, file_group_id, file_group_primary_path, tag_name, tag_id)
file_group.untag(session, tag.id)
@dataclasses.dataclass
@ -1241,7 +1236,7 @@ def add_files_to_refresh_count(count: int):
api_app.shared_ctx.refresh['counted_files'] = count
async def _move(destination: pathlib.Path, *sources: pathlib.Path, session: Session) \
async def _move(session: Session, destination: pathlib.Path, *sources: pathlib.Path) \
-> OrderedDict[pathlib.Path, pathlib.Path]:
media_directory = get_media_directory()
for source in sources:
@ -1432,8 +1427,7 @@ def delete_directory(directory: pathlib.Path, recursive: bool = False):
curs.execute(stmt, (str(directory),))
@optional_session
async def move(destination: pathlib.Path, *sources: pathlib.Path, session: Session = None) \
async def move(session: Session, destination: pathlib.Path, *sources: pathlib.Path) \
-> OrderedDict[pathlib.Path, pathlib.Path]:
"""Moves a file or directory (recursively), preserving applied Tags.
@ -1441,8 +1435,6 @@ async def move(destination: pathlib.Path, *sources: pathlib.Path, session: Sessi
Returns an OrderedDict, the key is the file's old location, the value is where the file was moved.
"""
if not session:
raise RuntimeError('`session` is required')
with flags.refreshing:
from wrolpi.api_utils import api_app
@ -1450,7 +1442,7 @@ async def move(destination: pathlib.Path, *sources: pathlib.Path, session: Sessi
# Move the files, if this fails it will revert itself and raise an error.
with timer('moving files', 'info'):
plan = await _move(destination, *sources, session=session)
plan = await _move(session, destination, *sources)
session.commit()
with flags.refresh_indexing:
@ -1485,12 +1477,9 @@ async def rename_file(path: pathlib.Path, new_name: str) -> pathlib.Path:
return new_path
async def rename_directory(directory: pathlib.Path, new_name: str, session: Session = None, send_events: bool = False) \
async def rename_directory(session: Session, directory: pathlib.Path, new_name: str, send_events: bool = False) \
-> pathlib.Path:
"""Rename a directory. This is done by moving all files into the new directory, and removing the old directory."""
if not session:
raise RuntimeError('`session` is required')
new_directory = directory.with_name(new_name)
if new_directory.exists():
raise FileConflict(f'Cannot rename {directory} to {new_directory} because it already exists.')
@ -1498,7 +1487,7 @@ async def rename_directory(directory: pathlib.Path, new_name: str, session: Sess
# Move all paths into the new directory.
paths = list(directory.iterdir())
try:
await move(new_directory, *paths, session=session)
await move(session, new_directory, *paths)
if send_events:
Events.send_file_move_completed(f'Directory has been renamed: {directory}')
except Exception:
@ -1511,14 +1500,10 @@ async def rename_directory(directory: pathlib.Path, new_name: str, session: Sess
return new_directory
@optional_session
async def rename(path: pathlib.Path, new_name: str, session: Session = None, send_events: bool = False) -> pathlib.Path:
async def rename(session: Session, path: pathlib.Path, new_name: str, send_events: bool = False) -> pathlib.Path:
"""Rename a directory or file. Preserve any tags."""
if not session:
raise RuntimeError('`session` is required')
if path.is_dir():
return await rename_directory(path, new_name, session=session, send_events=send_events)
return await rename_directory(session, path, new_name, send_events=send_events)
return await rename_file(path, new_name)
@ -1616,7 +1601,7 @@ async def upsert_file(file: pathlib.Path | str, tag_names: List[str] = None) ->
if file_group.url:
if download_manager.is_skipped(file_group.url):
download_manager.remove_from_skip_list(file_group.url)
if download := Download.get_by_url(file_group.url, session):
if download := Download.get_by_url(session, file_group.url):
# Mark download as completed
download.complete()
# Link the download to the location of the uploaded file.
@ -1627,8 +1612,8 @@ async def upsert_file(file: pathlib.Path | str, tag_names: List[str] = None) ->
if tag_names:
for tag_name in tag_names:
if tag_name not in file_group.tag_names:
tag = Tag.get_by_name(tag_name, session)
file_group.add_tag(tag.id, session=session)
tag = Tag.get_by_name(session, tag_name)
file_group.add_tag(session, tag.id)
# Commit all changes in one transaction
session.commit()
@ -1905,16 +1890,17 @@ async def _process_bulk_tag_job(job: dict):
# Pre-resolve tag IDs (using cached lookups)
add_tag_ids = set()
remove_tag_ids = set()
for tag_name in add_tag_names:
try:
add_tag_ids.add(Tag.get_id_by_name(tag_name))
except UnknownTag:
logger.warning(f'Unknown tag to add: {tag_name}')
for tag_name in remove_tag_names:
try:
remove_tag_ids.add(Tag.get_id_by_name(tag_name))
except UnknownTag:
logger.warning(f'Unknown tag to remove: {tag_name}')
with get_db_session() as session:
for tag_name in add_tag_names:
try:
add_tag_ids.add(Tag.get_id_by_name(session, tag_name))
except UnknownTag:
logger.warning(f'Unknown tag to add: {tag_name}')
for tag_name in remove_tag_names:
try:
remove_tag_ids.add(Tag.get_id_by_name(session, tag_name))
except UnknownTag:
logger.warning(f'Unknown tag to remove: {tag_name}')
if not add_tag_ids and not remove_tag_ids:
return
@ -1933,7 +1919,7 @@ async def _process_bulk_tag_job(job: dict):
for file_path in chunk:
try:
# Get or create FileGroup
file_group = FileGroup.get_by_path(file_path, session=session)
file_group = FileGroup.get_by_path(session, file_path)
if not file_group:
# Create FileGroup for unindexed file
paths_for_group = glob_shared_stem(file_path)

58
wrolpi/files/models.py
View file

@ -3,7 +3,6 @@ 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, event, Index
@ -13,7 +12,7 @@ from sqlalchemy.orm import deferred, relationship, Session
from wrolpi.common import Base, ModelHelper, tsvector, logger, recursive_map, get_media_directory, \
get_relative_to_media_directory, unique_by_predicate
from wrolpi.dates import TZDateTime, now, from_timestamp, strptime_ms, strftime
from wrolpi.db import optional_session, get_db_session
from wrolpi.db import get_db_session
from wrolpi.downloader import Download
from wrolpi.errors import FileGroupIsTagged, UnknownFile
from wrolpi.files import indexers
@ -181,12 +180,11 @@ class FileGroup(ModelHelper, Base):
self.viewed = viewed or now()
return self.viewed
@optional_session
def set_tags(self, tag_names_or_ids: Iterable[str | int], session: Session = None):
def set_tags(self, session: Session, tag_names_or_ids: Iterable[str | int]):
"""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}
tag_ids = {Tag.get_id_by_name(session, i) for i in tag_names_or_ids}
else:
tag_ids = set(tag_names_or_ids)
@ -202,6 +200,8 @@ class FileGroup(ModelHelper, Base):
TagFile.file_group_id == self.id,
TagFile.tag_id.in_(deleted_tags),
).delete(synchronize_session=False)
# Refresh the relationship after deleting
session.expire(self, ['tag_files'])
# Save changes to config.
save_tags_config.activate_switch()
@ -209,29 +209,21 @@ class FileGroup(ModelHelper, Base):
return self.tag_files
@singledispatchmethod
@optional_session
def add_tag(self, tag_id: int | str, session: Session = None) -> TagFile:
def add_tag(self, session: Session, tag_id: int | str) -> TagFile:
"""Add a tag to this FileGroup. `tag_id` can be an int (tag ID) or str (tag name)."""
if isinstance(tag_id, str):
tag_id = Tag.get_id_by_name(session, tag_id)
tags = [i.tag_id for i in self.tag_files]
tags.append(tag_id)
tag_files = self.set_tags(tags, session=session)
tag_files = self.set_tags(session, tags)
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):
def untag(self, session: Session, tag_id: int | str):
"""Remove a tag from this FileGroup. `tag_id` can be an int (tag ID) or str (tag name)."""
if isinstance(tag_id, str):
tag_id = Tag.get_id_by_name(session, tag_id)
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)
self.set_tags(session, tags)
def append_files(self, *paths: pathlib.Path):
"""Add all `paths` to this FileGroup.files."""
@ -315,7 +307,7 @@ class FileGroup(ModelHelper, Base):
if session := Session.object_session(self):
from wrolpi.downloader import download_manager
if self.url and (download := Download.get_by_url(self.url, session=session)):
if self.url and (download := Download.get_by_url(session, self.url)):
download.delete(add_to_skip_list)
elif self.url and add_to_skip_list is True:
download_manager.add_to_skip_list(self.url)
@ -382,22 +374,20 @@ class FileGroup(ModelHelper, Base):
return file_group
@staticmethod
def find_by_id(id_: int, session: Session = None) -> 'FileGroup':
def find_by_id(session: Session, id_: int) -> 'FileGroup':
fg = session.query(FileGroup).filter(FileGroup.id == id_).one_or_none()
if not fg:
raise UnknownFile(f'Unable to find FileGroup with id {id_}')
return fg
@staticmethod
@optional_session
def get_by_path(path, session) -> Optional['FileGroup']:
def get_by_path(session: Session, path) -> Optional['FileGroup']:
file_group = session.query(FileGroup).filter(FileGroup.primary_path == str(path)).one_or_none()
return file_group
@staticmethod
@optional_session
def find_by_path(path, session) -> 'FileGroup':
file_group = FileGroup.get_by_path(path, session)
def find_by_path(session: Session, path) -> 'FileGroup':
file_group = FileGroup.get_by_path(session, path)
if not file_group:
raise UnknownFile(f'Unable to find FileGroup with path {path}')
return file_group
@ -435,7 +425,7 @@ class FileGroup(ModelHelper, Base):
for tag_file in file_group.tag_files:
if tag_file.tag is not None and tag_file.tag.name not in self.tag_names:
# Preserve the created at.
self.add_tag(tag_file.tag.id).created_at = tag_file.created_at
self.add_tag(session, tag_file.tag.id).created_at = tag_file.created_at
session.delete(file_group)
self.files = collected_files
@ -489,7 +479,7 @@ class FileGroup(ModelHelper, Base):
if klass:
# Return the model location when possible.
with get_db_session() as session:
instance = klass.get_by_path(self.primary_path, session)
instance = klass.get_by_path(session, self.primary_path)
if instance and hasattr(instance, 'location'):
return instance.location
@ -522,7 +512,7 @@ class FileGroup(ModelHelper, Base):
klass = self.get_model_class()
if klass:
with get_db_session() as session:
return klass.get_by_path(self.primary_path, session)
return klass.get_by_path(session, self.primary_path)
return None
def do_model(self, session: Session) -> Optional[ModelHelper]:
@ -530,7 +520,7 @@ class FileGroup(ModelHelper, Base):
if model_class := self.get_model_class():
# Always call do_model() to ensure FileGroup.data is updated when files change.
# The model's do_model() should be idempotent (safe to call multiple times).
model = model_class.do_model(self, session)
model = model_class.do_model(session, self)
logger.debug(f'modeled: {self} {model}')
return model

18
wrolpi/files/test/test_api.py
View file

@ -243,8 +243,8 @@ async def test_files_search_any_tag(async_client, test_session, make_files_struc
assert bar.primary_path.name == 'bar.txt' \
and foo.primary_path.name == 'foo.txt' \
and foobar.primary_path.name == 'foo bar.txt'
foo.add_tag(one.id)
foobar.add_tag(two.id)
foo.add_tag(test_session, one.id)
foobar.add_tag(test_session, two.id)
test_session.commit()
# Only `foo` is tagged with `one`
@ -482,8 +482,8 @@ async def test_post_search_directories(test_session, async_client, make_files_st
from modules.videos.models import Channel
# Use from_config to create channels properly (creates Collection first)
channel1 = Channel.from_config({'directory': channel1_dir, 'name': 'Channel Name'}, session=test_session)
channel2 = Channel.from_config({'directory': channel2_dir, 'name': 'OtherChannel'}, session=test_session)
channel1 = Channel.from_config(test_session, {'directory': channel1_dir, 'name': 'Channel Name'})
channel2 = Channel.from_config(test_session, {'directory': channel2_dir, 'name': 'OtherChannel'})
from wrolpi.collections import Collection
domain_collection = Collection(directory=domain_dir, name='example.com', kind='domain')
test_session.add(domain_collection)
@ -635,7 +635,7 @@ async def test_post_upload(test_session, async_client, test_directory, make_file
assert file_group
# Remove Tag so file can be uploaded again.
file_group.set_tags([], test_session)
file_group.set_tags(test_session, [])
# An existing file can be overwritten.
forms = [
@ -987,10 +987,10 @@ async def test_bulk_tag_preview_api(test_session, async_client, make_files_struc
# Add shared and non-shared tags
tag1 = await tag_factory('shared_tag')
tag2 = await tag_factory('only_foo')
fg_foo.add_tag(tag1.id)
fg_foo.add_tag(tag2.id)
fg_bar.add_tag(tag1.id)
fg_baz.add_tag(tag1.id)
fg_foo.add_tag(test_session, tag1.id)
fg_foo.add_tag(test_session, tag2.id)
fg_bar.add_tag(test_session, tag1.id)
fg_baz.add_tag(test_session, tag1.id)
test_session.commit()
content = {'paths': ['foo.txt', 'bar.txt', 'baz.txt']}

4
wrolpi/files/test/test_downloader.py
View file

@ -19,7 +19,7 @@ async def test_file_downloader(test_session, make_files_structure, test_director
# Create download for the video file.
video_url = f'{url}/{video_file.name}'
test_download_manager.create_download(video_url, file_downloader.name, test_session,
test_download_manager.create_download(test_session, video_url, file_downloader.name,
destination=str(test_directory / 'downloads'))
test_session.commit()
@ -73,7 +73,7 @@ async def test_file_downloader_meta4(test_session, make_files_structure, test_di
assert response.status == HTTPStatus.OK, response.content
# Create download for the video file, this should automatically check for the .meta4 file and use it with aria2c.
test_download_manager.create_download(video_url, file_downloader.name, test_session,
test_download_manager.create_download(test_session, video_url, file_downloader.name,
destination=str(test_directory / 'downloads'))
test_session.commit()

4
wrolpi/files/test/test_ebooks.py
View file

@ -150,12 +150,12 @@ async def test_move_ebook(async_client, test_session, test_directory, example_ep
test_directory / 'ebook example.epub',
test_directory / 'ebook example.jpg',
]
ebook.file_group.add_tag(tag.id)
ebook.file_group.add_tag(test_session, tag.id)
test_session.commit()
new_directory = test_directory / 'new'
new_directory.mkdir()
await move(new_directory, ebook.file_group.primary_path)
await move(test_session, new_directory, ebook.file_group.primary_path)
ebook: EBook = test_session.query(EBook).one()
assert ebook.file_group.primary_path == test_directory / 'new/ebook example.epub'

42
wrolpi/files/test/test_lib.py
View file

@ -86,7 +86,7 @@ async def test_delete_file_names(async_client, test_session, make_files_structur
foo_fg = FileGroup.from_paths(test_session, foo)
foo1_fg = FileGroup.from_paths(test_session, foo1)
tag = await tag_factory()
foo1_fg.add_tag(tag.id)
foo1_fg.add_tag(test_session, tag.id)
test_session.commit()
assert foo.is_file()
assert foo1.is_file()
@ -114,7 +114,7 @@ async def test_delete_tagged(await_switches, test_session, make_files_structure,
await lib.refresh_files()
# Both files end up in a group.
bar = test_session.query(FileGroup).one()
bar.add_tag(tag.id)
bar.add_tag(test_session, tag.id)
await await_switches()
test_session.commit()
@ -381,7 +381,7 @@ async def test_file_group_tag(test_session, make_files_structure, test_directory
one = await tag_factory()
foo: FileGroup = test_session.query(FileGroup).one()
foo.add_tag(one.id)
foo.add_tag(test_session, one.id)
test_session.commit()
tag_file2: TagFile = test_session.query(TagFile).one()
@ -753,8 +753,8 @@ async def test_file_group_merge(async_client, test_session, test_directory, make
srt_group = FileGroup.from_paths(test_session, srt)
test_session.add_all([vid_group, srt_group])
test_session.flush([vid_group, srt_group])
vid_tag_file = vid_group.add_tag(one.name, test_session)
srt_tag_file = srt_group.add_tag(two.name, test_session)
vid_tag_file = vid_group.add_tag(test_session, one.name)
srt_tag_file = srt_group.add_tag(test_session, two.name)
test_session.flush([vid_tag_file, srt_tag_file])
tag_file_created_at = vid_tag_file.created_at
srt_file_created_at = srt_tag_file.created_at
@ -794,7 +794,7 @@ async def test_move(async_client, test_session, test_directory, make_files_struc
await lib.refresh_files()
# mv foo qux
plan = await lib.move(qux, foo, session=test_session)
plan = await lib.move(test_session, qux, foo)
plan = [(str(i.relative_to(test_directory)), str(j.relative_to(test_directory))) for i, j in plan.items()]
# The deepest files are moved first.
assert plan == [('foo/bar/baz/archive.html', 'qux/foo/bar/baz/archive.html'),
@ -823,7 +823,7 @@ async def test_move_files(async_client, test_session, test_directory, make_files
})
dest = test_directory / 'dest'
plan = await lib.move(dest, one, two, session=test_session)
plan = await lib.move(test_session, dest, one, two)
assert list(plan.items()) == [
(test_directory / 'foo/one.txt', test_directory / 'dest/one.txt'),
(test_directory / 'two.txt', test_directory / 'dest/two.txt'),
@ -855,7 +855,7 @@ async def test_move_deep_directory(async_client, test_session, test_directory, m
dest = test_directory / 'deep/dest'
# mv foo/bar foo/qux foo/quuz.mp4 foo/quuz.txt deep/dest
plan = await lib.move(dest, bar, qux, quuz_mp4, quuz_txt, session=test_session)
plan = await lib.move(test_session, dest, bar, qux, quuz_mp4, quuz_txt)
assert list(plan.items()) == [
(test_directory / 'foo/qux/quux.txt', test_directory / 'deep/dest/qux/quux.txt'),
(test_directory / 'foo/bar/baz.txt', test_directory / 'deep/dest/bar/baz.txt'),
@ -891,7 +891,7 @@ async def test_move_directory(async_client, test_session, test_directory, make_f
assert_directories({'foo'})
bar = test_directory / 'bar'
await lib.rename(test_directory / 'foo', 'bar', session=test_session)
await lib.rename(test_session, test_directory / 'foo', 'bar')
assert (bar / 'one.txt').read_text() == 'one'
assert_directories({'bar'})
@ -920,7 +920,7 @@ async def test_move_error(test_session, test_directory, make_files_structure, vi
return shutil.move(*args, **kwargs)
with mock.patch('wrolpi.files.lib.shutil.move', mock_shutil_move), pytest.raises(FileNotFoundError):
await lib.move(qux, foo, session=test_session)
await lib.move(test_session, qux, foo)
# The move failed, the files should be moved back.
# foo was not deleted.
assert foo.is_dir()
@ -966,14 +966,14 @@ async def test_move_tagged(async_client, test_session, test_directory, make_file
await lib.refresh_files()
bar_file_group, foo_file_group = test_session.query(FileGroup).order_by(FileGroup.primary_path)
bar_file_group.title = 'custom title' # Should not be overwritten.
bar_file_group.add_tag(tag.id)
bar_file_group.add_tag(test_session, tag.id)
test_session.commit()
qux = test_directory / 'qux'
qux.mkdir()
# Move both files into qux. The Tag should also be moved.
await lib.move(qux, bar, foo, session=test_session)
await lib.move(test_session, qux, bar, foo)
new_foo = qux / 'foo.txt'
new_bar = qux / 'bar.txt'
# Files were moved.
@ -994,7 +994,7 @@ async def test_move_tagged(async_client, test_session, test_directory, make_file
assert bar_file_group.title == 'custom title', 'Custom title should not have been overwritten.'
# Rename "bar.txt" to "baz.txt"
await lib.rename(new_bar, 'baz.txt', session=test_session)
await lib.rename(test_session, new_bar, 'baz.txt')
baz = new_bar.with_name('baz.txt')
assert baz.read_text() == 'bar'
@ -1154,7 +1154,7 @@ async def test_move_many_files(async_client, test_session, test_directory, make_
with timer('test_move_many_files'):
# mv foo bar
await lib.move(bar, foo, session=test_session)
await lib.move(test_session, bar, foo)
assert not foo.is_dir()
assert bar.is_dir()
@ -1193,10 +1193,10 @@ async def test_get_bulk_tag_preview_files(async_client, test_session, make_files
# Add shared and non-shared tags
tag1 = await tag_factory('shared')
tag2 = await tag_factory('only_foo')
fg_foo.add_tag(tag1.id)
fg_foo.add_tag(tag2.id)
fg_bar.add_tag(tag1.id)
fg_baz.add_tag(tag1.id)
fg_foo.add_tag(test_session, tag1.id)
fg_foo.add_tag(test_session, tag2.id)
fg_bar.add_tag(test_session, tag1.id)
fg_baz.add_tag(test_session, tag1.id)
test_session.commit()
# Preview for all three files
@ -1245,7 +1245,7 @@ async def test_get_bulk_tag_preview_multi_file_filegroup(async_client, test_sess
# Add a tag to this FileGroup
tag = await tag_factory('my_tag')
fg.add_tag(tag.id)
fg.add_tag(test_session, tag.id)
test_session.commit()
# Preview for the directory should find the FileGroup and its tag
@ -1323,8 +1323,8 @@ async def test_process_bulk_tag_job_remove_tags(async_client, test_session, make
test_session.commit()
tag = await tag_factory('to_remove')
fg_foo.add_tag(tag.id)
fg_bar.add_tag(tag.id)
fg_foo.add_tag(test_session, tag.id)
fg_bar.add_tag(test_session, tag.id)
test_session.commit()
job = {

29
wrolpi/files/test/test_models.py
View file

@ -1,9 +1,30 @@
import pytest
from wrolpi.errors import UnknownFile
from wrolpi.files import lib
from wrolpi.files.models import FileGroup
def test_file_group_find_by_id(test_session, make_files_structure):
"""FileGroup.find_by_id should accept session as first argument, id as second."""
make_files_structure({'test.txt': 'test contents'})
# Create a FileGroup from the file
from wrolpi.common import get_media_directory
fg = FileGroup.from_paths(test_session, get_media_directory() / 'test.txt')
test_session.add(fg)
test_session.flush()
# Session must be first argument (session-first pattern)
found = FileGroup.find_by_id(test_session, fg.id)
assert found is not None
assert found.id == fg.id
# Non-existent ID should raise UnknownFile
with pytest.raises(UnknownFile):
FileGroup.find_by_id(test_session, 99999)
@pytest.mark.asyncio
async def test_file_group_move(async_client, test_session, test_directory, video_factory):
"""Any new or renamed files are adopted during FileGroup.move."""
@ -33,12 +54,12 @@ async def test_add_tag_overload(async_client, test_session, make_files_structure
one, two = await tag_factory(), await tag_factory()
foo.add_tag(one.name)
foo.add_tag(test_session, one.name)
assert set(foo.tag_names) == {'one'}, 'Tag should have been found by name'
foo.add_tag(two.id)
foo.add_tag(test_session, two.id)
assert set(foo.tag_names) == {'one', 'two'}, 'Tag should have been found by ID.'
foo.untag(one.name)
foo.untag(test_session, one.name)
assert set(foo.tag_names) == {'two'}, 'Tag should have been found by name'
foo.untag(two.id)
foo.untag(test_session, two.id)
assert not foo.tag_names, 'Tag should have been found by ID.'

4
wrolpi/files/test/test_scrape_downloader.py
View file

@ -42,7 +42,7 @@ async def test_scrape_html_downloader(test_directory, test_session, test_downloa
'depth': 1,
'suffix': '.PDF', # Suffix case is ignored.
}
test_download_manager.create_download('https://example.com/dir', 'scrape_html', settings=settings,
test_download_manager.create_download(test_session, 'https://example.com/dir', 'scrape_html', settings=settings,
sub_downloader_name='file', destination=destination)
assert_download_urls(['https://example.com/dir', ])
@ -66,7 +66,7 @@ async def test_scrape_html_downloader_html(test_directory, test_session, test_do
test_download_manager.register_downloader(FileDownloader())
settings = {'suffix': '.html'}
test_download_manager.create_download('https://example.com/dir', 'scrape_html', settings=settings,
test_download_manager.create_download(test_session, 'https://example.com/dir', 'scrape_html', settings=settings,
sub_downloader_name='file', destination=test_directory)
assert_download_urls(['https://example.com/dir', ])

79
wrolpi/root_api.py
View file

@ -27,7 +27,6 @@ from wrolpi.common import logger, get_wrolpi_config, wrol_mode_enabled, get_medi
set_global_log_level, get_relative_to_media_directory, search_other_estimates
from wrolpi.config_api import config_bp
from wrolpi.dates import now
from wrolpi.db import get_db_session
from wrolpi.downloader import download_manager
from wrolpi.errors import WROLModeEnabled, HotspotError, HotspotPasswordTooShort, InvalidConfig, \
ValidationError
@ -274,18 +273,20 @@ def valid_regex(_: Request, body: schema.RegexRequest):
validate=True,
)
@wrol_mode_check
async def post_download(_: Request, body: schema.DownloadRequest):
async def post_download(request: Request, body: schema.DownloadRequest):
# Raises an InvalidDownload if the Downloader cannot be found.
download_manager.find_downloader_by_name(body.downloader)
session = request.ctx.session
kwargs = dict(downloader_name=body.downloader,
sub_downloader_name=body.sub_downloader, reset_attempts=True,
destination=body.destination, tag_names=body.tag_names,
settings=body.settings)
if body.frequency:
download_manager.recurring_download(body.urls[0], body.frequency, collection_id=body.collection_id, **kwargs)
download_manager.recurring_download(session, body.urls[0], body.frequency, collection_id=body.collection_id,
**kwargs)
else:
download_manager.create_downloads(body.urls, **kwargs)
download_manager.create_downloads(session, body.urls, **kwargs)
if download_manager.disabled.is_set() or download_manager.stopped.is_set():
# Downloads are disabled, warn the user.
Events.send_downloads_disabled('Download created. But, downloads are disabled.')
@ -300,26 +301,26 @@ async def post_download(_: Request, body: schema.DownloadRequest):
validate=True,
)
@wrol_mode_check
async def put_download(_: Request, download_id: int, body: schema.DownloadRequest):
async def put_download(request: Request, download_id: int, body: schema.DownloadRequest):
# Raises an InvalidDownload if the Downloader cannot be found.
download_manager.find_downloader_by_name(body.downloader)
if len(body.urls) != 1:
raise ValidationError('Only one URL can be specified when updating a Download')
with get_db_session(commit=True) as session:
download_manager.update_download(
id_=download_id,
url=body.urls[0],
downloader=body.downloader,
destination=body.destination,
frequency=body.frequency,
tag_names=body.tag_names,
sub_downloader=body.sub_downloader,
settings=body.settings,
collection_id=body.collection_id,
session=session,
)
session = request.ctx.session
download_manager.update_download(
session,
id_=download_id,
url=body.urls[0],
downloader=body.downloader,
destination=body.destination,
frequency=body.frequency,
tag_names=body.tag_names,
sub_downloader=body.sub_downloader,
settings=body.settings,
collection_id=body.collection_id,
)
if download_manager.disabled.is_set() or download_manager.stopped.is_set():
# Downloads are disabled, warn the user.
Events.send_downloads_disabled('Download created. But, downloads are disabled.')
@ -330,15 +331,17 @@ async def put_download(_: Request, download_id: int, body: schema.DownloadReques
@api_bp.delete('/download/<download_id:int>')
@openapi.description('Delete a download')
@wrol_mode_check
async def delete_download(_: Request, download_id: int):
deleted = download_manager.delete_download(download_id)
async def delete_download(request: Request, download_id: int):
session = request.ctx.session
deleted = download_manager.delete_download(session, download_id)
return response.empty(HTTPStatus.NO_CONTENT if deleted else HTTPStatus.NOT_FOUND)
@api_bp.post('/download/<download_id:int>/restart')
@openapi.description('Restart a download.')
async def restart_download(_: Request, download_id: int):
download_manager.restart_download(download_id)
async def restart_download(request: Request, download_id: int):
session = request.ctx.session
download_manager.restart_download(session, download_id)
return response.empty()
@ -382,15 +385,17 @@ async def enable_downloads(_: Request):
@api_bp.post('/download/clear_completed')
@openapi.description('Clear completed downloads')
async def clear_completed(_: Request):
download_manager.delete_completed()
async def clear_completed(request: Request):
session = request.ctx.session
download_manager.delete_completed(session)
return response.empty()
@api_bp.post('/download/clear_failed')
@openapi.description('Clear failed downloads')
async def clear_failed(_: Request):
download_manager.delete_failed()
async def clear_failed(request: Request):
session = request.ctx.session
download_manager.delete_failed(session)
return response.empty()
@ -403,8 +408,9 @@ async def retry_once(_: Request):
@api_bp.post('/download/delete_once')
@openapi.description('Delete all once downloads')
async def delete_once(_: Request):
download_manager.delete_once()
async def delete_once(request: Request):
session = request.ctx.session
download_manager.delete_once(session)
return response.empty()
@ -541,8 +547,9 @@ async def get_tags_request(_: Request):
body=schema.TagRequest,
validate=True,
)
async def post_tag(_: Request, body: schema.TagRequest, tag_id: int = None):
await tags.upsert_tag(body.name, body.color, tag_id)
async def post_tag(request: Request, body: schema.TagRequest, tag_id: int = None):
session = request.ctx.session
await tags.upsert_tag(session, body.name, body.color, tag_id)
if tag_id:
return response.empty(HTTPStatus.OK)
else:
@ -550,8 +557,9 @@ async def post_tag(_: Request, body: schema.TagRequest, tag_id: int = None):
@api_bp.delete('/tag/<tag_id:int>')
async def delete_tag_request(_: Request, tag_id: int):
Tag.find_by_id(tag_id).delete()
async def delete_tag_request(request: Request, tag_id: int):
session = request.ctx.session
Tag.find_by_id(session, tag_id).delete()
return response.empty()
@ -672,13 +680,14 @@ async def post_upgrade_start(_: Request):
body=schema.SearchSuggestionsRequest,
validate=True,
)
async def post_search_suggestions(_: Request, body: schema.SearchSuggestionsRequest):
async def post_search_suggestions(request: Request, body: schema.SearchSuggestionsRequest):
from modules.videos.channel.lib import search_channels_by_name
from modules.archive.lib import search_domains_by_name
session = request.ctx.session
channels, domains = await asyncio.gather(
search_channels_by_name(body.search_str, order_by_video_count=body.order_by_video_count),
search_domains_by_name(body.search_str),
search_channels_by_name(session, body.search_str, order_by_video_count=body.order_by_video_count),
search_domains_by_name(session, body.search_str),
)
ret = dict(

47
wrolpi/tags.py
View file

@ -14,7 +14,7 @@ from wrolpi import dates, flags
from wrolpi.common import ModelHelper, Base, logger, ConfigFile, get_media_directory, background_task, \
get_relative_to_media_directory, is_valid_hex_color, walk, INVALID_FILE_CHARS
from wrolpi.dates import TZDateTime
from wrolpi.db import optional_session, get_db_curs, get_db_session
from wrolpi.db import get_db_curs, get_db_session
from wrolpi.downloader import save_downloads_config
from wrolpi.errors import UnknownTag, UsedTag, InvalidTag, RefreshConflict, NoPrimaryFile
from wrolpi.events import Events
@ -44,20 +44,18 @@ class TagFile(ModelHelper, Base):
return f'<TagFile tag={self.tag.name} file_group={self.file_group.primary_path}>'
@staticmethod
@optional_session
def get_by_primary_keys(file_group_id: int, tag_id: int, session: Session = None) -> Optional['TagFile']:
def get_by_primary_keys(session: Session, file_group_id: int, tag_id: int) -> Optional['TagFile']:
return session.query(TagFile).filter_by(tag_id=tag_id, file_group_id=file_group_id).one_or_none()
@staticmethod
@optional_session
def get_by_tag_name(file_group_id: int, tag_name: str, session: Session = None) -> Optional['TagFile']:
def get_by_tag_name(session: Session, file_group_id: int, tag_name: str) -> Optional['TagFile']:
return session.query(TagFile).join(Tag) \
.filter(TagFile.file_group_id == file_group_id, Tag.name == tag_name) \
.one_or_none()
# Tag.get_id_by_name
def get_id_by_name_key(klass: Base, name: str, session: Session = None):
def get_id_by_name_key(klass: Base, session: Session, name: str):
return hashkey(name)
@ -65,7 +63,7 @@ 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):
def get_name_by_id_key(klass: Base, session: Session, id_: int):
return hashkey(id_)
@ -97,8 +95,7 @@ class Tag(ModelHelper, Base):
@classmethod
@cachetools.cached(cache=get_id_by_name_cache, key=get_id_by_name_key)
@optional_session
def get_id_by_name(cls, name: str, session: Session = None) -> int:
def get_id_by_name(cls, session: Session, name: str) -> int:
"""Returns a Tag's ID, if a Tag matches the provided name.
@raise UnknownTag: If no Tag matches the provided name."""
@ -108,8 +105,7 @@ class Tag(ModelHelper, Base):
@classmethod
@cachetools.cached(cache=get_name_by_id_cache, key=get_name_by_id_key)
@optional_session
def get_name_by_id(cls, id_: int, session: Session = None) -> str:
def get_name_by_id(cls, session: Session, id_: int) -> str:
"""Returns a Tag's name, if a Tag matches the provided ID.
@raise UnknownTag: If no Tag matches the provided ID."""
@ -123,29 +119,25 @@ class Tag(ModelHelper, Base):
get_name_by_id_cache.clear()
@staticmethod
@optional_session
def get_by_name(name: str, session: Session) -> Optional['Tag']:
def get_by_name(session: Session, name: str) -> Optional['Tag']:
tag = session.query(Tag).filter_by(name=name).one_or_none()
return tag
@staticmethod
@optional_session
def find_by_name(name: str, session: Session) -> 'Tag':
if tag := Tag.get_by_name(name, session):
def find_by_name(session: Session, name: str) -> 'Tag':
if tag := Tag.get_by_name(session, name):
return tag
raise UnknownTag(f'No Tag with name={name}')
@staticmethod
@optional_session
def get_by_id(id_: int, session: Session = None) -> Optional['Tag']:
def get_by_id(session: Session, id_: int) -> Optional['Tag']:
tag = session.query(Tag).filter_by(id=id_).one_or_none()
return tag
@staticmethod
@optional_session
def find_by_id(id_: int, session: Session = None) -> 'Tag':
if tag := Tag.get_by_id(id_, session):
def find_by_id(session: Session, id_: int) -> 'Tag':
if tag := Tag.get_by_id(session, id_):
return tag
raise UnknownTag(f'No Tag with id={id_}')
@ -158,7 +150,7 @@ class Tag(ModelHelper, Base):
has_collections = session.query(Collection).filter(Collection.tag_id == self.id).first() is not None
return bool(any(self.tag_files) or any(self.tag_zim_entries) or has_collections)
async def update_tag(self, name: str, color: str | None, session: Session):
async def update_tag(self, session: Session, name: str, color: str | None):
"""Change name/color of tag. Ensures safeness of new values, checks for conflicts."""
if ',' in name:
raise InvalidTag('Tag name cannot have comma')
@ -169,7 +161,7 @@ class Tag(ModelHelper, Base):
name = name.replace('/', '⧸')
name = INVALID_FILE_CHARS.sub('', name)
if (other := self.get_by_name(name, session)) and other.id != self.id:
if (other := self.get_by_name(session, name)) and other.id != self.id:
raise InvalidTag('Tag name already taken')
if flags.refreshing.is_set():
@ -667,14 +659,13 @@ def get_tags() -> List[dict]:
return tags
@optional_session
async def upsert_tag(name: str, color: str, tag_id: int = None, session: Session = None) -> Tag:
async def upsert_tag(session: Session, name: str, color: str, tag_id: int = None) -> Tag:
if tag_id:
tag = Tag.find_by_id(tag_id, session)
await tag.update_tag(name, color, session)
tag = Tag.find_by_id(session, tag_id)
await tag.update_tag(session, name, color)
else:
tag = Tag()
await tag.update_tag(name, color, session)
await tag.update_tag(session, name, color)
session.add(tag)
tag.flush()

3
wrolpi/test/test_common.py
View file

@ -1073,7 +1073,8 @@ async def test_initialize_config_files(async_client, test_session, test_director
# Create some DB entries so the configs will not be imported.
tag1, tag2 = await tag_factory(), await tag_factory()
download = test_download_manager.create_download('https://example.com', downloader_name=test_downloader.name)
download = test_download_manager.create_download(test_session, 'https://example.com',
downloader_name=test_downloader.name)
test_session.commit()
# Write some config files, they should not be re-created.
get_wrolpi_config().get_file().write_text('wrolpi')

88
wrolpi/test/test_db.py
View file

@ -1,87 +1 @@
import mock
import pytest
from sqlalchemy.orm import Session
from wrolpi.db import optional_session
from wrolpi.flags import WROLPiFlag
@pytest.mark.asyncio
async def test_optional_session(test_session):
"""
`optional_session` wraps a function and provides a session if one was not passed.
"""
url = test_session.get_bind().url
called = mock.Mock()
def check_call(func):
# Function can be called with various arguments.
func()
func(session=test_session)
func(test_session)
# Function actually ran.
called.assert_called()
@optional_session()
def sync_with_wrapper(session: Session):
assert session.get_bind().url == url
called()
check_call(sync_with_wrapper)
# Commit can be passed.
@optional_session(commit=True)
def sync_with_commit(session: Session):
assert session.get_bind().url == url
called()
check_call(sync_with_commit)
# Wrapper is not called.
@optional_session
def sync_without_wrapper(session: Session):
assert session.get_bind().url == url
called()
check_call(sync_without_wrapper)
async def check_async_call(func):
# Function can be called with various arguments.
await func()
await func(session=test_session)
await func(test_session)
# Function actually ran.
called.assert_called()
# Async is called.
@optional_session
async def async_without_wrapper(session: Session):
assert session.get_bind().url == url
called()
await check_async_call(async_without_wrapper)
# Async is called with wrapper.
@optional_session(commit=True)
async def async_with_wrapper(session: Session):
assert session.get_bind().url == url
called()
await check_async_call(async_with_wrapper)
@pytest.mark.asyncio
async def test_optional_session_commit(test_session):
"""It is possible to request that a Commit happens after your wrapped function completes."""
assert test_session.query(WROLPiFlag).count() == 0
@optional_session(commit=True)
async def func(session: Session = None):
flag = WROLPiFlag(refresh_complete=True)
session.add(flag)
# No `session.commit()` here.
await func()
assert test_session.query(WROLPiFlag).one().refresh_complete is True
# Placeholder test file - optional_session tests removed since decorator is deprecated.

68
wrolpi/test/test_downloader.py
View file

@ -27,23 +27,23 @@ async def test_delete_old_once_downloads(test_session, test_download_manager, te
mock_now.return_value = datetime(2020, 6, 5, 0, 0, tzinfo=pytz.UTC)
# Recurring downloads should not be deleted.
d1 = test_download_manager.create_download('https://example.com/1', test_downloader.name)
d2 = test_download_manager.create_download('https://example.com/2', test_downloader.name)
d1 = test_download_manager.create_download(test_session, 'https://example.com/1', test_downloader.name)
d2 = test_download_manager.create_download(test_session, 'https://example.com/2', test_downloader.name)
d1.frequency = 1
d2.frequency = 1
d2.started()
# Should be deleted.
d3 = test_download_manager.create_download('https://example.com/3', test_downloader.name)
d4 = test_download_manager.create_download('https://example.com/4', test_downloader.name)
d3 = test_download_manager.create_download(test_session, 'https://example.com/3', test_downloader.name)
d4 = test_download_manager.create_download(test_session, 'https://example.com/4', test_downloader.name)
d3.complete()
d4.complete()
d3.last_successful_download = datetime(2020, 1, 1, 0, 0, 0, tzinfo=pytz.UTC)
d4.last_successful_download = datetime(2020, 5, 1, 0, 0, 0, tzinfo=pytz.UTC)
# Not a month old.
d5 = test_download_manager.create_download('https://example.com/5', test_downloader.name)
d5 = test_download_manager.create_download(test_session, 'https://example.com/5', test_downloader.name)
d5.last_successful_download = datetime(2020, 6, 1, 0, 0, 0, tzinfo=pytz.UTC)
# An old, but pending download should not be deleted.
d6 = test_download_manager.create_download('https://example.com/6', test_downloader.name)
d6 = test_download_manager.create_download(test_session, 'https://example.com/6', test_downloader.name)
d6.last_successful_download = datetime(2020, 4, 1, 0, 0, 0, tzinfo=pytz.UTC)
d6.started()
@ -68,17 +68,20 @@ async def test_recreate_download(test_session, test_download_manager, test_downl
settings = {'tag_names': [tag.name, ]}
# Create download with settings initially.
test_download_manager.create_downloads(['https://example.com/1', ], test_downloader.name, settings=settings)
test_download_manager.create_downloads(test_session, ['https://example.com/1', ], test_downloader.name,
settings=settings)
download: Download = test_session.query(Download).one()
assert download.settings == settings
# Restarting download does not remove settings.
test_download_manager.create_downloads(['https://example.com/1', ], test_downloader.name, settings=None)
test_download_manager.create_downloads(test_session, ['https://example.com/1', ], test_downloader.name,
settings=None)
download: Download = test_session.query(Download).one()
assert download.settings == settings
# Settings can be overwritten when not None.
test_download_manager.create_downloads(['https://example.com/1', ], test_downloader.name, settings=dict())
test_download_manager.create_downloads(test_session, ['https://example.com/1', ], test_downloader.name,
settings=dict())
download: Download = test_session.query(Download).one()
assert download.settings == dict()
@ -90,12 +93,12 @@ async def test_download_rename_tag(async_client, test_session, test_download_man
tag_names = [tag.name, ]
test_download_manager.create_downloads(
['https://example.com/1', ], test_downloader.name, tag_names=tag_names)
test_session, ['https://example.com/1', ], test_downloader.name, tag_names=tag_names)
download = test_session.query(Download).one()
assert download.tag_names == ['one', ]
await tag.update_tag('new name', tag.color, test_session)
await tag.update_tag(test_session, 'new name', tag.color)
download = test_session.query(Download).one()
assert download.tag_names == ['new name', ]
@ -134,11 +137,11 @@ async def test_download_get_downloader(test_session, test_download_manager, test
"""
A Download can find it's Downloader.
"""
download1 = test_download_manager.create_download('https://example.com', test_downloader.name)
download1 = test_download_manager.create_download(test_session, 'https://example.com', test_downloader.name)
assert download1.get_downloader() == test_downloader
with pytest.raises(InvalidDownload):
download2 = test_download_manager.create_download('https://example.com', 'bad downloader')
download2 = test_download_manager.create_download(test_session, 'https://example.com', 'bad downloader')
@pytest.mark.asyncio
@ -160,7 +163,7 @@ async def test_calculate_next_download(test_session, test_download_manager, fake
]
for attempts, expected in attempts_expected:
download.attempts = attempts
result = test_download_manager.calculate_next_download(download)
result = test_download_manager.calculate_next_download(test_session, download)
assert result == expected, f'{attempts} != {result}'
d1 = Download(url='https://example.com/1', frequency=DownloadFrequency.weekly)
@ -169,9 +172,9 @@ async def test_calculate_next_download(test_session, test_download_manager, fake
test_session.add_all([d1, d2, d3])
test_session.commit()
# Downloads are ead out over the next week.
assert test_download_manager.calculate_next_download(d1) == datetime(2000, 1, 8, tzinfo=pytz.UTC)
assert test_download_manager.calculate_next_download(d2) == datetime(2000, 1, 11, 12, tzinfo=pytz.UTC)
assert test_download_manager.calculate_next_download(d3) == datetime(2000, 1, 9, 18, tzinfo=pytz.UTC)
assert test_download_manager.calculate_next_download(test_session, d1) == datetime(2000, 1, 8, tzinfo=pytz.UTC)
assert test_download_manager.calculate_next_download(test_session, d2) == datetime(2000, 1, 11, 12, tzinfo=pytz.UTC)
assert test_download_manager.calculate_next_download(test_session, d3) == datetime(2000, 1, 9, 18, tzinfo=pytz.UTC)
@pytest.mark.asyncio
@ -180,7 +183,8 @@ async def test_recurring_downloads(test_session, test_download_manager, fake_now
test_downloader.set_test_success()
# Download every hour.
test_download_manager.recurring_download('https://example.com/recurring', Seconds.hour, test_downloader.name)
test_download_manager.recurring_download(test_session, 'https://example.com/recurring', Seconds.hour,
test_downloader.name)
# One download is scheduled.
downloads = test_download_manager.get_new_downloads(test_session)
@ -249,19 +253,19 @@ async def test_max_attempts(test_session, test_download_manager, test_downloader
test_downloader.set_test_success()
test_download_manager.create_download('https://example.com', test_downloader.name)
test_download_manager.create_download(test_session, 'https://example.com', test_downloader.name)
await test_download_manager.wait_for_all_downloads()
download = session.query(Download).one()
assert download.attempts == 1
test_download_manager.create_download('https://example.com', test_downloader.name)
test_download_manager.create_download(test_session, 'https://example.com', test_downloader.name)
await test_download_manager.wait_for_all_downloads()
download = session.query(Download).one()
assert download.attempts == 2
# There are no further attempts.
test_downloader.set_test_unrecoverable_exception()
test_download_manager.create_download('https://example.com', test_downloader.name)
test_download_manager.create_download(test_session, 'https://example.com', test_downloader.name)
await test_download_manager.wait_for_all_downloads()
download = session.query(Download).one()
assert download.attempts == 3
@ -277,7 +281,7 @@ async def test_skip_urls(test_session, test_download_manager, assert_download_ur
test_downloader.set_test_success()
test_download_manager.create_downloads([
test_download_manager.create_downloads(test_session, [
'https://example.com/1',
'https://example.com/skipme',
'https://example.com/2',
@ -286,10 +290,11 @@ async def test_skip_urls(test_session, test_download_manager, assert_download_ur
assert_download_urls({'https://example.com/1', 'https://example.com/2'})
assert get_download_manager_config().skip_urls == ['https://example.com/skipme']
test_download_manager.delete_once()
test_download_manager.delete_once(test_session)
# The user can start a download even if the URL is in the skip list.
test_download_manager.create_download('https://example.com/skipme', test_downloader.name, reset_attempts=True)
test_download_manager.create_download(test_session, 'https://example.com/skipme', test_downloader.name,
reset_attempts=True)
assert_download_urls({'https://example.com/skipme'})
assert get_download_manager_config().skip_urls == []
@ -410,18 +415,18 @@ async def test_crud_download(async_client, test_session, test_download_manager,
async def test_downloads_config(test_session, test_download_manager, test_download_manager_config,
test_downloader, assert_downloads, tag_factory, await_switches):
# Can import with an empty config.
await import_downloads_config()
await import_downloads_config(test_session)
assert_downloads([])
test_downloader.set_test_success()
tag = await tag_factory()
download1, download2 = test_download_manager.create_downloads([
download1, download2 = test_download_manager.create_downloads(test_session, [
'https://example.com/1',
'https://example.com/2',
], downloader_name=test_downloader.name)
test_download_manager.recurring_download('https://example.com/3', frequency=DownloadFrequency.weekly,
test_download_manager.recurring_download(test_session, 'https://example.com/3', frequency=DownloadFrequency.weekly,
downloader_name=test_downloader.name)
download2.next_download = datetime(2000, 1, 2, 0, 0, 0, tzinfo=pytz.UTC)
download2.destination = 'some directory'
@ -474,7 +479,7 @@ async def test_downloads_config(test_session, test_download_manager, test_downlo
get_download_manager_config().initialize()
# Import the saved downloads.
await import_downloads_config()
await import_downloads_config(test_session)
expected = [
dict(
url='https://example.com/2',
@ -499,7 +504,7 @@ async def test_downloads_config(test_session, test_download_manager, test_downlo
assert_downloads(expected)
# Import again, no duplicates.
await import_downloads_config()
await import_downloads_config(test_session)
assert_downloads(expected)
@ -522,8 +527,8 @@ async def test_download_excluded_urls(test_session, test_download_manager, test_
]
)
settings = dict(excluded_urls='example.org,example.gov')
feed_download: Download = test_download_manager.create_download('https://example.com/feed', rss_downloader.name,
test_session,
feed_download: Download = test_download_manager.create_download(test_session, 'https://example.com/feed',
rss_downloader.name,
sub_downloader_name=test_downloader.name,
settings=settings)
await test_download_manager.wait_for_all_downloads()
@ -557,6 +562,7 @@ async def test_download_with_collection_id(test_session, test_download_manager,
# Creating a recurring download with collection_id should work
download = test_download_manager.recurring_download(
test_session,
'https://example.com/test',
DownloadFrequency.weekly,
test_downloader.name,

2
wrolpi/test/test_root_api.py
View file

@ -519,7 +519,7 @@ async def test_empty_post_download(async_client):
async def test_restart_download(test_session, async_client, test_download_manager, test_downloader, await_switches):
"""A Download can be restarted."""
# Create a download, fail it, it should be restarted.
download = test_download_manager.create_download('https://example.com', test_downloader.name)
download = test_download_manager.create_download(test_session, 'https://example.com', test_downloader.name)
download.fail()
test_session.commit()
await await_switches()

11
wrolpi/test/test_rss.py
View file

@ -3,8 +3,8 @@ from typing import List
import mock
import pytest
from sqlalchemy.orm import Session
from wrolpi.db import optional_session
from wrolpi.downloader import Download, DownloadResult, Downloader, DownloadManager, RSSDownloader
@ -14,8 +14,7 @@ class RSSHTTPDownloader(Downloader):
async def do_download(self, download: Download) -> DownloadResult:
return DownloadResult(success=True)
@optional_session
def already_downloaded(self, *urls: List[str], session=None):
def already_downloaded(self, session: Session, *urls: List[str]):
downloads = list(session.query(Download).filter(Download.url.in_(urls)))
return downloads
@ -46,7 +45,7 @@ async def test_rss_download(test_session, test_download_manager, await_switches)
]
)
test_download_manager: DownloadManager
test_download_manager.create_download('https://example.com/feed', rss_downloader.name,
test_download_manager.create_download(test_session, 'https://example.com/feed', rss_downloader.name,
sub_downloader_name='rss_http')
await test_download_manager.wait_for_all_downloads()
@ -66,7 +65,7 @@ async def test_rss_download(test_session, test_download_manager, await_switches)
dict(link='https://example.com/d'),
]
)
test_download_manager.create_download('https://example.com/feed', rss_downloader.name,
test_download_manager.create_download(test_session, 'https://example.com/feed', rss_downloader.name,
sub_downloader_name='rss_http')
await test_download_manager.wait_for_all_downloads()
@ -90,7 +89,7 @@ async def test_rss_no_entries(test_session, test_download_manager, await_switche
with mock.patch('wrolpi.downloader.parse_feed') as mock_parse_feed:
mock_parse_feed.return_value = dict(bozo=0, ) # missing `entries`
test_download_manager.create_download('https://example.com/feed', rss_downloader.name, test_session,
test_download_manager.create_download(test_session, 'https://example.com/feed', rss_downloader.name,
sub_downloader_name='rss_http')
await test_download_manager.wait_for_all_downloads()

80
wrolpi/test/test_tags.py
View file

@ -20,16 +20,16 @@ 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()
file_group.add_tag(tag_one.id)
file_group.add_tag(test_session, tag_one.id)
assert file_group.__json__()['tags'] == ['one']
file_group.add_tag(tag_two.name)
file_group.add_tag(test_session, tag_two.name)
assert file_group.__json__()['tags'] == ['one', 'two']
file_group.untag(tag_one.id)
file_group.untag(test_session, tag_one.id)
assert file_group.__json__()['tags'] == ['two']
file_group.untag(tag_two.name)
file_group.untag(test_session, tag_two.name)
assert file_group.__json__()['tags'] == []
@ -39,17 +39,17 @@ async def test_tags_model(async_client, test_session, make_files_structure, tag_
tag1 = await tag_factory()
assert tag1.__json__()
assert Tag.find_by_id(tag1.id)
assert Tag.find_by_name(tag1.name)
assert Tag.find_by_id(test_session, tag1.id)
assert Tag.find_by_name(test_session, tag1.name)
assert Tag.get_by_id(123) is None, '.get should return None with bad id'
assert Tag.get_by_name('bad name') is None, '.get should return None with bad name'
assert Tag.get_by_id(test_session, 123) is None, '.get should return None with bad id'
assert Tag.get_by_name(test_session, 'bad name') is None, '.get should return None with bad name'
with pytest.raises(UnknownTag):
Tag.find_by_id(123)
Tag.find_by_id(test_session, 123)
with pytest.raises(UnknownTag):
Tag.find_by_name('bad name')
Tag.find_by_name(test_session, 'bad name')
@pytest.mark.asyncio
@ -65,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
video_group.add_tag(tag_one.id)
video_group.add_tag(test_session, 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()
video_group.add_tag(tag_two.id)
video_group.add_tag(test_session, tag_two.id)
test_session.commit()
assert len(tag_one.tag_files) == 1
assert len(tag_two.tag_files) == 1
@ -92,8 +92,8 @@ async def test_tags_config_file(test_session, test_directory, tag_factory, examp
test_tags_config):
"""Test that the config is updated when a FileGroup is tagged."""
await files_lib.refresh_files()
pdf: FileGroup = FileGroup.get_by_path(example_pdf, test_session)
video: FileGroup = FileGroup.get_by_path(video_file, test_session)
pdf: FileGroup = FileGroup.get_by_path(test_session, example_pdf)
video: FileGroup = FileGroup.get_by_path(test_session, video_file)
tag1 = await tag_factory()
tag2 = await tag_factory()
test_session.commit()
@ -104,31 +104,31 @@ async def test_tags_config_file(test_session, test_directory, tag_factory, examp
assert tag1.name in test_tags_config.read_text()
assert tag2.name in test_tags_config.read_text()
video.add_tag(tag1.id)
video.add_tag(test_session, tag1.id)
await await_switches()
assert str(video_file.relative_to(test_directory)) in test_tags_config.read_text()
assert str(example_pdf.relative_to(test_directory)) not in test_tags_config.read_text()
# Tag PDF twice.
pdf.add_tag(tag1.id)
pdf.add_tag(tag2.id)
pdf.add_tag(test_session, tag1.id)
pdf.add_tag(test_session, tag2.id)
await await_switches()
assert str(video_file.relative_to(test_directory)) in test_tags_config.read_text()
assert str(example_pdf.relative_to(test_directory)) in test_tags_config.read_text()
# One TagFile still exists.
pdf.untag(tag1.id)
pdf.untag(test_session, tag1.id)
await await_switches()
assert str(video_file.relative_to(test_directory)) in test_tags_config.read_text()
assert str(example_pdf.relative_to(test_directory)) in test_tags_config.read_text()
# PDF is no longer tagged.
pdf.untag(tag2.id)
pdf.untag(test_session, tag2.id)
await await_switches()
assert str(video_file.relative_to(test_directory)) in test_tags_config.read_text()
assert str(example_pdf.relative_to(test_directory)) not in test_tags_config.read_text()
video.untag(tag1.id)
video.untag(test_session, tag1.id)
await await_switches()
assert str(video_file.relative_to(test_directory)) not in test_tags_config.read_text()
assert str(example_pdf.relative_to(test_directory)) not in test_tags_config.read_text()
@ -138,7 +138,7 @@ async def test_tags_config_file(test_session, test_directory, tag_factory, examp
assert isinstance(tags_config.tag_files, list) and len(tags_config.tag_files) == 0
# Removing non-existent tag does not error.
video.untag(tag1.id)
video.untag(test_session, tag1.id)
@pytest.mark.asyncio
@ -166,7 +166,7 @@ async def test_tags_crud(async_client, test_session, example_pdf, assert_tags_co
# Apply the tag to the PDF.
tag = test_session.query(tags.Tag).one()
pdf.add_tag(tag.name, test_session)
pdf.add_tag(test_session, tag.name)
test_session.commit()
# The tag can be updated.
@ -182,7 +182,7 @@ async def test_tags_crud(async_client, test_session, example_pdf, assert_tags_co
content = dict(name='ガーデン', color='#ffffff')
request, response = await async_client.post('/api/tag/1', content=json.dumps(content))
assert response.status_code == HTTPStatus.OK
assert Tag.find_by_name('ガーデン').color == '#ffffff'
assert Tag.find_by_name(test_session, 'ガーデン').color == '#ffffff'
# Conflicting names return an error.
content = dict(name='ガーデン', color='#111111')
@ -194,7 +194,7 @@ async def test_tags_crud(async_client, test_session, example_pdf, assert_tags_co
request, response = await async_client.delete('/api/tag/1')
assert response.status_code == HTTPStatus.BAD_REQUEST
pdf.untag(tag.id, test_session)
pdf.untag(test_session, tag.id)
test_session.commit()
# Can delete unused Tag.
@ -212,7 +212,7 @@ async def test_delete_tagged_file(test_session, example_pdf, tag_factory):
tag = await tag_factory()
pdf: FileGroup = test_session.query(FileGroup).one()
pdf.add_tag(tag.id)
pdf.add_tag(test_session, tag.id)
test_session.commit()
with pytest.raises(FileGroupIsTagged):
@ -294,9 +294,9 @@ async def test_import_tags_delete_missing(test_session, test_directory, test_tag
foo, = make_files_structure({'foo': 'text'})
foo = FileGroup.from_paths(test_session, foo)
test_session.flush([foo, ])
foo.add_tag(tag1.id)
foo.add_tag(test_session, tag1.id)
# use tag4 for Zim entry.
await zim_lib.add_tag(tag4.name, test_zim.id, 'home')
await zim_lib.add_tag(test_session, tag4.name, test_zim.id, 'home')
tags.save_tags_config.activate_switch()
await await_switches()
@ -339,13 +339,13 @@ async def test_import_tags_delete_missing(test_session, test_directory, test_tag
@pytest.mark.asyncio
async def test_invalid_tag(async_client, test_session, test_directory):
with pytest.raises(InvalidTag):
await tags.upsert_tag('cannot use comma: ,', '#fff')
await tags.upsert_tag(test_session, 'cannot use comma: ,', '#fff')
with pytest.raises(InvalidTag):
await tags.upsert_tag('invalid color', 'foo')
await tags.upsert_tag(test_session, 'invalid color', 'foo')
with pytest.raises(UnknownTag):
await tags.upsert_tag('Tag ID does not exist', '#fff', 1)
await tags.upsert_tag(test_session, 'Tag ID does not exist', '#fff', 1)
@pytest.mark.asyncio
@ -382,12 +382,12 @@ async def test_tags_directory(test_session, test_directory, tag_factory, video_f
assert readme.is_file() and readme.stat().st_size > 0
# Video1 tagged with First Aid and Special/name.
vid1.add_tag(tag1.id)
vid1.add_tag(tag2.id)
vid1.add_tag(test_session, tag1.id)
vid1.add_tag(test_session, tag2.id)
# Video2 tagged with First Aid.
vid2.add_tag(tag1.id)
vid2.add_tag(test_session, tag1.id)
# PDF tagged with Special/name.
pdf.add_tag(tag2.id)
pdf.add_tag(test_session, tag2.id)
await await_switches()
assert_tag_links([
'First Aid, Special⧸name/vid1.en.vtt',
@ -407,7 +407,7 @@ async def test_tags_directory(test_session, test_directory, tag_factory, video_f
assert test_session.query(TagFile).count() == 4
# Video1 tagged with First Aid only.
vid1.untag(tag2.id)
vid1.untag(test_session, tag2.id)
await await_switches()
assert_tag_links([
'First Aid/vid1.en.vtt',
@ -424,7 +424,7 @@ async def test_tags_directory(test_session, test_directory, tag_factory, video_f
assert test_session.query(TagFile).count() == 3
# PDF no longer tagged.
pdf.untag(tag2.id)
pdf.untag(test_session, tag2.id)
await await_switches()
assert_tag_links([
'First Aid/vid1.en.vtt',
@ -441,8 +441,8 @@ async def test_tags_directory(test_session, test_directory, tag_factory, video_f
assert test_session.query(TagFile).count() == 2
# No more tagged files.
vid1.untag(tag1.id)
vid2.untag(tag1.id)
vid1.untag(test_session, tag1.id)
vid2.untag(test_session, tag1.id)
await await_switches()
assert test_session.query(TagFile).count() == 0
# Readme, and un-deletable file and directory exist.
@ -471,7 +471,7 @@ async def test_update_tag(test_session, test_directory, video_factory, tag_facto
assert (test_directory / 'tags/one/video.mp4').is_file()
# Invalid characters are replaced during rename.
await tags.upsert_tag('new/name%', tag.color, tag.id)
await tags.upsert_tag(test_session, 'new/name%', tag.color, tag.id)
await await_switches()
# Video file was moved.
@ -494,7 +494,7 @@ async def test_tag_rename_with_channel(test_session, test_directory, video_facto
assert (test_directory / 'videos/one/Channel Name').is_dir()
assert (test_directory / 'videos/one/Channel Name/video.mp4').is_file()
await tag.update_tag('New Tag Name', None, session=test_session)
await tag.update_tag(test_session, 'New Tag Name', None)
assert (test_directory / 'tags').is_dir()
assert (test_directory / 'videos/New Tag Name/Channel Name').is_dir()