diff --git a/exposurelib/database.py b/exposurelib/database.py index f1e0200bf5c912b8f15157c7ba1c8adfeb681620..5aae9d1c0b2e6e796e40f93028c32b5a6a602917 100644 --- a/exposurelib/database.py +++ b/exposurelib/database.py @@ -1068,14 +1068,27 @@ class AbstractExposure(AbstractDatabase, ABC): asset_table = "Asset" entity_table = "Entity" - values = ",".join([f"'{quadkey}'" for quadkey in quadkeys]) + values = ",".join([f"('{quadkey}')" for quadkey in quadkeys]) + # Get all entity IDs first based on the Quadkey list. The statement is split into two + # parts, so the Quadkey index is always used. sql_statement = f""" - SELECT SUM(night), SUM(structural) + WITH Q (quadkey) AS (VALUES {values}) + SELECT id FROM {entity_table} - INNER JOIN {asset_table} - ON {entity_table}.id = {asset_table}.entity_id - WHERE {entity_table}.quadkey IN ({values}) + INNER JOIN Q USING (quadkey) + """ + self.cursor.execute(sql_statement) + values = ",".join([f"({entity_id})" for entity_id, in self.cursor]) + if not values: + return 0, 0 + + # Sum the values based on the list of entity IDs. + sql_statement = f""" + WITH E (id) AS (VALUES {values}) + SELECT SUM(night), SUM(structural) + FROM E + INNER JOIN {asset_table} ON E.id = {asset_table}.entity_id """ self.cursor.execute(sql_statement) total_population, total_structural_value = self.cursor.fetchone() @@ -1100,15 +1113,30 @@ class AbstractExposure(AbstractDatabase, ABC): assets. """ - values = ",".join([f"'{quadkey}'" for quadkey in quadkeys]) + values = ",".join([f"('{quadkey}')" for quadkey in quadkeys]) + # Get all entity IDs first based on the Quadkey list. The statement is split into two + # parts, so the Quadkey index is always used. sql_statement = f""" + WITH Q (quadkey) AS (VALUES {values}) + SELECT id + FROM entity + INNER JOIN Q USING (quadkey) + """ + self.cursor.execute(sql_statement) + + values = ",".join([f"({entity_id})" for entity_id, in self.cursor.fetchall()]) + if not values: + return + + # Update the values based on the list of entity IDs. + sql_statement = f""" + WITH E (id) AS (VALUES {values}) UPDATE Asset SET structural = structural * {structural_value_factor}, night = night * {population_factor} - FROM Entity - WHERE Entity.quadkey IN ({values}) - AND Entity.id = Asset.entity_id + FROM E + WHERE E.id = Asset.entity_id """ self.cursor.execute(sql_statement)