colouring-montreal/etl/1_load_geometries.py
2018-09-09 11:32:27 +01:00

106 lines
2.4 KiB
Python

"""Load geometries from GeoJSON to Postgres
- create 'geometry' record with {
id: <polygon-guid>,
doc: {source_id: <toid>},
geom: <geom-wkb_hex>
}
- create corresponding 'building' record with {
id: <building-guid>,
doc: {},
geom_id: <polygon-guid>
}
"""
import glob
import json
import os
import sys
import fiona
import psycopg2
from shapely.geometry import shape
def main(source_dir, config_path):
"""Load config, read files and save features to the database
"""
conf = read_config(config_path)
dbconf = conf['database']
conn = psycopg2.connect(**dbconf)
source_files = glob.glob("{}/*.geojson".format(source_dir))
loaded = {}
for source_file in source_files:
with fiona.open(source_file, 'r') as source:
with conn.cursor() as cur:
for feature in source:
fid = feature['properties']['fid']
if fid not in loaded:
save_feature(cur, feature)
loaded[fid] = True
else:
print("Skipping", fid)
conn.commit()
conn.close()
def save_feature(cur, feature):
"""Save a feature with geometry and source id
"""
cur.execute(
"""INSERT INTO geometries
(
geometry_doc,
geometry_geom
)
VALUES
(
%s::jsonb,
ST_SetSRID(%s::geometry, %s)
)
RETURNING geometry_id
""", (
json.dumps({
'source_id': feature['properties']['fid']
}),
shape(feature['geometry']).wkb_hex,
3857
)
)
geom_id, = cur.fetchone()
cur.execute(
"""INSERT INTO buildings
(
building_doc,
geometry_id
)
VALUES
(
%s::jsonb,
%s
)
""", (
json.dumps({}),
geom_id
)
)
def read_config(config_path):
"""Read a JSON config file containing database connection details
"""
with open(config_path, 'r') as conf_fh:
conf = json.load(conf_fh)
return conf
if __name__ == '__main__':
if len(sys.argv) != 3:
print("Usage: {} ./path/to/source/dir ./path/to/dbconfig.json".format(
os.path.basename(__file__)
))
exit()
main(sys.argv[1], sys.argv[2])