diff --git a/buildinginitializer/buildinginitializer.py b/buildinginitializer/buildinginitializer.py index 7333ce4acc24c4a0edba6719e489b9a39f8a73da..2cd0287bdf0b9c15728eae9e39d8dcb4b1196ea5 100644 --- a/buildinginitializer/buildinginitializer.py +++ b/buildinginitializer/buildinginitializer.py @@ -68,6 +68,10 @@ class BuildingInitializer: self.building_id_column = building_id_column self.building_geometry_column = building_geometry_column + self.outputs = set() + for rule in self.rule_handler: + self.outputs.update(rule.outputs) + if source_db is None: raise ValueError("No source database object is found.") if building_db is None: @@ -251,12 +255,8 @@ class BuildingInitializer: Connected cursor to the OSM Replication database. """ - # Define all unique column names - columns = set() - for _, args in result.items(): - for key in args.keys(): - columns.add(key) - columns = list(columns) + # The rule outputs define the columns that should be upserted. + columns = list(self.outputs) # Format values so that they can be uploaded as a batch on a PostGIS database. values = [] @@ -270,15 +270,15 @@ class BuildingInitializer: columns = ",".join(columns) # Format values values = ",".join(values) - # Add all relevant columns to the database. If the building ID already exists, update # the attributes. sql_statement = f""" - INSERT INTO {self.building_table} ({self.building_id_column}, {columns}) - VALUES {values} - ON CONFLICT ({self.building_id_column}) DO UPDATE - SET {on_conflict} - """ + INSERT INTO {self.building_table} ({self.building_id_column}, {columns}) + VALUES {values} + ON CONFLICT ({self.building_id_column}) DO UPDATE + SET {on_conflict} + """ + cursor.execute(sql_statement) @staticmethod @@ -438,7 +438,7 @@ class BuildingInitializer: SELECT ol.key, ol.type, buildings.id AS building_id FROM osm_lands AS ol INNER JOIN buildings ON ST_Intersects(buildings.geom, ol.geometry) - WHERE buildings.geom && ol.geometry; + WHERE buildings.geom && ol.geometry """ cursor.execute(sql_statement) for key, value, member in cursor: diff --git a/tests/test_buildinginitializer.py b/tests/test_buildinginitializer.py index d7b95f8091afc0c34143101635fb4b6ed46baa14..06101fec0addbc5f7cc22252e93ef541491c5b71 100644 --- a/tests/test_buildinginitializer.py +++ b/tests/test_buildinginitializer.py @@ -87,17 +87,20 @@ def test_upsert_building_information(source_db, building_db): """ building_initializer = OSMInitializer({}, source_db=source_db, building_db=building_db) + building_initializer.outputs = {"occupancy", "storeys", "relation_id"} building_information = { 173039036: {"occupancy": "TEST-1", "storeys": 10, "relation_id": None}, 352224937: {"storeys": 8, "relation_id": 5, "occupancy": "TEST-2"}, } building_db.connect() building_initializer.upsert(building_information, building_db.cursor) - building_db.cursor.execute(""" + building_db.cursor.execute( + """ SELECT osm_id, occupancy, storeys, relation_id FROM obm_buildings WHERE osm_id IN (173039036,352224937) - """) + """ + ) for building_id, occupancy, storeys, relation_id in building_db.cursor: if building_id == 173039036: assert occupancy == "TEST-1" @@ -121,11 +124,13 @@ def test_run(source_db, building_db, rules): building_ids = [352224937, 173039036, 27139591] building_initializer.run(building_ids=building_ids, batch_size=4) building_db.connect() - building_db.cursor.execute(""" + building_db.cursor.execute( + """ SELECT osm_id, storeys FROM obm_buildings WHERE osm_id IN (173039036,352224937) - """) + """ + ) for building_id, storeys in building_db.cursor: if building_id == 173039036: assert storeys == 4