wrolpi/modules/map/lib.py

import asyncio
import subprocess
from multiprocessing import Manager
from pathlib import Path
from typing import List

from sqlalchemy.orm import Session

from modules.map.models import MapFile
from wrolpi import flags
from wrolpi.cmd import SUDO_BIN
from wrolpi.common import get_media_directory, walk, logger
from wrolpi.dates import now, timedelta_to_timestamp, seconds_to_timestamp
from wrolpi.db import optional_session, get_db_session
from wrolpi.events import Events
from wrolpi.vars import PROJECT_DIR, IS_RPI5

logger = logger.getChild(__name__)

IMPORTING = Manager().dict(dict(
    pending=None,
))


def get_map_directory() -> Path:
    map_directory = get_media_directory() / 'map'
    if not map_directory.is_dir():
        map_directory.mkdir()
    return map_directory


def is_pbf_file(pbf: Path) -> bool:
    """Uses file command to check type of a file.  Returns True if a file is an OpenStreetMap PBF file."""
    cmd = ('/usr/bin/file', pbf)
    try:
        output = subprocess.check_output(cmd)
    except FileNotFoundError:
        return False

    return 'OpenStreetMap Protocolbuffer Binary Format' in output.decode()


def is_dump_file(path: Path) -> bool:
    """Uses file command to check type of a file.  Returns True if a file is a Postgresql dump file."""
    cmd = ('/usr/bin/file', path)
    try:
        output = subprocess.check_output(cmd)
    except FileNotFoundError:
        return False

    return 'PostgreSQL custom database dump' in output.decode()


def get_map_paths() -> List[Path]:
    """Find all pbf/dump files in the map directory."""
    map_directory = get_map_directory()
    paths = walk(map_directory)

    def is_valid(path: Path) -> bool:
        if path.is_file():
            if str(path).endswith('.osm.pbf') and is_pbf_file(path):
                return True
            elif path.suffix == '.dump' and is_dump_file(path):
                return True
        return False

    return list(filter(is_valid, paths))


def get_or_create_map_file(pbf_path: Path, session: Session) -> MapFile:
    """Finds the MapFile row in the DB, or creates one."""
    map_file: MapFile = session.query(MapFile).filter_by(path=str(pbf_path)).one_or_none()
    if map_file:
        return map_file

    map_file = MapFile(path=pbf_path, size=pbf_path.stat().st_size)
    session.add(map_file)
    return map_file


import_logger = logger.getChild('import')


async def import_files(paths: List[str]):
    if flags.map_importing.is_set():
        import_logger.error('Map import already running...')
        return

    with flags.map_importing:
        paths = [get_media_directory() / i for i in paths]
        import_logger.warning(f'Importing: {", ".join(map(str, paths))}')

        dumps = [i for i in paths if i.suffix == '.dump']
        pbfs = [i for i in paths if i.suffix == '.pbf']

        total_elapsed = 0

        any_success = False
        try:
            if pbfs:
                success = False
                try:
                    IMPORTING.update(dict(
                        pending=list(pbfs),
                    ))
                    total_elapsed += await run_import_command(*pbfs)
                    success = True
                    any_success = True
                except Exception as e:
                    import_logger.warning('Failed to run import', exc_info=e)
                finally:
                    IMPORTING.update(dict(
                        pending=None,
                    ))

                if success:
                    with get_db_session(commit=True) as session:
                        # Any previously imported PBFs are no longer imported.
                        for pbf_path in pbfs:
                            pbf_file = get_or_create_map_file(pbf_path, session)
                            pbf_file.imported = True
                        for pbf_file in session.query(MapFile):
                            pbf_file.imported = pbf_file.path in pbfs

            for path in dumps:
                # Import each dump individually.
                if not path.is_file():
                    import_logger.fatal(f'Map file does not exist! {path}')
                    continue

                with get_db_session() as session:
                    map_file = session.query(MapFile).filter_by(path=path).one_or_none()
                    if map_file and map_file.imported:
                        # Don't import a map file twice.
                        import_logger.debug(f'{path} is already imported')
                        continue

                success = False
                try:
                    IMPORTING.update(dict(
                        pending=str(path),
                    ))
                    total_elapsed += await run_import_command(path)
                    success = True
                    any_success = True
                except Exception as e:
                    import_logger.warning('Failed to run import', exc_info=e)
                finally:
                    IMPORTING.update(dict(
                        pending=None,
                    ))

                if success:
                    with get_db_session(commit=True) as session:
                        map_file = get_or_create_map_file(path, session)
                        map_file.imported = True
        finally:
            if any_success:
                Events.send_map_import_complete(f'Map import completed; took {seconds_to_timestamp(total_elapsed)}')
            else:
                Events.send_map_import_failed(f'Map import failed!  See server logs.')


