Skip to content

Commit f8bfc57

Browse files
committed
fix: revision of global functions
1 parent 9ff26ea commit f8bfc57

File tree

5 files changed

+647
-469
lines changed

5 files changed

+647
-469
lines changed

soil_id/db.py

Lines changed: 90 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,11 @@
2525
# local libraries
2626
import soil_id.config
2727

28+
import os
29+
from dotenv import load_dotenv
30+
31+
# Load .env file
32+
load_dotenv()
2833

2934
def get_datastore_connection():
3035
"""
@@ -33,17 +38,25 @@ def get_datastore_connection():
3338
Returns:
3439
Connection object if successful, otherwise exits the program.
3540
"""
41+
conn = None # Initialize variable
3642
try:
43+
# conn = psycopg.connect(
44+
# host=os.getenv("DB_HOST"),
45+
# user=os.getenv("DB_USERNAME"),
46+
# password=os.getenv("DB_PASSWORD"),
47+
# dbname=os.getenv("DB_NAME"),
48+
# )
3749
conn = psycopg.connect(
3850
host=soil_id.config.DB_HOST,
3951
user=soil_id.config.DB_USERNAME,
4052
password=soil_id.config.DB_PASSWORD,
4153
dbname=soil_id.config.DB_NAME,
4254
)
55+
logging.info("Database connection successful.")
4356
return conn
4457
except Exception as err:
45-
logging.error(err)
46-
sys.exit(str(err))
58+
logging.error(f"Database connection failed: {err}")
59+
raise
4760

4861

