mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 10:51:36 -04:00
239 lines
8.5 KiB
Python
239 lines
8.5 KiB
Python
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
|
|
|
|
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
|
|
|
|
|
|
# Bps calculated using many tests on a well-cooled RPi4.
|
|
RPI4_PBF_BYTES_PER_SECOND = 61879
|
|
|
|
|
|
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 size_in_bytes > 1_000_000_000:
|
|
# Use exponential curve for large files. These magic numbers are from testing many imports on an RPi 4.
|
|
a = 7.17509261732342e-14 * size_in_bytes ** 2
|
|
b = 6.6590760410412e-5 * size_in_bytes
|
|
c = 10283
|
|
return int(a - b + c)
|
|
return max(int(size_in_bytes // RPI4_PBF_BYTES_PER_SECOND), 0)
|