import pathlib from pathlib import Path from typing import List, Dict, Union from sqlalchemy import or_, func, desc, asc from sqlalchemy.orm import Session from wrolpi import flags 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.downloader import save_downloads_config, download_manager, Download from wrolpi.errors import APIError, ValidationError, RefreshConflict from wrolpi.vars import PYTEST from .. import schema from ..common import check_for_channel_conflicts from ..errors import UnknownChannel, ChannelDirectoryConflict from ..lib import save_channels_config from ..models import Channel, Video logger = logger.getChild(__name__) async def get_minimal_channels() -> List[dict]: """ Get the minimum amount of information necessary about all channels. """ with get_db_curs() as curs: # Get all channels, even if they don't have videos. Also get the minimum frequency download because this is the # one that will consume the most resources. stmt = ''' SELECT c.id, c.name AS "name", directory, c.url, t.name AS "tag_name", COUNT(v.id) as video_count, SUM(fg.size)::BIGINT AS size, ( select min(d.frequency) from download d where d.channel_id = c.id ) AS minimum_frequency FROM channel AS c LEFT JOIN video v on c.id = v.channel_id LEFT JOIN file_group fg on fg.id = v.file_group_id LEFT JOIN tag t ON t.id = c.tag_id GROUP BY 1, 2, 3, 4, 5 ''' curs.execute(stmt) channels = sorted([dict(i) for i in curs.fetchall()], key=lambda i: i['name'].lower()) return channels async def get_channels_video_count() -> Dict[int, int]: """ Add video counts to all channels """ with get_db_curs() as curs: stmt = 'SELECT id FROM channel' curs.execute(stmt) # Get all channel IDs, start them with a count of 0. video_counts = {int(i['id']): 0 for i in curs.fetchall()} stmt = ''' SELECT c.id, COUNT(v.id) AS video_count FROM channel AS c LEFT JOIN video AS v ON v.channel_id = c.id GROUP BY 1 ''' curs.execute(stmt) # Replace all the counts of those channels with videos. video_counts.update({int(i['id']): int(i['video_count']) for i in curs.fetchall()}) return video_counts 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: """ Attempt to find a Channel using the provided params. The params are in order of reliability. Raises UnknownChannel if no channel is found. """ channel: COD = None # noqa # Try to find the channel by the most reliable methods first. if channel_id: channel = session.query(Channel).filter_by(id=channel_id).one_or_none() if not channel and source_id: channel = session.query(Channel).filter_by(source_id=source_id).one_or_none() if not channel and url: channel = session.query(Channel).filter_by(url=url).one_or_none() if not channel and directory: directory = Path(directory) if not directory.is_absolute(): directory = get_media_directory() / directory channel = session.query(Channel).filter_by(directory=directory).one_or_none() if not channel and name: channel = session.query(Channel).filter_by(name=name).one_or_none() if not channel: raise UnknownChannel(f'No channel matches {channel_id=} {source_id=} {url=} {directory=}') logger.debug(f'Found {channel=} using {channel_id=} {source_id=} {url=} {directory=}') session.refresh(channel) if return_dict: statistics = channel.get_statistics() channel = channel.dict() channel['statistics'] = statistics 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) # Verify that the URL/Name/directory aren't taken check_for_channel_conflicts( session, id_=channel.id, url=data.url, name=data.name, directory=data.directory, ) data = data.__dict__ old_directory = pathlib.Path(channel.directory) new_directory = get_media_directory() / data.pop('directory') # Apply the changes now that we've OK'd them channel.update(data) session.commit() async def _(): await channel.move_channel(new_directory, session, send_events=True) if new_directory != old_directory: if not new_directory.is_dir(): new_directory.mkdir(parents=True) if PYTEST: await _() else: background_task(_()) save_channels_config.activate_switch() 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. """ try: # Verify that the URL/Name/directory aren't taken check_for_channel_conflicts( session, url=data.url or None, name=data.name, directory=data.directory, source_id=data.source_id, ) except APIError as e: raise ValidationError() from e directory = pathlib.Path(data.directory) directory = get_media_directory() / directory if not directory.is_absolute() else directory directory.mkdir(parents=True, exist_ok=True) channel = Channel() session.add(channel) session.flush([channel]) # Apply the changes now that we've OK'd them channel.update(data.__dict__) session.commit() save_channels_config.activate_switch() 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_dict = channel.dict() channel.delete_with_videos() save_channels_config.activate_switch() return channel_dict @optional_session async def search_channels_by_name(name: str, limit: int = 5, session: Session = None, order_by_video_count: bool = False) -> List[Channel]: name = name or '' name_no_spaces = ''.join(name.split(' ')) if order_by_video_count: stmt = session.query(Channel, func.count(Video.id).label('video_count')) \ .filter(or_( Channel.name.ilike(f'%{name}%'), Channel.name.ilike(f'%{name_no_spaces}%'), )) \ .outerjoin(Video, Video.channel_id == Channel.id) \ .group_by(Channel.id, Channel.name) \ .order_by(desc('video_count'), asc(Channel.name)) \ .limit(limit) channels = [i[0] for i in stmt] else: stmt = session.query(Channel) \ .filter(or_( Channel.name.ilike(f'%{name}%'), Channel.name.ilike(f'%{name_no_spaces}%'), )) \ .order_by(asc(Channel.name)) \ .limit(limit) channels = stmt.all() return channels @wrol_mode_check 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) download.settings = settings save_channels_config.activate_switch() save_downloads_config.activate_switch() return download @wrol_mode_check async def update_channel_download(channel_id: int, download_id: int, url: str, frequency: int, settings: dict): """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) download.url = url download.frequency = frequency download.settings = settings session.commit() download_manager.remove_from_skip_list(url) save_channels_config.activate_switch() save_downloads_config.activate_switch() return download @wrol_mode_check @optional_session async def tag_channel(tag_name: str | None, directory: pathlib.Path | None, channel_id: int, session: Session = None): """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.""" 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) # May also clear the tag if `tag_name` is None. channel.set_tag(tag_name) channel.flush() session.commit() # Only move Channel when requested. if directory: # Move to newly defined directory only if necessary. directory.mkdir(parents=True, exist_ok=True) if channel.directory != directory: coro = channel.move_channel(directory, session, send_events=True) if PYTEST: await coro else: background_task(coro) else: save_channels_config.activate_switch() else: logger.info(f'Tagging {channel} with {tag_name}') @optional_session async def search_channels(tag_names: List[str], session: Session) -> List[Channel]: """Search Tagged Channels.""" from wrolpi.tags import Tag channels = session.query(Channel).join(Tag).filter(Tag.name.in_(tag_names)).all() return channels