colouring-montreal/etl/3_join-data.py
2018-08-01 17:05:02 +01:00

145 lines
3.7 KiB
Python

"""Join shapefile data to buildings
- read through shapes
- locate building in current database (by centroid)
- update building with data
"""
import json
import os
import sys
import fiona
import psycopg2
from shapely.geometry import shape
def main(source_file, config_path, transform_config_path):
"""Load config, read files and save features to the database
"""
conf = read_config(config_path)
transform_config = read_config(transform_config_path)
data_mapping = [
# from_fieldname, to_fieldname, mapping(old_val, source_val)->new_val
(from_, to_, eval(transform))
for from_, to_, transform in transform_config['mapping']
]
dbconf = conf['database']
conn = psycopg2.connect(**dbconf)
with fiona.open(source_file, 'r') as source:
epsg_code = transform_config['crs']
with conn.cursor() as cur:
for feature in source:
geometry_id = find_geom(cur, feature, epsg_code)
if geometry_id is not None:
save_data(
cur, feature['properties'], data_mapping, geometry_id)
else:
print("Skipping", feature['properties'])
conn.commit()
conn.close()
def save_data(cur, props, data_conf, geometry_id):
"""Save data to a building
"""
cur.execute(
"""SELECT building_id, building_doc FROM buildings
WHERE
geometry_id = %s
""", (
geometry_id,
)
)
building = cur.fetchone()
if building is None:
doc = update_from_props({}, props, data_conf)
cur.execute(
"""INSERT INTO buildings
(
building_doc,
geometry_id
)
VALUES
(
%s::jsonb,
%s
)
""", (
json.dumps(doc),
geometry_id
)
)
else:
building_id, old_doc = building
doc = update_from_props(old_doc, props, data_conf)
cur.execute(
"""UPDATE buildings
SET
building_doc = %s::jsonb
WHERE
building_id = %s
""", (
json.dumps(doc),
building_id
)
)
def find_geom(cur, feature, epsg_code):
"""Find a building geometry
"""
cur.execute(
"""SELECT geometry_id FROM geometries
WHERE
ST_Intersects(
ST_Transform(
ST_SetSRID(%s::geometry, %s),
3857
),
geometry_geom
)
""", (
shape(feature['geometry']).wkb_hex,
epsg_code
)
)
result = cur.fetchone()
if result is not None:
return result[0]
else:
return result
def update_from_props(doc, props, mapping):
"""Expect mapping to be a list of transforms
- from_fieldname (expect to find in source feature['properties'])
- to_fieldname (expect to create or find in existing doc)
- transform(old_val, new_val) function/lambda to do any processing
"""
for from_, to_, transform in mapping:
if to_ not in doc:
doc[to_] = None
doc[to_] = transform(doc[to_], props[from_])
return doc
def read_config(config_path):
"""Read a JSON config file containing database connection details
"""
with open(config_path, 'r') as fh:
conf = json.load(fh)
return conf
if __name__ == '__main__':
if len(sys.argv) != 4:
print(
"Usage: {} ./path/to/source/file.csv ./path/to/dbconfig.json ./path/to/mapping.json".format(
os.path.basename(__file__)
))
exit()
main(sys.argv[1], sys.argv[2], sys.argv[3])