feat: pass db connection for global as param

otherwise API clients need to set special
environment variables for this library

Author: shrouxm

soil_id/db.py

@@ -21,11 +21,11 @@
 
 # Third-party libraries
 import psycopg
-from shapely import wkb
-from shapely.geometry import Point
 
 # local libraries
 import soil_id.config
+from shapely import wkb
+from shapely.geometry import Point
 
 
 def get_datastore_connection():
@@ -162,7 +162,7 @@ def save_soilgrids_output(plot_id, model_version, soilgrids_blob):
         conn.close()
 
 
-def get_hwsd2_profile_data(conn, hwsd2_mu_select):
+def get_hwsd2_profile_data(connection, hwsd2_mu_select):
     """
     Retrieve HWSD v2 data based on selected hwsd2 (map unit) values.
     This version reuses an existing connection.
@@ -179,7 +179,7 @@ def get_hwsd2_profile_data(conn, hwsd2_mu_select):
         return pd.DataFrame()
 
     try:
-        with conn.cursor() as cur:
+        with connection.cursor() as cur:
             # Create placeholders for the SQL IN clause
             placeholders = ", ".join(["%s"] * len(hwsd2_mu_select))
             sql_query = f"""
@@ -222,7 +222,7 @@ def get_hwsd2_profile_data(conn, hwsd2_mu_select):
         return pd.DataFrame()
 
 
-def extract_hwsd2_data(lon, lat, buffer_dist, table_name):
+def extract_hwsd2_data(connection, lon, lat, buffer_dist, table_name):
     """
     Fetches HWSD soil data from a PostGIS table within a given buffer around a point,
     performing distance and intersection calculations directly on geographic coordinates.
@@ -236,133 +236,121 @@ def extract_hwsd2_data(lon, lat, buffer_dist, table_name):
     Returns:
         DataFrame: Merged data from hwsdv2 and hwsdv2_data.
     """
-    # Use a single connection for both queries.
-    with get_datastore_connection() as conn:
-        # Compute the buffer polygon (in WKT) around the problem point.
-        # Here, we use the geography type to compute a buffer in meters,
-        # then cast it back to geometry in EPSG:4326.
-        buffer_query = """
-            WITH buffer AS (
-              SELECT ST_AsText(
-                       ST_Buffer(
-                         ST_SetSRID(ST_Point(%s, %s), 4326)::geography,
-                         %s
-                       )::geometry
-                     ) AS wkt
-            )
-            SELECT wkt FROM buffer;
-        """
-        with conn.cursor() as cur:
-            cur.execute(buffer_query, (lon, lat, buffer_dist))
-            buffer_wkt = cur.fetchone()[0]
-            print("Buffer WKT:", buffer_wkt)
-
-        # Build the main query that uses the computed buffer.
-        # Distance is computed by casting geometries to geography,
-        # which returns the geodesic distance in meters.
-        main_query = f"""
-            WITH
-            -- Step 1: Get the polygon that contains the point
-            point_poly AS (
-            SELECT ST_MakeValid(geom) AS geom
-            FROM {table_name}
-            WHERE ST_Intersects(
-                geom,
-                ST_SetSRID(ST_Point({lon}, {lat}), 4326)
-            )
-            ),
-
-            -- Step 2: Get polygons that intersect the buffer
-            valid_geom AS (
-            SELECT
-                hwsd2,
-                ST_MakeValid(geom) AS geom
-            FROM {table_name}
-            WHERE geom && ST_GeomFromText('{buffer_wkt}', 4326)
-                AND ST_Intersects(geom, ST_GeomFromText('{buffer_wkt}', 4326))
-            )
-
-            -- Step 3: Filter to those that either contain the point or border the point's polygon
-            SELECT
-            vg.hwsd2,
-            ST_AsEWKB(vg.geom) AS geom,
-            ST_Distance(
-                vg.geom::geography,
-                ST_SetSRID(ST_Point({lon}, {lat}), 4326)::geography
-            ) AS distance,
-            ST_Intersects(
-                vg.geom,
-                ST_SetSRID(ST_Point({lon}, {lat}), 4326)
-            ) AS pt_intersect
-            FROM valid_geom vg, point_poly pp
-            WHERE
-            ST_Intersects(vg.geom, ST_SetSRID(ST_Point({lon}, {lat}), 4326))
-            OR ST_Intersects(vg.geom, pp.geom);
-        """
+    # Compute the buffer polygon (in WKT) around the problem point.
+    # Here, we use the geography type to compute a buffer in meters,
+    # then cast it back to geometry in EPSG:4326.
+    buffer_query = """
+        WITH buffer AS (
+          SELECT ST_AsText(
+                   ST_Buffer(
+                     ST_SetSRID(ST_Point(%s, %s), 4326)::geography,
+                     %s
+                   )::geometry
+                 ) AS wkt
+        )
+        SELECT wkt FROM buffer;
+    """
+    with connection.cursor() as cur:
+        cur.execute(buffer_query, (lon, lat, buffer_dist))
+        buffer_wkt = cur.fetchone()[0]
+        print("Buffer WKT:", buffer_wkt)
+
+    # Build the main query that uses the computed buffer.
+    # Distance is computed by casting geometries to geography,
+    # which returns the geodesic distance in meters.
+    main_query = f"""
+        WITH
+        -- Step 1: Get the polygon that contains the point
+        point_poly AS (
+        SELECT ST_MakeValid(geom) AS geom
+        FROM {table_name}
+        WHERE ST_Intersects(
+            geom,
+            ST_SetSRID(ST_Point({lon}, {lat}), 4326)
+        )
+        ),
+
+        -- Step 2: Get polygons that intersect the buffer
+        valid_geom AS (
+        SELECT
+            hwsd2,
+            ST_MakeValid(geom) AS geom
+        FROM {table_name}
+        WHERE geom && ST_GeomFromText('{buffer_wkt}', 4326)
+            AND ST_Intersects(geom, ST_GeomFromText('{buffer_wkt}', 4326))
+        )
 
+        -- Step 3: Filter to those that either contain the point or border the point's polygon
+        SELECT
+        vg.hwsd2,
+        ST_AsEWKB(vg.geom) AS geom,
+        ST_Distance(
+            vg.geom::geography,
+            ST_SetSRID(ST_Point({lon}, {lat}), 4326)::geography
+        ) AS distance,
+        ST_Intersects(
+            vg.geom,
+            ST_SetSRID(ST_Point({lon}, {lat}), 4326)
+        ) AS pt_intersect
+        FROM valid_geom vg, point_poly pp
+        WHERE
+        ST_Intersects(vg.geom, ST_SetSRID(ST_Point({lon}, {lat}), 4326))
+        OR ST_Intersects(vg.geom, pp.geom);
+    """
 
