118 lines
2.9 KiB
Python
118 lines
2.9 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
|
||
|
|
from dataclasses import asdict, dataclass
|
||
|
|
import asyncio
|
||
|
|
from os.path import basename, splitext
|
||
|
|
import sys
|
||
|
|
import logging
|
||
|
|
|
||
|
|
import msgpack
|
||
|
|
import psycopg
|
||
|
|
from psycopg.types import TypeInfo
|
||
|
|
from psycopg.types.shapely import register_shapely
|
||
|
|
from shapely.wkb import loads
|
||
|
|
|
||
|
|
from obs.api.app import app
|
||
|
|
from obs.api.db import ZoneType, connect_db, make_session
|
||
|
|
from obs.api.utils import chunk
|
||
|
|
|
||
|
|
log = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
ROAD_BUFFER = 1000
|
||
|
|
AREA_BUFFER = 100
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class Region:
|
||
|
|
relation_id: int
|
||
|
|
name: str
|
||
|
|
admin_level: int
|
||
|
|
geometry: bytes
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class Road:
|
||
|
|
way_id: int
|
||
|
|
name: str
|
||
|
|
zone: str
|
||
|
|
directionality: int
|
||
|
|
oneway: int
|
||
|
|
geometry: bytes
|
||
|
|
|
||
|
|
|
||
|
|
def read_file(filename, type_only=False):
|
||
|
|
"""
|
||
|
|
Reads a file iteratively, yielding Road and Region objects as they
|
||
|
|
appear. Those may be mixed.
|
||
|
|
"""
|
||
|
|
with open(filename, "rb") as f:
|
||
|
|
unpacker = msgpack.Unpacker(f)
|
||
|
|
try:
|
||
|
|
while True:
|
||
|
|
type_id = unpacker.unpack()
|
||
|
|
data = unpacker.unpack()
|
||
|
|
|
||
|
|
if type_id == b"\x01":
|
||
|
|
yield Road(*data)
|
||
|
|
|
||
|
|
# elif type_id == b"\x02":
|
||
|
|
# yield Region(*data)
|
||
|
|
except msgpack.OutOfData:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
async def import_osm(connection, filename, import_group=None):
|
||
|
|
if import_group is None:
|
||
|
|
import_group = splitext(basename(filename))[0]
|
||
|
|
|
||
|
|
# Pass 1: Find IDs only
|
||
|
|
road_ids = []
|
||
|
|
for item in read_file(filename):
|
||
|
|
road_ids.append(item.way_id)
|
||
|
|
|
||
|
|
async with connection.cursor() as cursor:
|
||
|
|
print(f"Pass 1: Delete previous")
|
||
|
|
|
||
|
|
print(f"Deleting import group {import_group}")
|
||
|
|
await cursor.execute(
|
||
|
|
"DELETE FROM road WHERE import_group = %s", (import_group,)
|
||
|
|
)
|
||
|
|
|
||
|
|
print(f"Deleting by ID")
|
||
|
|
for ids in chunk(road_ids, 10000):
|
||
|
|
await cursor.execute("DELETE FROM road WHERE way_id = ANY(%s)", (ids,))
|
||
|
|
|
||
|
|
# Pass 2: Import
|
||
|
|
print(f"Pass 2: Import")
|
||
|
|
|
||
|
|
async with cursor.copy(
|
||
|
|
"COPY road (way_id, name, zone, directionality, oneway, geometry, import_group) FROM STDIN"
|
||
|
|
) as copy:
|
||
|
|
for item in read_file(filename):
|
||
|
|
await copy.write_row(
|
||
|
|
(
|
||
|
|
item.way_id,
|
||
|
|
item.name,
|
||
|
|
item.zone,
|
||
|
|
item.directionality,
|
||
|
|
item.oneway,
|
||
|
|
bytes.hex(item.geometry),
|
||
|
|
import_group,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
async def main():
|
||
|
|
url = app.config.POSTGRES_URL
|
||
|
|
url = url.replace("+asyncpg", "")
|
||
|
|
|
||
|
|
async with await psycopg.AsyncConnection.connect(url) as connection:
|
||
|
|
for filename in sys.argv[1:]:
|
||
|
|
print("Loading file", filename)
|
||
|
|
await import_osm(connection, filename)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
asyncio.run(main())
|