import asyncio import subprocess from multiprocessing import Event, Manager from pathlib import Path from typing import List from sqlalchemy.orm import Session from modules.map.models import MapFile from wrolpi.cmd import BASH_BIN from wrolpi.common import get_media_directory, walk, logger, wrol_mode_check from wrolpi.dates import now, timedelta_to_timestamp from wrolpi.db import optional_session, get_db_session from wrolpi.vars import PYTEST, PROJECT_DIR logger = logger.getChild(__name__) IMPORT_EVENT = Event() 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 IMPORT_EVENT.is_set(): import_logger.error('Map import already running...') return 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'] any_success = False try: IMPORT_EVENT.set() if pbfs: success = False try: IMPORTING.update(dict( pending=list(pbfs), )) 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), )) 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: # A map was imported, remove the tile cache files. await clear_mod_tile() IMPORT_EVENT.clear() async def run_import_command(*paths: Path): """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. """ 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) cmd = f'{BASH_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}') @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 MOD_TILE_CACHE_DIR = Path('/var/lib/mod_tile/ajt') @wrol_mode_check async def clear_mod_tile(): """Remove all cached map tile files""" if PYTEST: return logger.warning('Clearing map tile cache files') if MOD_TILE_CACHE_DIR.is_dir(): await asyncio.create_subprocess_shell(f'rm -r {MOD_TILE_CACHE_DIR}')