/
osrm2pgsql.py
executable file
·105 lines (85 loc) · 3.27 KB
/
osrm2pgsql.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
Script to upload OSRM binary data into PostGIS
'''
import argparse
import logging
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from stanalysis.osrmbinary import unpack_osrm_nodes, unpack_osrm_edges
from stanalysis.models import OSRMNode, OSRMEdge, Base
log = logging.getLogger(__name__)
def upload_osrm_binary(binaryfd, dbsession, merge=False, commit_every=1000):
''' Upload OSRM data into the db
:param: binaryfd - a file descript pointing to the .osrm data
:param: dbsession - an active :class:`sqlalchemy.Session`
'''
insert_method = dbsession.add if not merge else dbsession.merge
log.info("Loading OSRM node data")
nodes_added = 0
for node in unpack_osrm_nodes(binaryfd):
ormified = OSRMNode(node.id, node.lat, node.lon,
node.bollard, node.traffic_light)
insert_method(ormified)
nodes_added += 1
if nodes_added % commit_every == 0:
log.info("Added %i nodes", nodes_added)
dbsession.commit()
dbsession.commit()
log.info("Added %i nodes", nodes_added)
log.info("Loading OSRM edge data")
existing_edges = set()
edges_added = 0
for edge in unpack_osrm_edges(binaryfd):
ormified = OSRMEdge(
edge.node_a,
edge.node_b,
edge.distance,
edge.weight,
edge.bidirectional,
)
# make sure we don't j
key = OSRMEdge.hash_edge(edge.node_a, edge.node_b)
if key in existing_edges:
log.warning("Skipping existing edge: %s", repr(edge))
continue
existing_edges.add(key)
insert_method(ormified)
edges_added += 1
if edges_added % commit_every == 0:
log.info("Added %i edges", edges_added)
dbsession.commit()
dbsession.commit()
log.info("Added %i edges", edges_added)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('osrm_file', metavar='input.osrm',
type=argparse.FileType('rb'),
help='Input .osrm binary file')
parser.add_argument(
'--connection',
default='postgresql://osm:osm@localhost/osrm',
help='Postgres connection string. Default %(default)s'
)
parser.add_argument('--mode', choices=['append', 'recreate', 'update'],
default='recreate',
help='If "recreate", the tables will be dropped')
parser.add_argument('--verbose', action='store_true',
help='Increase logging level')
args = parser.parse_args()
logging.basicConfig(
level=logging.INFO if args.verbose else logging.WARNING)
log.info("Creating DB engine")
engine = create_engine(args.connection, echo=False)
log.info("Creating DB session")
Session = sessionmaker(bind=engine)
session = Session()
if args.mode == 'recreate':
log.info("Dropping existing OSRM tables")
Base.metadata.drop_all(engine)
if args.mode in ('append', 'recreate'):
log.info("Creating OSRM tables")
Base.metadata.create_all(engine)
upload_osrm_binary(args.osrm_file, session, args.mode == 'update')
log.info("Done.")