4962
# us, global
@@ -53,6 +66,7 @@ def save_model_output(
5366
"""
5467
Save the output of the model to the 'landpks_soil_model' table.
5568
"""
69+
conn = None
5670
try:
5771
conn = get_datastore_connection()
5872
cur = conn.cursor()
@@ -87,6 +101,7 @@ def save_rank_output(record_id, model_version, rank_blob):
87101
"""
88102
Update the rank of the soil model in the 'landpks_soil_model' table.
89103
"""
104+
conn = None
90105
try:
91106
conn = get_datastore_connection()
92107
cur = conn.cursor()
@@ -110,6 +125,7 @@ def load_model_output(plot_id):
110125
"""
111126
Load model output based on plot ID and model version.
112127
"""
128+
conn = None
113129
try:
114130
conn = get_datastore_connection()
115131
cur = conn.cursor()
@@ -135,6 +151,7 @@ def save_soilgrids_output(plot_id, model_version, soilgrids_blob):
135151
"""
136152
Save the output of the soil grids to the 'landpks_soilgrids_model' table.
137153
"""
154+
conn = None
138155
try:
139156
conn = get_datastore_connection()
140157
cur = conn.cursor()
@@ -160,20 +177,25 @@ def get_WISE30sec_data(MUGLB_NEW_Select):
160177
"""
161178
Retrieve WISE 30 second data based on selected MUGLB_NEW values.
162179
"""
180+
if not MUGLB_NEW_Select: # Handle empty list case
181+
logging.warning("MUGLB_NEW_Select is empty. Returning empty DataFrame.")
182+
return pd.DataFrame() # Return an empty DataFrame
183+
184+
conn = None
163185
try:
164186
conn = get_datastore_connection()
165187
cur = conn.cursor()
166188

167189
# Create placeholders for the SQL IN clause
168190
placeholders = ", ".join(["%s"] * len(MUGLB_NEW_Select))
169-
sql = f"""SELECT MUGLB_NEW, COMPID, id, MU_GLOBAL, NEWSUID, SCID, PROP, CLAF,
170-
PRID, Layer, TopDep, BotDep, CFRAG, SDTO, STPC, CLPC, CECS,
171-
PHAQ, ELCO, SU_name, FAO_SYS
172-
FROM wise_soil_data
173-
WHERE MUGLB_NEW IN ({placeholders})"""
174-
175-
# Execute the query with the parameters
176-
cur.execute(sql, tuple(MUGLB_NEW_Select))
191+
sql_query = f"""SELECT MUGLB_NEW, COMPID, id, MU_GLOBAL, NEWSUID, SCID, PROP, CLAF,
192+
PRID, Layer, TopDep, BotDep, CFRAG, SDTO, STPC, CLPC, CECS,
193+
PHAQ, ELCO, SU_name, FAO_SYS
194+
FROM wise_soil_data
195+
WHERE MUGLB_NEW IN ({placeholders})"""
196+
197+
# Execute the query only if the list is non-empty
198+
cur.execute(sql_query, tuple(MUGLB_NEW_Select))
177199
results = cur.fetchall()
178200

179201
# Convert the results to a pandas DataFrame
@@ -214,12 +236,12 @@ def get_WISE30sec_data(MUGLB_NEW_Select):
214236
if conn:
215237
conn.close()
216238

217-
218239
# global
219240
def get_WRB_descriptions(WRB_Comp_List):
220241
"""
221242
Retrieve WRB descriptions based on provided WRB component list.
222243
"""
244+
conn = None
223245
try:
224246
conn = get_datastore_connection()
225247
cur = conn.cursor()
@@ -264,30 +286,68 @@ def get_WRB_descriptions(WRB_Comp_List):
264286

265287
# global only
266288
def getSG_descriptions(WRB_Comp_List):
289+
"""
290+
Fetch WRB descriptions from a PostgreSQL database using wrb2006_to_fao90
291+
and wrb_fao90_desc tables. Returns a pandas DataFrame with columns:
292+
[WRB_tax, Description_en, Management_en, Description_es, ...]
293+
294+
Args:
295+
WRB_Comp_List (list[str]): List of WRB_2006_Full values (e.g. ["Chernozem","Gleysol"]).
296+
297+
Returns:
298+
pandas.DataFrame or None if an error occurs.
299+
"""
300+
301+
conn = None
267302
try:
303+
# 1. Get a connection to your datastore (replace with your actual function):
268304
conn = get_datastore_connection()
269-
270-
# Execute a SQL query and return the results
305+
271306
def execute_query(query, params):
272307
with conn.cursor() as cur:
308+
# Execute the query with the parameters
273309
cur.execute(query, params)
274310
return cur.fetchall()
275-
276-
# First SQL query
277-
sql1 = """SELECT lu.WRB_1984_Full
278-
FROM wrb2006_to_fao90 AS lu
279-
WHERE lu.WRB_2006_Full IN %s"""
280-
names = execute_query(sql1, (tuple(WRB_Comp_List),))
281-
WRB_Comp_List = [item for t in names for item in t]
282-
283-
# Second SQL query
284-
sql2 = """SELECT WRB_tax, Description_en, Management_en, Description_es, Management_es,
285-
Description_ks, Management_ks, Description_fr, Management_fr
286-
FROM wrb_fao90_desc
287-
WHERE WRB_tax IN %s"""
288-
results = execute_query(sql2, (tuple(WRB_Comp_List),))
289-
290-
# Convert results to DataFrame
311+
312+
# 2. Map WRB_2006_Full -> WRB_1984_Full using wrb2006_to_fao90
313+
# Make sure we pass (tuple(WRB_Comp_List),) so psycopg2 can fill IN ('A','B','C')
314+
# Example: WHERE lu.WRB_2006_Full IN ('Chernozem','Gleysol',...)
315+
sql1 = """
316+
SELECT lu.WRB_1984_Full
317+
FROM wrb2006_to_fao90 AS lu
318+
WHERE lu.WRB_2006_Full = ANY(%s)
319+
"""
320+
names = execute_query(sql1, ([WRB_Comp_List],))
321+
322+
# Flatten from [(x,), (y,)] => [x, y]
323+
WRB_Comp_List_mapped = [item for (item,) in names]
324+
325+
if not WRB_Comp_List_mapped:
326+
# If no mapping found, return an empty DataFrame or None
327+
logging.warning("No mapped WRB_1984_Full names found for given WRB_2006_Full values.")
328+
return pd.DataFrame(columns=[
329+
"WRB_tax", "Description_en", "Management_en", "Description_es",
330+
"Management_es", "Description_ks", "Management_ks",
331+
"Description_fr", "Management_fr"
332+
])
333+
334+
# 3. Get descriptions from wrb_fao90_desc where WRB_tax IN ...
335+
sql2 = """
336+
SELECT WRB_tax,
337+
Description_en,
338+
Management_en,
339+
Description_es,
340+
Management_es,
341+
Description_ks,
342+
Management_ks,
343+
Description_fr,
344+
Management_fr
345+
FROM wrb_fao90_desc
346+
WHERE WRB_tax = ANY(%s)
347+
"""
348+
results = execute_query(sql2, ([WRB_Comp_List_mapped],))
349+
350+
# 4. Convert the raw query results to a DataFrame
291351
data = pd.DataFrame(
292352
results,
293353
columns=[
@@ -311,3 +371,4 @@ def execute_query(query, params):
311371
finally:
312372
if conn:
313373
conn.close()
374+

0 commit comments

Comments
 (0)