feat: pass db connection for global as param

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

Author: garobrik

soil_id/db.py

@@ -160,7 +160,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.
@@ -177,7 +177,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"""
@@ -220,7 +220,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.
@@ -234,26 +234,24 @@ 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)
+    # 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,
@@ -344,7 +342,7 @@ def extract_hwsd2_data(lon, lat, buffer_dist, table_name):
         # """
 
         # Use GeoPandas to execute the main query and load results into a GeoDataFrame.
-        hwsd = gpd.read_postgis(main_query, conn, geom_col="geom")
+        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.
@@ -354,7 +352,7 @@ def extract_hwsd2_data(lon, lat, buffer_dist, table_name):
         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)
+        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()
@@ -365,46 +363,37 @@ def extract_hwsd2_data(lon, lat, buffer_dist, table_name):
 
 
 # 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()
-
-        # 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})"""
+        with connection.cursor() as cur:
 
-        # 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(
@@ -423,18 +412,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:
@@ -447,13 +431,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()
@@ -524,7 +505,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",
@@ -444,6 +440,7 @@ def list_soils_global(lon, lat):
 
     # Merge component descriptions
     WRB_Comp_Desc = get_WRB_descriptions(
+        connection,
         mucompdata_cond_prob["compname_grp"].drop_duplicates().tolist()
     )
 
@@ -582,6 +579,7 @@ def convert_to_serializable(obj):
 #                                   rankPredictionGlobal                                     #
 ##############################################################################################
 def rank_soils_global(
+    connection,
     lon,
     lat,
     list_output_data: SoilListOutputData,
@@ -949,7 +947,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:
@@ -54,9 +57,11 @@
         else:
             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,

soil_id/tests/global/test_global.py

@@ -17,6 +17,7 @@
 import time
 import pytest
 
+from soil_id.db import get_datastore_connection
 from soil_id.global_soil import list_soils_global, rank_soils_global, sg_list
 
 test_locations = [
@@ -28,21 +29,22 @@
 
 @pytest.mark.skip
 def test_soil_location():
-    for item in test_locations:
-        logging.info(f"Testing {item['lon']}, {item['lat']}")
-        start_time = time.perf_counter()
-        list_soils_result = list_soils_global(item["lon"], item["lat"])
-        logging.info(f"...time: {(time.perf_counter() - start_time):.2f}s")
-        rank_soils_global(
-            item["lon"],
-            item["lat"],
-            list_output_data=list_soils_result,
-            soilHorizon=["Loam"],
-            topDepth=[15],
-            bottomDepth=[45],
-            rfvDepth=[20],
-            lab_Color=[[41.23035939, 3.623018224, 13.27654356]],
-            bedrock=None,
-            cracks=None,
-        )
-        sg_list(item["lon"], item["lat"])
+    with get_datastore_connection() as connection:
+        for item in test_locations:
+            logging.info(f"Testing {item['lon']}, {item['lat']}")
+            start_time = time.perf_counter()
+            list_soils_result = list_soils_global(connection, item["lon"], item["lat"])
+            logging.info(f"...time: {(time.perf_counter()-start_time):.2f}s")
+            rank_soils_result = rank_soils_global(
+                connection,
+                item["lon"],
+                item["lat"],
+                list_output_data=list_soils_result,
+                soilHorizon=["Loam"],
+                horizonDepth=[15],
+                rfvDepth=[20],
+                lab_Color=[[41.23035939, 3.623018224, 13.27654356]],
+                bedrock=None,
+                cracks=None,
+            )
+            sg_soils_result = sg_list(item["lon"], item["lat"])