wrolpi/modules/files/lib.py

import asyncio
import re
from functools import wraps
from pathlib import Path
from typing import List
from uuid import uuid4

import magic
import psycopg2
from sqlalchemy.orm import Session

from modules.files.models import File
from wrolpi.common import get_media_directory, wrol_mode_check, walk, chunks, logger
from wrolpi.dates import from_timestamp
from wrolpi.db import get_db_session, get_db_curs, get_ranked_models
from wrolpi.errors import InvalidFile
from wrolpi.vars import PYTEST

logger = logger.getChild(__name__)

__all__ = ['list_files', 'delete_file', 'split_file_name', 'upsert_file', 'refresh_files', 'search']


def filter_parent_directories(directories: List[Path]) -> List[Path]:
    """
    Remove parent directories if their children are in the list.

    >>> filter_parent_directories([Path('foo'), Path('foo/bar'), Path('baz')])
    [Path('foo/bar'), Path('baz')]
    """
    unique_children = set()
    for directory in sorted(directories):
        for parent in directory.parents:
            # Remove any parent of this child.
            if parent in unique_children:
                unique_children.remove(parent)
        unique_children.add(directory)

    # Restore the original order.
    new_directories = [i for i in directories if i in unique_children]
    return new_directories


def list_files(directories: List[str]) -> List[Path]:
    """
    List all files down to the directories provided.  This includes all parent directories of the directories.
    """
    media_directory = get_media_directory()

    # Always display the media_directory files.
    paths = list(media_directory.iterdir())

    if directories:
        directories = [media_directory / i for i in directories if i]
        directories = filter_parent_directories(directories)
        for directory in directories:
            directory = media_directory / directory
            for parent in directory.parents:
                is_relative_to = str(media_directory).startswith(str(parent))
                if parent == Path('.') or is_relative_to:
                    continue
                paths.extend(parent.iterdir())
            paths.extend(directory.iterdir())

    return paths


@wrol_mode_check
def delete_file(file: str):
    """Delete a file in the media directory."""
    file = get_media_directory() / file
    if file.is_dir() or not file.is_file():
        raise InvalidFile(f'Invalid file {file}')
    file.unlink()


FILE_NAME_REGEX = re.compile(r'[_ ]')


def split_file_name(path: Path) -> List[str]:
    """
    Split a file name into words.

    >>> split_file_name(Path('foo.mp4'))
    ['foo']
    >>> split_file_name(Path('foo bar.mp4'))
    ['foo', 'bar']
    """
    words = [str(i) for i in FILE_NAME_REGEX.split(path.stem)]
    return words


def upsert_file(path: Path, session: Session) -> File:
    """Update/insert a File in the DB.  Gather metadata about it."""
    file = session.query(File).filter_by(path=path).one_or_none()
    if not file:
        file = File(path=path)
        session.add(file)
    if not file.mimetype:
        file.mimetype = magic.from_file(str(path), mime=True)
    if not file.title:
        file.title = split_file_name(path)
    if not file.size or not file.modification_datetime:
        stat = path.stat()
        file.size = stat.st_size
        file.modification_datetime = file.modification_datetime or from_timestamp(stat.st_mtime)

    return file


def _refresh_files():
    """Find and index all files"""
    logger.info('Refreshing Files')
    # Mark all files as unverified.  Any records that are unverified after refresh will be deleted.
    with get_db_curs(commit=True) as curs:
        curs.execute('UPDATE file SET idempotency=null')

    paths = filter(lambda i: i.is_file(), walk(get_media_directory()))
    idempotency = str(uuid4())
    for chunk in chunks(paths, 20):
        with get_db_session(commit=True) as session:
            for path in chunk:
                file = upsert_file(path, session)
                file.idempotency = idempotency

    # Remove any records where the file no longer exists.
    with get_db_curs(commit=True) as curs:
        curs.execute('DELETE FROM file WHERE idempotency IS NULL')

    logger.info('Done refreshing Files')


@wraps(_refresh_files)
def refresh_files():
    """Schedule a refresh task if not testing.  If testing, do a synchronous refresh."""
    if PYTEST:
        return _refresh_files()

    async def _():
        return _refresh_files()

    asyncio.create_task(_())


def search(search_str: str, limit: int, offset: int):
    """
    Search the "title" of each file.  The title is the words in the file's stem.
    """
    with get_db_curs() as curs:
        stmt = '''
            SELECT id, ts_rank_cd(textsearch, websearch_to_tsquery(%(search_str)s)), COUNT(*) OVER() AS total
            FROM file
            WHERE textsearch @@ websearch_to_tsquery(%(search_str)s)
            OFFSET %(offset)s LIMIT %(limit)s
        '''
        params = dict(search_str=search_str, offset=offset, limit=limit)
        curs.execute(stmt, params)

        try:
            results = list(map(dict, curs.fetchall()))
        except psycopg2.ProgrammingError:
            # No files.
            return [], 0

        total = results[0]['total'] if results else 0
        ranked_ids = [i['id'] for i in results]

    with get_db_session() as session:
        results = get_ranked_models(ranked_ids, File, session)
        results = [i.__json__() for i in results]

    return results, total