-        # Use GeoPandas to execute the main query and load results into a GeoDataFrame.
-        hwsd = gpd.read_postgis(main_query, conn, geom_col="geom")
-        print("Main query returned", len(hwsd), "rows.")
+    # Use GeoPandas to execute the main query and load results into a GeoDataFrame.
+    hwsd = gpd.read_postgis(main_query, connection, geom_col="geom")
+    print("Main query returned", len(hwsd), "rows.")
 
-        # Remove the geometry column (if not needed) from this dataset.
-        hwsd = hwsd.drop(columns=["geom"])
+    # Remove the geometry column (if not needed) from this dataset.
+    hwsd = hwsd.drop(columns=["geom"])
 
-        # Get the list of hwsd2 identifiers.
-        hwsd2_mu_select = hwsd["hwsd2"].tolist()
+    # Get the list of hwsd2 identifiers.
+    hwsd2_mu_select = hwsd["hwsd2"].tolist()
 
-        # Call get_hwsd2_profile_data using the same connection.
-        hwsd_data = get_hwsd2_profile_data(conn, hwsd2_mu_select)
+    # Call get_hwsd2_profile_data using the same connection.
+    hwsd_data = get_hwsd2_profile_data(connection, hwsd2_mu_select)
 
-        # Merge the two datasets.
-        merged = pd.merge(hwsd_data, hwsd, on="hwsd2", how="left").drop_duplicates()
-        return merged
+    # Merge the two datasets.
+    merged = pd.merge(hwsd_data, hwsd, on="hwsd2", how="left").drop_duplicates()
+    return merged
 
 
 # global
 
 
 # Function to fetch data from a PostgreSQL table
-def fetch_table_from_db(table_name):
-    conn = None
+def fetch_table_from_db(connection, table_name):
     try:
-        conn = get_datastore_connection()
-        cur = conn.cursor()
-
-        query = f"SELECT * FROM {table_name} ORDER BY id ASC;"
-        cur.execute(query)
-        rows = cur.fetchall()
+        with connection.cursor() as cur:
+            query = f"SELECT * FROM {table_name} ORDER BY id ASC;"
+            cur.execute(query)
+            rows = cur.fetchall()
 
-        return rows
+            return rows
 
     except Exception as err:
         logging.error(f"Error querying PostgreSQL: {err}")
         return None
 
-    finally:
-        if conn:
-            conn.close()
 
-
-def get_WRB_descriptions(WRB_Comp_List):
+def get_WRB_descriptions(connection, WRB_Comp_List):
     """
     Retrieve WRB descriptions based on provided WRB component list.
     """
-    conn = None
     try:
-        conn = get_datastore_connection()
-        cur = conn.cursor()
+        with connection.cursor() as cur:
 
