fix: revision of global functions

Author: jjmaynard

soil_id/db.py

@@ -25,6 +25,11 @@
 # local libraries
 import soil_id.config
 
+import os
+from dotenv import load_dotenv
+
+# Load .env file
+load_dotenv()
 
 def get_datastore_connection():
     """
@@ -33,17 +38,25 @@ def get_datastore_connection():
     Returns:
         Connection object if successful, otherwise exits the program.
     """
+    conn = None  # Initialize variable
     try:
+        # conn = psycopg.connect(
+        #     host=os.getenv("DB_HOST"),
+        #     user=os.getenv("DB_USERNAME"),
+        #     password=os.getenv("DB_PASSWORD"),
+        #     dbname=os.getenv("DB_NAME"),
+        # )
         conn = psycopg.connect(
             host=soil_id.config.DB_HOST,
             user=soil_id.config.DB_USERNAME,
             password=soil_id.config.DB_PASSWORD,
             dbname=soil_id.config.DB_NAME,
         )
+        logging.info("Database connection successful.")
         return conn
     except Exception as err:
-        logging.error(err)
-        sys.exit(str(err))
+        logging.error(f"Database connection failed: {err}")
+        raise
 
 
 # us, global
@@ -53,6 +66,7 @@ def save_model_output(
     """
     Save the output of the model to the 'landpks_soil_model' table.
     """
+    conn = None
     try:
         conn = get_datastore_connection()
         cur = conn.cursor()
@@ -87,6 +101,7 @@ def save_rank_output(record_id, model_version, rank_blob):
     """
     Update the rank of the soil model in the 'landpks_soil_model' table.
     """
+    conn = None
     try:
         conn = get_datastore_connection()
         cur = conn.cursor()
@@ -110,6 +125,7 @@ def load_model_output(plot_id):
     """
     Load model output based on plot ID and model version.
     """
+    conn = None
     try:
         conn = get_datastore_connection()
         cur = conn.cursor()
@@ -135,6 +151,7 @@ def save_soilgrids_output(plot_id, model_version, soilgrids_blob):
     """
     Save the output of the soil grids to the 'landpks_soilgrids_model' table.
     """
+    conn = None
     try:
         conn = get_datastore_connection()
         cur = conn.cursor()
@@ -160,20 +177,25 @@ def get_WISE30sec_data(MUGLB_NEW_Select):
     """
     Retrieve WISE 30 second data based on selected MUGLB_NEW values.
     """
+    if not MUGLB_NEW_Select:  # Handle empty list case
+        logging.warning("MUGLB_NEW_Select is empty. Returning empty DataFrame.")
+        return pd.DataFrame()  # Return an empty DataFrame
+
+    conn = None
     try:
         conn = get_datastore_connection()
         cur = conn.cursor()
 
         # Create placeholders for the SQL IN clause
         placeholders = ", ".join(["%s"] * len(MUGLB_NEW_Select))
-        sql = f"""SELECT MUGLB_NEW, COMPID, id, MU_GLOBAL, NEWSUID, SCID, PROP, CLAF,
-                         PRID, Layer, TopDep, BotDep,  CFRAG,  SDTO,  STPC,  CLPC, CECS,
-                         PHAQ, ELCO, SU_name, FAO_SYS
-                  FROM wise_soil_data
-                  WHERE MUGLB_NEW IN ({placeholders})"""
-
-        # Execute the query with the parameters
-        cur.execute(sql, tuple(MUGLB_NEW_Select))
+        sql_query = f"""SELECT MUGLB_NEW, COMPID, id, MU_GLOBAL, NEWSUID, SCID, PROP, CLAF,
+                               PRID, Layer, TopDep, BotDep, CFRAG, SDTO, STPC, CLPC, CECS,
+                               PHAQ, ELCO, SU_name, FAO_SYS
+                        FROM wise_soil_data
+                        WHERE MUGLB_NEW IN ({placeholders})"""
+
+        # Execute the query only if the list is non-empty
+        cur.execute(sql_query, tuple(MUGLB_NEW_Select))
         results = cur.fetchall()
 
         # Convert the results to a pandas DataFrame
@@ -214,12 +236,12 @@ def get_WISE30sec_data(MUGLB_NEW_Select):
         if conn:
             conn.close()
 
-
 # global
 def get_WRB_descriptions(WRB_Comp_List):
     """
     Retrieve WRB descriptions based on provided WRB component list.
     """
+    conn = None
     try:
         conn = get_datastore_connection()
         cur = conn.cursor()
@@ -264,30 +286,68 @@ def get_WRB_descriptions(WRB_Comp_List):
 
 # global only
 def getSG_descriptions(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:
+    [WRB_tax, Description_en, Management_en, Description_es, ...]
+    
+    Args:
+        WRB_Comp_List (list[str]): List of WRB_2006_Full values (e.g. ["Chernozem","Gleysol"]).
+    
+    Returns:
+        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()
-
-        # Execute a SQL query and return the results
+        
         def execute_query(query, params):
             with conn.cursor() as cur:
+                # Execute the query with the parameters
                 cur.execute(query, params)
                 return cur.fetchall()
-
-        # First SQL query
-        sql1 = """SELECT lu.WRB_1984_Full
-                  FROM wrb2006_to_fao90 AS lu
-                  WHERE lu.WRB_2006_Full IN %s"""
-        names = execute_query(sql1, (tuple(WRB_Comp_List),))
-        WRB_Comp_List = [item for t in names for item in t]
-
-        # Second SQL query
-        sql2 = """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 %s"""
-        results = execute_query(sql2, (tuple(WRB_Comp_List),))
-
-        # Convert results to DataFrame
+        
+        # 2. Map WRB_2006_Full -> WRB_1984_Full using wrb2006_to_fao90
+        #    Make sure we pass (tuple(WRB_Comp_List),) so psycopg2 can fill IN ('A','B','C')
+        #    Example: WHERE lu.WRB_2006_Full IN ('Chernozem','Gleysol',...)
+        sql1 = """
+            SELECT lu.WRB_1984_Full
+            FROM wrb2006_to_fao90 AS lu
+            WHERE lu.WRB_2006_Full = ANY(%s)
+        """
+        names = execute_query(sql1, ([WRB_Comp_List],))
+
+        # Flatten from [(x,), (y,)] => [x, y]
+        WRB_Comp_List_mapped = [item for (item,) in names]
+        
+        if not WRB_Comp_List_mapped:
+            # If no mapping found, return an empty DataFrame or None
+            logging.warning("No mapped WRB_1984_Full names found for given WRB_2006_Full values.")
+            return pd.DataFrame(columns=[
+                "WRB_tax", "Description_en", "Management_en", "Description_es",
+                "Management_es", "Description_ks", "Management_ks",
+                "Description_fr", "Management_fr"
+            ])
+
+        # 3. Get descriptions from wrb_fao90_desc where WRB_tax IN ...
+        sql2 = """
+            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 = ANY(%s)
+        """
+        results = execute_query(sql2, ([WRB_Comp_List_mapped],))
+        
+        # 4. Convert the raw query results to a DataFrame
         data = pd.DataFrame(
             results,
             columns=[
@@ -311,3 +371,4 @@ def execute_query(query, params):
     finally:
         if conn:
             conn.close()
+