fix connection reuse

app.py

@@ -5,7 +5,8 @@ import pandas as pd
 import logging
 
 from data_access import get_questions, get_source_finders, get_run_ids, get_baseline_rankers, \
-    get_unified_sources, get_source_text, calculate_cumulative_statistics_for_all_questions, get_metadata
+    get_unified_sources, get_source_text, calculate_cumulative_statistics_for_all_questions, get_metadata, \
+    get_async_connection
 
 logger = logging.getLogger(__name__)
 
@@ -20,7 +21,7 @@ baseline_ranker_options = []
 run_ids = []
 available_run_id_dict = {}
 finder_options = []
-previous_run_id = None
+previous_run_id = "initial_run"
 
 run_id_dropdown = None
 
@@ -29,13 +30,13 @@ run_id_dropdown = None
 # Initialize data in a single async function
 async def initialize_data():
     global questions, source_finders, questions_dict, source_finders_dict, question_options, finder_options, baseline_rankers_dict, source_finders_dict, baseline_ranker_options
+    async with get_async_connection() as conn:
+        # Get questions and source finders
+        questions = await get_questions(conn)
+        source_finders = await get_source_finders(conn)
+        baseline_rankers = await get_baseline_rankers(conn)
 
-    questions = await get_questions()
-    source_finders = await get_source_finders()
-
-    baseline_rankers = await get_baseline_rankers()
     baseline_rankers_dict = {f["name"]: f["id"] for f in baseline_rankers}
-
     # Convert to dictionaries for easier lookup
     questions_dict = {q["text"]: q["id"] for q in questions}
     baseline_rankers_dict = {f["name"]: f["id"] for f in baseline_rankers}