async def run_import_command(*paths: Path) -> int:
    """Run the map import script on the provided paths.

    Can only import a single *.dump file, or a list of *.osm.pbf files.  They cannot be mixed.

    @return: The seconds elapsed during import.
    """
    paths = [i.absolute() for i in paths]
    dumps = [i for i in paths if i.suffix == '.dump']
    pbfs = [i for i in paths if i.suffix == '.pbf']

    # Only import a single dump, or a list of pbfs.  No other combinations acceptable.
    if not paths:
        raise ValueError('Must import a file')
    if not dumps and not pbfs:
        raise ValueError('Cannot import unknown file!')
    if dumps and pbfs:
        raise ValueError('Cannot import mixed dumps and pbfs.')
    if len(dumps) > 1:
        raise ValueError('Cannot import more than one dump')
    if dumps and not is_dump_file(dumps[0]):
        import_logger.warning(f'Could not import non-dump file: {dumps}')
        raise ValueError('Invalid dump file')
    if pbfs:
        for path in pbfs:
            if not is_pbf_file(path):
                import_logger.warning(f'Could not import non-pbf file: {path}')
                raise ValueError('Invalid PBF file')

    paths = ' '.join(str(i) for i in paths)
    # Run with sudo so renderd can be restarted.
    cmd = f'{SUDO_BIN} {PROJECT_DIR}/scripts/import_map.sh {paths}'
    import_logger.warning(f'Running map import command: {cmd}')
    start = now()
    proc = await asyncio.create_subprocess_shell(cmd, stderr=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE)

    stdout, stderr = await proc.communicate()
    elapsed = now() - start
    success = 'Successful' if proc.returncode == 0 else 'Unsuccessful'
    import_logger.info(f'{success} import of {repr(paths)} took {timedelta_to_timestamp(elapsed)}')
    if proc.returncode != 0:
        # Log all lines.  Truncate long lines.
        for line in stdout.decode().splitlines():
            import_logger.debug(line[:500])
        for line in stderr.decode().splitlines():
            import_logger.error(line[:500])
        raise ValueError(f'Importing map file failed with return code {proc.returncode}')

    return int(elapsed.total_seconds())


@optional_session
def get_import_status(session: Session = None) -> List[MapFile]:
    paths = get_map_paths()
    map_paths = []
    for path in paths:
        map_file = get_or_create_map_file(path, session)
        map_paths.append(map_file)

    session.commit()
    return map_paths


# Calculated using many tests on a well-cooled RPi4 (4GB).
RPI4_PBF_BYTES_PER_SECOND = 61879
RPI4_A = 7.17509261732342e-14
RPI4_B = 6.6590760410412e-5
RPI4_C = 10283
# Calculated using many tests on a well-cooled RPi5 (8GB).
RPI5_PBF_BYTES_PER_SECOND = 136003
RPI5_A = 1.6681382703586e-15
RPI5_B = 2.35907824676145e-6
RPI5_C = 53.727


def seconds_to_import_rpi4(size_in_bytes: int) -> int:
    if size_in_bytes > 1_000_000_000:
        # Use exponential curve for large files.
        a = RPI4_A * size_in_bytes ** 2
        b = RPI4_B * size_in_bytes
        return int(a - b + RPI4_C)
    # Use simpler equation for small files.
    return max(int(size_in_bytes // RPI4_PBF_BYTES_PER_SECOND), 0)


def seconds_to_import_rpi5(size_in_bytes: int) -> int:
    if size_in_bytes > 5_000_000_000:
        # Use exponential curve for large files.
        a = RPI5_A * size_in_bytes ** 2
        b = RPI5_B * size_in_bytes
        return int(a - b + RPI5_C)
    # Use simpler equation for small files.
    return max(int(size_in_bytes // RPI5_PBF_BYTES_PER_SECOND), 0)


def seconds_to_import(size_in_bytes: int) -> int:
    """Attempt to predict how long it will take an RPi4 to import a given PBF file."""
    if IS_RPI5:
        return seconds_to_import_rpi5(size_in_bytes)

    # Default to RPi4, because it's the most conservative estimate.
    return seconds_to_import_rpi4(size_in_bytes)