From 8e35718b21ab30993b726c06823a0c6f338f3b95 Mon Sep 17 00:00:00 2001
From: Laurens Oostwegel <laurens@gfz-potsdam.de>
Date: Wed, 10 Apr 2024 08:30:10 +0200
Subject: [PATCH] Restructure the Quadkey selection statement of the country
 calibration SQL query

---
 exposurelib/database.py | 46 +++++++++++++++++++++++++++++++++--------
 1 file changed, 37 insertions(+), 9 deletions(-)

diff --git a/exposurelib/database.py b/exposurelib/database.py
index f1e0200b..5aae9d1c 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)
 
-- 
GitLab