@@ -52,7 +53,7 @@ def update_sources_list(question_option, source_finder_id, run_id: str, baseline
     if evt:
         logger.info(f"event: {evt.target.elem_id}")
         if evt.target.elem_id == "run_id_dropdown" and (type(run_id) == list or run_id == previous_run_id):
-            return gr.skip(), gr.skip(), gr.skip(), gr.skip()
+            return gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip()
 
     if type(run_id) == str:
         previous_run_id = run_id
@@ -65,55 +66,56 @@ async def update_sources_list_async(question_option, source_finder_name, run_id,
     if not question_option:
         return gr.skip(), gr.skip(), gr.skip(), "No question selected", ""
     logger.info("processing update")
-    if type(baseline_ranker_name) == list:
-        baseline_ranker_name = baseline_ranker_name[0]
-
-    baseline_ranker_id_int = 1 if len(baseline_ranker_name) == 0 else baseline_rankers_dict.get(baseline_ranker_name)
+    async with get_async_connection() as conn:
+        if type(baseline_ranker_name) == list:
+            baseline_ranker_name = baseline_ranker_name[0]
 
-    if len(source_finder_name):
-        finder_id_int = source_finders_dict.get(source_finder_name)
-    else:
-        finder_id_int = None
+        baseline_ranker_id_int = 1 if len(baseline_ranker_name) == 0 else baseline_rankers_dict.get(baseline_ranker_name)
 
-    if question_option == "All questions":
-        if finder_id_int and type(run_id) == str:
-            run_id_int = available_run_id_dict.get(run_id)
-            all_stats = await calculate_cumulative_statistics_for_all_questions(run_id_int, baseline_ranker_id_int)
+        if len(source_finder_name):
+            finder_id_int = source_finders_dict.get(source_finder_name)
         else:
-            all_stats = None
-        return None, all_stats, gr.skip(), "Select Run Id and source finder to see results", ""
+            finder_id_int = None
 
-    # Extract question ID from selection
-    question_id = questions_dict.get(question_option)
+        if question_option == "All questions":
+            if finder_id_int and type(run_id) == str:
+                run_id_int = available_run_id_dict.get(run_id)
+                all_stats = await calculate_cumulative_statistics_for_all_questions(conn, run_id_int, baseline_ranker_id_int)
+            else:
+                all_stats = None
+            return None, all_stats, gr.skip(), "Select Run Id and source finder to see results", ""
 
-    available_run_id_dict = await get_run_ids(question_id, finder_id_int)
-    run_id_options = list(available_run_id_dict.keys())
-    if run_id not in run_id_options:
-        run_id = run_id_options[0]
+        # Extract question ID from selection
+        question_id = questions_dict.get(question_option)
 
-    run_id_int = available_run_id_dict.get(run_id)
+        available_run_id_dict = await get_run_ids(conn, question_id, finder_id_int)
+        run_id_options = list(available_run_id_dict.keys())
+        if run_id not in run_id_options:
+            run_id = run_id_options[0]
 
+        run_id_int = available_run_id_dict.get(run_id)
 
 
-    source_runs = None
-    stats = None
-    # Get source runs data
-    if finder_id_int:
-        source_runs, stats = await get_unified_sources(question_id, run_id_int, baseline_ranker_id_int)
-        # Create DataFrame for display
-        df = pd.DataFrame(source_runs)
 
-    if not source_runs:
-        return None, None, run_id_options, "No results found for the selected filters",
+        source_runs = None
+        stats = None
+        # Get source runs data
+        if finder_id_int:
+            source_runs, stats = await get_unified_sources(conn, question_id, run_id_int, baseline_ranker_id_int)
+            # Create DataFrame for display
+            df = pd.DataFrame(source_runs)
 
-    # Format table columns
-    columns_to_display = ['sugya_id', 'in_baseline', 'baseline_rank', 'in_source_run', 'source_run_rank', 'tractate',
-                          'folio', 'reason']
-    df_display = df[columns_to_display] if all(col in df.columns for col in columns_to_display) else df
+        if not source_runs:
+            return None, None, run_id_options, "No results found for the selected filters",
 
-    # CSV for download
-    # csv_data = df.to_csv(index=False)
-    metadata = await get_metadata(question_id, run_id_int)
+        # Format table columns
+        columns_to_display = ['sugya_id', 'in_baseline', 'baseline_rank', 'in_source_run', 'source_run_rank', 'tractate',
+                              'folio', 'reason']
+        df_display = df[columns_to_display] if all(col in df.columns for col in columns_to_display) else df
+
+        # CSV for download
+        # csv_data = df.to_csv(index=False)
+        metadata = await get_metadata(conn, question_id, run_id_int)
 
     result_message = f"Found {len(source_runs)} results"
     return df_display, stats, gr.Dropdown(choices=run_id_options, value=run_id), result_message, metadata
@@ -128,7 +130,8 @@ async def handle_row_selection_async(evt: gr.SelectData):
         # Get the ID from the selected row
         tractate_chunk_id = evt.row_value[0]
         # Get the source text
-        text = await get_source_text(tractate_chunk_id)
+        async with get_async_connection() as conn:
+            text = await get_source_text(conn, tractate_chunk_id)
         return text
     except Exception as e:
         return f"Error retrieving source text: {str(e)}"
@@ -248,7 +251,6 @@ async def main():
         question_dropdown.change(
             update_sources_list,
             inputs=[question_dropdown, source_finder_dropdown, run_id_dropdown, baseline_rankers_dropdown],
-            # outputs=[run_id_dropdown, results_table, result_text, download_button]
             outputs=[results_table, statistics_table, run_id_dropdown, result_text, metadata_text]
         )
 

data_access.py

@@ -30,85 +30,80 @@ async def get_async_connection(schema="talmudexplore"):
         await conn.close()
 
 
-async def get_questions():
-    async with get_async_connection() as conn:
-        questions = await conn.fetch("SELECT id, question_text FROM questions ORDER BY id")
-        return [{"id": q["id"], "text": q["question_text"]} for q in questions]
-
-async def get_metadata(question_id: int, source_finder_id_run_id: int):
-    async with get_async_connection() as conn:
-        metadata = await conn.fetchrow('''
-            SELECT metadata 
-            FROM source_finder_run_question_metadata sfrqm                           
-            WHERE sfrqm.question_id = $1 and sfrqm.source_finder_run_id = $2;
-        ''', question_id, source_finder_id_run_id)
-        if metadata is None:
-            return ""
-        return metadata.get('metadata')
+async def get_questions(conn: asyncpg.Connection):
+    questions = await conn.fetch("SELECT id, question_text FROM questions ORDER BY id")
+    return [{"id": q["id"], "text": q["question_text"]} for q in questions]
+
+async def get_metadata(conn: asyncpg.Connection, question_id: int, source_finder_id_run_id: int):
+    metadata = await conn.fetchrow('''
+        SELECT metadata 
+        FROM source_finder_run_question_metadata sfrqm                           
+        WHERE sfrqm.question_id = $1 and sfrqm.source_finder_run_id = $2;
+    ''', question_id, source_finder_id_run_id)
+    if metadata is None:
+        return ""
+    return metadata.get('metadata')
 
 
 # Get distinct source finders
-async def get_source_finders():
-    async with get_async_connection() as conn:
-        finders = await conn.fetch("SELECT id, source_finder_type as name FROM source_finders ORDER BY id")
-        return [{"id": f["id"], "name": f["name"]} for f in finders]
+async def get_source_finders(conn: asyncpg.Connection):
+    finders = await conn.fetch("SELECT id, source_finder_type as name FROM source_finders ORDER BY id")
+    return [{"id": f["id"], "name": f["name"]} for f in finders]
 
 
 # Get distinct run IDs for a question
-async def get_run_ids(question_id: int, source_finder_id: int):
-    async with get_async_connection() as conn:
-        query = """
-        select distinct sfr.description, srs.source_finder_run_id as run_id
-            from talmudexplore.source_run_results srs
-            join talmudexplore.source_finder_runs sfr on srs.source_finder_run_id = sfr.id
-            join talmudexplore.source_finders sf on sfr.source_finder_id = sf.id
-            where sfr.source_finder_id = $1
-            and srs.question_id = $2
-        """
-        run_ids = await conn.fetch(query, source_finder_id, question_id)
-        return {r["description"]:r["run_id"] for r in run_ids}
-
-
-async def get_baseline_rankers():
-    async with get_async_connection() as conn:
-        rankers = await conn.fetch("SELECT id, ranker FROM rankers ORDER BY id")
-        return [{"id": f["id"], "name": f["ranker"]} for f in rankers]
-
-async def calculate_baseline_vs_source_stats_for_question(baseline_sources , source_runs_sources):
+async def get_run_ids(conn: asyncpg.Connection, question_id: int, source_finder_id: int):
+    query = """
+    select distinct sfr.description, srs.source_finder_run_id as run_id
+        from talmudexplore.source_run_results srs
+        join talmudexplore.source_finder_runs sfr on srs.source_finder_run_id = sfr.id
+        join talmudexplore.source_finders sf on sfr.source_finder_id = sf.id
+        where sfr.source_finder_id = $1
+        and srs.question_id = $2
+    """
+    run_ids = await conn.fetch(query, source_finder_id, question_id)
+    return {r["description"]:r["run_id"] for r in run_ids}
+
+
+async def get_baseline_rankers(conn: asyncpg.Connection):
+    rankers = await conn.fetch("SELECT id, ranker FROM rankers ORDER BY id")
+    return [{"id": f["id"], "name": f["ranker"]} for f in rankers]
+
+async def calculate_baseline_vs_source_stats_for_question(conn: asyncpg.Connection, baseline_sources , source_runs_sources):
     # for a given question_id and source_finder_id and run_id calculate the baseline vs source stats
     # e.g. overlap, high ranked overlap, etc.
-    async with get_async_connection() as conn:
-        actual_sources_set = {s["id"] for s in source_runs_sources}
-        baseline_sources_set = {s["id"] for s in baseline_sources}
-
-        # Calculate overlap
-        overlap = actual_sources_set.intersection(baseline_sources_set)
-        # only_in_1 = actual_sources_set - baseline_sources_set
-        # only_in_2 = baseline_sources_set - actual_sources_set
-
-        # Calculate high-ranked overlap (rank >= 4)
-        actual_high_ranked = {s["id"] for s in source_runs_sources if int(s["source_rank"]) >= 4}
-        baseline_high_ranked = {s["id"] for s in baseline_sources if int(s["baseline_rank"]) >= 4}
-
-        high_ranked_overlap = actual_high_ranked.intersection(baseline_high_ranked)
-
-        results = {
-            "total_baseline_sources": len(baseline_sources),
-            "total_found_sources": len(source_runs_sources),
-            "overlap_count": len(overlap),
-            "overlap_percentage": round(len(overlap) * 100 / max(len(actual_sources_set), len(baseline_sources_set)),
-                                        2) if max(len(actual_sources_set), len(baseline_sources_set)) > 0 else 0,
-            "num_high_ranked_baseline_sources": len(baseline_high_ranked),
-            "num_high_ranked_found_sources": len(actual_high_ranked),
-            "high_ranked_overlap_count": len(high_ranked_overlap),
-            "high_ranked_overlap_percentage": round(len(high_ranked_overlap) * 100 / max(len(actual_high_ranked), len(baseline_high_ranked)), 2) if max(len(actual_high_ranked), len(baseline_high_ranked)) > 0 else 0
-        }
-        #convert results to dataframe
-        results_df = pd.DataFrame([results])
-        return results_df
 
-
-async def calculate_cumulative_statistics_for_all_questions(source_finder_run_id: int, ranker_id: int):
+    actual_sources_set = {s["id"] for s in source_runs_sources}
+    baseline_sources_set = {s["id"] for s in baseline_sources}
+
+    # Calculate overlap
+    overlap = actual_sources_set.intersection(baseline_sources_set)
+    # only_in_1 = actual_sources_set - baseline_sources_set
+    # only_in_2 = baseline_sources_set - actual_sources_set
+
+    # Calculate high-ranked overlap (rank >= 4)
+    actual_high_ranked = {s["id"] for s in source_runs_sources if int(s["source_rank"]) >= 4}
+    baseline_high_ranked = {s["id"] for s in baseline_sources if int(s["baseline_rank"]) >= 4}
+
+    high_ranked_overlap = actual_high_ranked.intersection(baseline_high_ranked)
+
+    results = {
+        "total_baseline_sources": len(baseline_sources),
+        "total_found_sources": len(source_runs_sources),
+        "overlap_count": len(overlap),
+        "overlap_percentage": round(len(overlap) * 100 / max(len(actual_sources_set), len(baseline_sources_set)),
+                                    2) if max(len(actual_sources_set), len(baseline_sources_set)) > 0 else 0,
+        "num_high_ranked_baseline_sources": len(baseline_high_ranked),
+        "num_high_ranked_found_sources": len(actual_high_ranked),
+        "high_ranked_overlap_count": len(high_ranked_overlap),
+        "high_ranked_overlap_percentage": round(len(high_ranked_overlap) * 100 / max(len(actual_high_ranked), len(baseline_high_ranked)), 2) if max(len(actual_high_ranked), len(baseline_high_ranked)) > 0 else 0
+    }
+    #convert results to dataframe
+    results_df = pd.DataFrame([results])
+    return results_df
+
+
+async def calculate_cumulative_statistics_for_all_questions(conn: asyncpg.Connection, source_finder_run_id: int, ranker_id: int):
     """
     Calculate cumulative statistics across all questions for a specific source finder, run, and ranker.
 
@@ -119,83 +114,75 @@ async def calculate_cumulative_statistics_for_all_questions(source_finder_run_id
     Returns:
         pd.DataFrame: DataFrame containing aggregated statistics
     """
-    async with get_async_connection() as conn:
-        # Get all questions
-        query = "SELECT id FROM questions ORDER BY id"
-        questions = await conn.fetch(query)
-        question_ids = [q["id"] for q in questions]
-
-        # Initialize aggregates
-        total_baseline_sources = 0
-        total_found_sources = 0
-        total_overlap = 0
-        total_high_ranked_baseline = 0
-        total_high_ranked_found = 0
-        total_high_ranked_overlap = 0
-
-        # Process each question
-        valid_questions = 0
-        for question_id in question_ids:
-            try:
-                # Get unified sources for this question
-                stats, sources = await get_stats(conn, question_id, ranker_id, source_finder_run_id)
-
-                if sources and len(sources) > 0:
-                    valid_questions += 1
-                    stats_dict = stats.iloc[0].to_dict()
-
-                    # Add to running totals
-                    total_baseline_sources += stats_dict.get('total_baseline_sources', 0)
-                    total_found_sources += stats_dict.get('total_found_sources', 0)
-                    total_overlap += stats_dict.get('overlap_count', 0)
-                    total_high_ranked_baseline += stats_dict.get('num_high_ranked_baseline_sources', 0)
-                    total_high_ranked_found += stats_dict.get('num_high_ranked_found_sources', 0)
-                    total_high_ranked_overlap += stats_dict.get('high_ranked_overlap_count', 0)
-            except Exception as e:
-                # Skip questions with errors
-                continue
-
-        # Calculate overall percentages
-        overlap_percentage = round(total_overlap * 100 / max(total_baseline_sources, total_found_sources), 2) \
-            if max(total_baseline_sources, total_found_sources) > 0 else 0
-
-        high_ranked_overlap_percentage = round(
-            total_high_ranked_overlap * 100 / max(total_high_ranked_baseline, total_high_ranked_found), 2) \
-            if max(total_high_ranked_baseline, total_high_ranked_found) > 0 else 0
-
-        # Compile results
-        cumulative_stats = {
-            "total_questions_analyzed": valid_questions,
-            "total_baseline_sources": total_baseline_sources,
-            "total_found_sources": total_found_sources,
-            "total_overlap_count": total_overlap,
-            "overall_overlap_percentage": overlap_percentage,
-            "total_high_ranked_baseline_sources": total_high_ranked_baseline,
-            "total_high_ranked_found_sources": total_high_ranked_found,
-            "total_high_ranked_overlap_count": total_high_ranked_overlap,
-            "overall_high_ranked_overlap_percentage": high_ranked_overlap_percentage,
-            "avg_baseline_sources_per_question": round(total_baseline_sources / valid_questions,
-                                                       2) if valid_questions > 0 else 0,
-            "avg_found_sources_per_question": round(total_found_sources / valid_questions,
-                                                    2) if valid_questions > 0 else 0
-        }
-
-        return pd.DataFrame([cumulative_stats])
-
-
-async def get_unified_sources(question_id: int, source_finder_run_id: int, ranker_id: int):
+    # Get all questions
+    query = "SELECT id FROM questions ORDER BY id"
+    questions = await conn.fetch(query)
+    question_ids = [q["id"] for q in questions]
+
+    # Initialize aggregates
+    total_baseline_sources = 0
+    total_found_sources = 0
+    total_overlap = 0
+    total_high_ranked_baseline = 0
+    total_high_ranked_found = 0
+    total_high_ranked_overlap = 0
+
+    # Process each question
+    valid_questions = 0
+    for question_id in question_ids:
+        try:
+            # Get unified sources for this question
+            sources, stats = await get_unified_sources(conn, question_id, ranker_id, source_finder_run_id)
+
+            if sources and len(sources) > 0:
+                valid_questions += 1
+                stats_dict = stats.iloc[0].to_dict()
+
+                # Add to running totals
+                total_baseline_sources += stats_dict.get('total_baseline_sources', 0)
+                total_found_sources += stats_dict.get('total_found_sources', 0)
+                total_overlap += stats_dict.get('overlap_count', 0)
+                total_high_ranked_baseline += stats_dict.get('num_high_ranked_baseline_sources', 0)
+                total_high_ranked_found += stats_dict.get('num_high_ranked_found_sources', 0)
+                total_high_ranked_overlap += stats_dict.get('high_ranked_overlap_count', 0)
+        except Exception as e:
+            # Skip questions with errors
+            continue
+
+    # Calculate overall percentages
+    overlap_percentage = round(total_overlap * 100 / max(total_baseline_sources, total_found_sources), 2) \
+        if max(total_baseline_sources, total_found_sources) > 0 else 0
+
+    high_ranked_overlap_percentage = round(
+        total_high_ranked_overlap * 100 / max(total_high_ranked_baseline, total_high_ranked_found), 2) \
+        if max(total_high_ranked_baseline, total_high_ranked_found) > 0 else 0
+
+    # Compile results
+    cumulative_stats = {
+        "total_questions_analyzed": valid_questions,
+        "total_baseline_sources": total_baseline_sources,
+        "total_found_sources": total_found_sources,
+        "total_overlap_count": total_overlap,
+        "overall_overlap_percentage": overlap_percentage,
+        "total_high_ranked_baseline_sources": total_high_ranked_baseline,
+        "total_high_ranked_found_sources": total_high_ranked_found,
+        "total_high_ranked_overlap_count": total_high_ranked_overlap,
+        "overall_high_ranked_overlap_percentage": high_ranked_overlap_percentage,
+        "avg_baseline_sources_per_question": round(total_baseline_sources / valid_questions,
+                                                   2) if valid_questions > 0 else 0,
+        "avg_found_sources_per_question": round(total_found_sources / valid_questions,
+                                                2) if valid_questions > 0 else 0
+    }
+
+    return pd.DataFrame([cumulative_stats])
+
+
+async def get_unified_sources(conn: asyncpg.Connection, question_id: int, source_finder_run_id: int, ranker_id: int):
     """
     Create unified view of sources from both baseline_sources and source_runs
     with indicators of where each source appears and their respective ranks.
     """
-    async with get_async_connection() as conn:
-        stats_df, unified_results = await get_stats(conn, question_id, ranker_id, source_finder_run_id)
-
-        return unified_results, stats_df
 
-
-async def get_stats(conn, question_id, ranker_id, source_finder_run_id):
-    # Get sources from source_runs
     query_runs = """
                  SELECT tb.tractate_chunk_id as id,
                         sr.rank              as source_rank,
@@ -217,7 +204,7 @@ async def get_stats(conn, question_id, ranker_id, source_finder_run_id):
                        AND bs.ranker_id = $2 
                      """
     baseline_sources = await conn.fetch(query_baseline, question_id, ranker_id)
-    stats_df = await calculate_baseline_vs_source_stats_for_question(baseline_sources, source_runs)
+    stats_df = await calculate_baseline_vs_source_stats_for_question(conn, baseline_sources, source_runs)
     # Convert to dictionaries for easier lookup
     source_runs_dict = {s["id"]: dict(s) for s in source_runs}
     baseline_dict = {s["id"]: dict(s) for s in baseline_sources}
@@ -244,21 +231,22 @@ async def get_stats(conn, question_id, ranker_id, source_finder_run_id):
             "metadata": source_runs_dict.get(sugya_id, {}).get("metadata", "")
         }
         unified_results.append(result)
-    return stats_df, unified_results
+
+    return unified_results, stats_df
 
 
-async def get_source_text(tractate_chunk_id: int):
+async def get_source_text(conn: asyncpg.Connection, tractate_chunk_id: int):
     """
     Retrieves the text content for a given tractate chunk ID.
     """
-    async with get_async_connection() as conn:
-        query = """
-        SELECT tb.text_with_nikud as text
-        FROM talmud_bavli tb
-        WHERE tb.tractate_chunk_id = $1
-        """
-        result = await conn.fetchrow(query, tractate_chunk_id)
-        return result["text"] if result else "Source text not found"
+
+    query = """
+    SELECT tb.text_with_nikud as text
+    FROM talmud_bavli tb
+    WHERE tb.tractate_chunk_id = $1
+    """
+    result = await conn.fetchrow(query, tractate_chunk_id)
+    return result["text"] if result else "Source text not found"
 
 def get_pg_sync_connection(schema="talmudexplore"):
     conn = psycopg2.connect(dbname=os.getenv("pg_dbname"),

eval_tables.py

@@ -92,6 +92,11 @@ def create_eval_database():
     );
     ''')
 
+    cursor.execute('''alter table source_run_results
+                        add constraint source_run_results_pk
+                            unique (source_finder_run_id, question_id, sugya_id);
+    ''')
+
     conn.commit()
     conn.close()
 

tests/test_db_layer.py

@@ -1,13 +1,15 @@
 import pandas as pd
 import pytest
 
-from data_access import calculate_cumulative_statistics_for_all_questions, get_metadata, get_run_ids
+from data_access import calculate_cumulative_statistics_for_all_questions, get_metadata, get_run_ids, \
+    get_async_connection
 from data_access import get_unified_sources
 
 
 @pytest.mark.asyncio
 async def test_get_unified_sources():
-    results, stats = await get_unified_sources(2, 2, 1, 1)
+    async with get_async_connection() as conn:
+        results, stats = await get_unified_sources(conn,2, 2, 1)
     assert results is not None
     assert stats is not None
 
@@ -23,12 +25,12 @@ async def test_get_unified_sources():
 @pytest.mark.asyncio
 async def test_calculate_cumulative_statistics_for_all_questions():
     # Test with known source_finder_id, run_id, and ranker_id
-    source_finder_id = 2
-    run_id = 1
+    source_finder_run_id = 2
     ranker_id = 1
 
     # Call the function to test
-    result = await calculate_cumulative_statistics_for_all_questions(source_finder_id, run_id, ranker_id)
+    async with get_async_connection() as conn:
+        result = await calculate_cumulative_statistics_for_all_questions(conn, source_finder_run_id, ranker_id)
 
     # Check basic structure of results
     assert isinstance(result, pd.DataFrame), "Result should be a pandas DataFrame"
@@ -65,12 +67,12 @@ async def test_calculate_cumulative_statistics_for_all_questions():
 @pytest.mark.asyncio
 async def test_get_metadata_none_returned():
     # Test with known source_finder_id, run_id, and ranker_id
-    source_finder_id = 1
-    run_id = 1
+    source_finder_run_id = 1
     question_id = 1
 
     # Call the function to test
-    result = await get_metadata(question_id, source_finder_id, run_id)
+    async with get_async_connection() as conn:
+        result = await get_metadata(conn, question_id, source_finder_run_id)
 
     assert result == "", "Should return empty string when no metadata is found"
 
@@ -81,7 +83,8 @@ async def test_get_metadata():
     question_id = 1
 
     # Call the function to test
-    result = await get_metadata(question_id, source_finder_run_id)
+    async with get_async_connection() as conn:
+        result = await get_metadata(conn, question_id, source_finder_run_id)
 
     assert result is not None, "Should return metadata when it exists"
 
@@ -93,16 +96,17 @@ async def test_get_run_ids():
     source_finder_id = 2  # Using a source finder ID that exists in the test database
 
     # Call the function to test
-    result = await get_run_ids(question_id, source_finder_id)
+    async with get_async_connection() as conn:
+        result = await get_run_ids(conn, question_id, source_finder_id)
 
-    # Verify the result is a dictionary
-    assert isinstance(result, dict), "Result should be a dictionary"
+        # Verify the result is a dictionary
+        assert isinstance(result, dict), "Result should be a dictionary"
 
-    # Check that the dictionary is not empty (assuming there are run IDs for this question/source finder)
-    assert len(result) > 0, "Should return at least one run ID"
+        # Check that the dictionary is not empty (assuming there are run IDs for this question/source finder)
+        assert len(result) > 0, "Should return at least one run ID"
 
-    # Test with a non-existent question_id
-    non_existent_question_id = 9999
-    empty_result = await get_run_ids(non_existent_question_id, source_finder_id)
+        # Test with a non-existent question_id
+        non_existent_question_id = 9999
+        empty_result = await get_run_ids(conn, non_existent_question_id, source_finder_id)
     assert isinstance(empty_result, dict), "Should return an empty dictionary for non-existent question"
     assert len(empty_result) == 0, "Should return empty dictionary for non-existent question"