wrolpi/modules/map/lib.py

import asyncio
import subprocess
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.api_utils import api_app
from wrolpi.cmd import SUDO_BIN
from wrolpi.common import get_media_directory, walk, logger, get_wrolpi_config
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__)


def get_map_directory() -> Path:
    map_directory = get_media_directory() / get_wrolpi_config().map_directory
    if not map_directory.is_dir():
        map_directory.mkdir(parents=True)
    return map_directory


def is_pbf_file(pbf: Path) -> bool:
    """Uses file command to check type of 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:
                    api_app.shared_ctx.map_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:
                    api_app.shared_ctx.map_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:
                    api_app.shared_ctx.map_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:
                    api_app.shared_ctx.map_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)