diff --git a/hub/persistence/db_control.py b/hub/persistence/db_control.py index 6a1c3ac5..c6740336 100644 --- a/hub/persistence/db_control.py +++ b/hub/persistence/db_control.py @@ -123,22 +123,20 @@ class DBControl: if result_sets is None: continue results[scenario_name] = [] - for result_set in result_sets: - city_id = result_set[0].id - for building_name in scenario[scenario_name]: - _building = self._city_object.get_by_name_or_alias_and_city(building_name, city_id) - if _building is None: - continue - city_object_id = _building.id - _ = self._simulation_results.get_simulation_results_by_city_id_city_object_id_and_names( - city_id, - city_object_id, - result_names) + city_ids = [r[0].id for r in result_sets] + for building_name in scenario[scenario_name]: + _building = self._city_object.get_by_name_or_alias_in_cities(building_name, city_ids) + if _building is None: + continue + city_object_id = _building.id + _ = self._simulation_results.get_simulation_results_by_city_object_id_and_names( + city_object_id, + result_names) - for value in _: - values = json.loads(value.values) - values["building"] = building_name - results[scenario_name].append(values) + for value in _: + values = json.loads(value.values) + values["building"] = building_name + results[scenario_name].append(values) return results def persist_city(self, city: City, pickle_path, scenario, application_id: int, user_id: int): diff --git a/hub/persistence/repositories/city_object.py b/hub/persistence/repositories/city_object.py index 30287411..5c805d15 100644 --- a/hub/persistence/repositories/city_object.py +++ b/hub/persistence/repositories/city_object.py @@ -131,3 +131,35 @@ class CityObject(Repository): except IndexError as err: logging.error('Error while fetching city object by name and city, empty result %s', err) raise IndexError from err + + def get_by_name_or_alias_in_cities(self, name, city_ids) -> Model: + """ + Fetch a city object based on name and city ids + :param name: city object name + :param city_ids: a list of city identifiers + :return: [CityObject] with the provided name or alias belonging to the city with id city_id + """ + try: + # search by name first + with Session(self.engine) as session: + city_object = session.execute(select(Model).where(Model.name == name, Model.city_id in city_ids)).first() + if city_object is not None: + return city_object[0] + # name not found, so search by alias instead + city_objects = session.execute( + select(Model).where(Model.aliases.contains(name), Model.city_id in city_ids) + ).all() + for city_object in city_objects: + aliases = city_object[0].aliases.replace('{', '').replace('}', '').split(',') + for alias in aliases: + if alias == name: + # force the name as the alias + city_object[0].name = name + return city_object[0] + return None + except SQLAlchemyError as err: + logging.error('Error while fetching city object by name and city: %s', err) + raise SQLAlchemyError from err + except IndexError as err: + logging.error('Error while fetching city object by name and city, empty result %s', err) + raise IndexError from err diff --git a/hub/persistence/repositories/simulation_results.py b/hub/persistence/repositories/simulation_results.py index 3f6faa48..b2757f40 100644 --- a/hub/persistence/repositories/simulation_results.py +++ b/hub/persistence/repositories/simulation_results.py @@ -167,3 +167,27 @@ class SimulationResults(Repository): except SQLAlchemyError as err: logging.error('Error while fetching city by city_id: %s', err) raise SQLAlchemyError from err + + def get_simulation_results_by_city_object_id_and_names(self, city_object_id, result_names=None) -> [Model]: + """ + Fetch the simulation results based in the city_object_id with the given names or all + :param city_object_id: the city object id + :param result_names: if given filter the results + :return: [SimulationResult] + """ + try: + with Session(self.engine) as session: + result_set = session.execute(select(Model).where( + Model.city_object_id == city_object_id + )) + results = [r[0] for r in result_set] + if not result_names: + return results + filtered_results = [] + for result in results: + if result.name in result_names: + filtered_results.append(result) + return filtered_results + except SQLAlchemyError as err: + logging.error('Error while fetching city by city_id: %s', err) + raise SQLAlchemyError from err \ No newline at end of file