-        # Create placeholders for the SQL IN clause
-        placeholders = ", ".join(["%s"] * len(WRB_Comp_List))
-        sql = f"""SELECT WRB_tax, Description_en, Management_en, Description_es, Management_es,
-                         Description_ks, Management_ks, Description_fr, Management_fr
-                  FROM wrb_fao90_desc
-                  WHERE WRB_tax IN ({placeholders})"""
-
-        # Execute the query with the parameters
-        cur.execute(sql, tuple(WRB_Comp_List))
-        results = cur.fetchall()
+            # Create placeholders for the SQL IN clause
+            placeholders = ", ".join(["%s"] * len(WRB_Comp_List))
+            sql = f"""SELECT WRB_tax, Description_en, Management_en, Description_es, Management_es,
+                             Description_ks, Management_ks, Description_fr, Management_fr
+                      FROM wrb_fao90_desc
+                      WHERE WRB_tax IN ({placeholders})"""
+
+            # Execute the query with the parameters
+            cur.execute(sql, tuple(WRB_Comp_List))
+            results = cur.fetchall()
 
         # Convert the results to a pandas DataFrame
         data = pd.DataFrame(
@@ -381,18 +369,13 @@ def get_WRB_descriptions(WRB_Comp_List):
         )
 
         return data
-
     except Exception as err:
         logging.error(f"Error querying PostgreSQL: {err}")
         return None
 
-    finally:
-        if conn:
-            conn.close()
-
 
 # global only
-def getSG_descriptions(WRB_Comp_List):
+def getSG_descriptions(connection, WRB_Comp_List):
     """
     Fetch WRB descriptions from a PostgreSQL database using wrb2006_to_fao90
     and wrb_fao90_desc tables. Returns a pandas DataFrame with columns:
@@ -405,13 +388,10 @@ def getSG_descriptions(WRB_Comp_List):
         pandas.DataFrame or None if an error occurs.
     """
 
-    conn = None
     try:
-        # 1. Get a connection to your datastore (replace with your actual function):
-        conn = get_datastore_connection()
 
         def execute_query(query, params):
-            with conn.cursor() as cur:
+            with connection.cursor() as cur:
                 # Execute the query with the parameters
                 cur.execute(query, params)
                 return cur.fetchall()
@@ -482,7 +462,3 @@ def execute_query(query, params):
     except Exception as err:
         logging.error(f"Error querying PostgreSQL: {err}")
         return None
-
-    finally:
-        if conn:
-            conn.close()

soil_id/global_soil.py

@@ -27,12 +27,7 @@
 from scipy.stats import norm
 
 from .color import calculate_deltaE2000
-from .db import (
-    extract_hwsd2_data,
-    fetch_table_from_db,
-    get_WRB_descriptions,
-    getSG_descriptions,
-)
+from .db import extract_hwsd2_data, fetch_table_from_db, get_WRB_descriptions, getSG_descriptions
 from .services import get_soilgrids_classification_data, get_soilgrids_property_data
 from .utils import (
     adjust_depth_interval,
@@ -77,10 +72,11 @@ class SoilListOutputData:
 ##################################################################################################
 #                                 getSoilLocationBasedGlobal                                     #
 ##################################################################################################
-def list_soils_global(lon, lat):
+def list_soils_global(connection, lon, lat):
     # Extract HWSD2 Data
     try:
         hwsd2_data = extract_hwsd2_data(
+            connection,
             lon,
             lat,
             table_name="hwsdv2",
@@ -580,6 +576,7 @@ def convert_to_serializable(obj):
 #                                   rankPredictionGlobal                                     #
 ##############################################################################################
 def rank_soils_global(
+    connection,
     lon,
     lat,
     list_output_data: SoilListOutputData,
@@ -949,7 +946,7 @@ def rank_soils_global(
     ysf = []
 
     # Load color distribution data from NormDist2 (FAO90) table
-    rows = fetch_table_from_db("NormDist2")
+    rows = fetch_table_from_db(connection, "NormDist2")
     row_id = 0
     for row in rows:
         # row is a tuple; iterate over its values.

soil_id/tests/global/generate_bulk_test_results.py

@@ -20,7 +20,7 @@
 import traceback
 
 import pandas
-
+from soil_id.db import get_datastore_connection
 from soil_id.global_soil import list_soils_global, rank_soils_global
 
 test_data_df = pandas.read_csv(
@@ -36,7 +36,10 @@
 
 print(f"logging results to {result_file_name}")
 # buffering=1 is line buffering
-with open(result_file_name, "w", buffering=1) as result_file:
+with (
+    open(result_file_name, "w", buffering=1) as result_file,
+    get_datastore_connection() as connection,
+):
     result_agg = {}
 
     for pedon_key, pedon in pedons:
@@ -55,9 +58,10 @@
             start_time = time.perf_counter()
             try:
 
-                list_result = list_soils_global(lat=lat, lon=lon)
+                list_result = list_soils_global(connection=connection, lat=lat, lon=lon)
 
                 result_record["rank_result"] = rank_soils_global(
+                    connection=connection,
                     lat=lat,
                     lon=lon,
                     list_output_data=list_result,