fix(cache): guard empty repolist and add bounded cache polling

8Knot/cache_manager/cache_facade.py

@@ -22,6 +22,7 @@
 """
 
 import logging
+import time
 from uuid import uuid4
 import psycopg2 as pg
 from psycopg2.extras import execute_values
@@ -57,10 +58,11 @@ def cache_query_results(
         client_pagination (int, optional): _description_. Defaults to 2000.
     """
     logging.warning(f"{target_table} -- CQR CACHE_QUERY_RESULTS BEGIN")
-    with pg.connect(
+    augur_conn = pg.connect(
         db_connection_string,
         options=f"-c search_path={env_augur_schema}",
-    ) as augur_conn:
+    )
+    try:
         with augur_conn.cursor(name=f"{target_table}-{uuid4()}") as augur_cur:
             # set number of rows we want from primary db at a time
             augur_cur.itersize = server_pagination
@@ -72,7 +74,8 @@ def cache_query_results(
 
             logging.warning(f"{target_table} -- CQR STARTING TRANSACTION")
             # connect to cache
-            with pg.connect(cache_cx_string) as cache_conn:
+            cache_conn = pg.connect(cache_cx_string)
+            try:
                 logging.warning(f"{target_table} -- CQR COMPOSING SQL")
                 # compose SQL w/ table name
                 # ref: https://www.psycopg.org/docs/sql.html
@@ -112,29 +115,33 @@ def cache_query_results(
 
                 logging.warning(f"{target_table} -- CQR COMMITTING TRANSACTION")
                 cache_conn.commit()
+            finally:
+                cache_conn.close()
 
         # don't need to commit on primary db
         logging.warning(f"{target_table} -- CQR SUCCESS")
+    finally:
+        augur_conn.close()
 
 
-def get_uncached(func_name: str, repolist: list[int]) -> list[int]:  # or None
+def get_uncached(func_name: str, repolist: list[int]) -> list[int]:
     """
     Checks bookkeeping data to find, for a given querying function, which
     repos' data are noted as present in cache vs. those that arent.
 
     Returns a list of repos that AREN'T resident in cache.
     """
-    with pg.connect(cache_cx_string) as cache_conn:
+    if not repolist:
+        return []
+
+    cache_conn = pg.connect(cache_cx_string)
+    try:
         with cache_conn.cursor() as cache_cur:
-            composed_query = pg_sql.SQL(
-                """
+            composed_query = pg_sql.SQL("""
                 SELECT cb.repo_id
                 FROM cache_bookkeeping cb
                 WHERE cb.cache_func = '{cache_func_name}' AND cb.repo_id in %s
-                """.format(
-                    cache_func_name=func_name
-                )
-            ).as_string(cache_conn)
+                """.format(cache_func_name=func_name)).as_string(cache_conn)
 
             # exec query
             cache_cur.execute(query=composed_query, vars=(tuple(repolist),))
@@ -144,13 +151,51 @@ def get_uncached(func_name: str, repolist: list[int]) -> list[int]:  # or None
 
             # process list of single-value tuples to get list of values.
             # looks like: [(val1,), (val2,), ...]
-            already_cached: set[int] = set([v[0] for v in already_cached])
+            already_cached_set: set[int] = set([v[0] for v in already_cached])
 
             # repos that are already cached will be removed from repolist set,
             # leaving uncached remaining.
-            not_cached: list[int] = list(set(repolist) - already_cached)
+            not_cached: list[int] = list(set(repolist) - already_cached_set)
 
             return not_cached
+    finally:
+        cache_conn.close()
+
+
+def wait_for_cache(
+    func_name: str,
+    repolist: list[int],
+    timeout: int = 60,
+    caller: str = "",
+) -> None:
+    """Wait for data to appear in cache with bounded polling and backoff.
+
+    Replaces the per-visualization `while not_cached := get_uncached(...): time.sleep(0.5)`
+    polling loops. Uses progressive backoff (0.5s -> 1.0s -> 2.0s capped at 5.0s) to reduce
+    load on postgres-cache while queries are running.
+
+    Args:
+        func_name: name of the query function (used for bookkeeping lookup).
+        repolist: repo IDs expected in cache.
+        timeout: max seconds to wait before raising TimeoutError.
+        caller: optional label for log messages.
+    """
+    if not repolist:
+        return
+
+    label = caller or func_name
+    delay = 0.5
+    max_delay = 5.0
+    deadline = time.monotonic() + timeout
+
+    while not_cached := get_uncached(func_name=func_name, repolist=repolist):
+        if time.monotonic() >= deadline:
+            raise TimeoutError(
+                f"{label}: timed out after {timeout}s waiting for {len(not_cached)} repos in {func_name}"
+            )
+        logging.warning(f"{label} - WAITING ON {len(not_cached)} REPOS FOR {func_name}")
+        time.sleep(delay)
+        delay = min(delay * 2, max_delay)
 
 
 def caching_wrapper(func_name: str, query: str, repolist: list[int], n_repolist_uses=1) -> None:
@@ -173,10 +218,14 @@ def caching_wrapper(func_name: str, query: str, repolist: list[int], n_repolist_
     try:
         # STEP 1: Which repos need to be queried for?
         #           some might already be in cache.
-        uncached_repos: list[int] | None = get_uncached(func_name=func_name, repolist=repolist)
+        if not repolist:
+            logging.warning(f"{func_name} COLLECTION - NO REPOS REQUESTED")
+            return
+
+        uncached_repos: list[int] = get_uncached(func_name=func_name, repolist=repolist)
         if not uncached_repos:
             logging.warning(f"{func_name} COLLECTION - ALL REQUESTED REPOS IN CACHE")
-            return 0
+            return
         else:
             logging.warning(f"{func_name} COLLECTION - CACHING {len(uncached_repos)} NEW REPOS")
 
@@ -211,27 +260,31 @@ def retrieve_from_cache(
     Results are retrieved by a DataFrame, so column names
     may need to be overridden by calling function.
     """
+    if not repolist:
+        logging.warning(f"{tablename} - EMPTY REPOLIST PROVIDED, RETURNING EMPTY DATAFRAME")
+        return pd.DataFrame()
 
     # GET ALL DATA FROM POSTGRES CACHE
-    df = None
-    with pg.connect(cache_cx_string) as cache_conn:
+    cache_conn = pg.connect(cache_cx_string)
+    try:
         with cache_conn.cursor() as cache_cur:
+            composed_query = pg_sql.SQL("SELECT * FROM {tablename} t WHERE t.repo_id IN %s;").format(
+                tablename=pg_sql.Identifier(tablename)
+            )
+
             cache_cur.execute(
-                """
-                SELECT *
-                FROM {tablename} t
-                WHERE t.repo_id IN %s;
-                """.format(
-                    tablename=tablename
-                ),
+                composed_query,
                 (tuple(repolist),),
             )
 
             logging.warning(f"{tablename} - LOADING DATA FROM CACHE")
+            rows = cache_cur.fetchall()
             df = pd.DataFrame(
-                cache_cur.fetchall(),
+                rows,
                 # get df column names from the database columns
-                columns=[desc[0] for desc in cache_cur.description],
+                columns=[desc[0] for desc in cache_cur.description] if cache_cur.description else [],
             )
             logging.warning(f"{tablename} - DATA LOADED - {df.shape} rows,cols")
             return df
+    finally:
+        cache_conn.close()

pyproject.toml

@@ -37,3 +37,7 @@ packages = []
 
 [tool.setuptools.dynamic]
 version = { attr = "metadata.__version__" }
+
+[tool.pytest.ini_options]
+pythonpath = ["8Knot"]
+

tests/test_cache_facade.py

@@ -0,0 +1,103 @@
+import os
+
+# Set dummy environment variables required by cache_manager.cx_common at import time
+os.environ.setdefault("AUGUR_DATABASE", "test_augur")
+os.environ.setdefault("AUGUR_HOST", "localhost")
+os.environ.setdefault("AUGUR_PASSWORD", "test_pass")
+os.environ.setdefault("AUGUR_PORT", "5432")
+os.environ.setdefault("AUGUR_USERNAME", "test_user")
+
+import pytest
+from unittest.mock import MagicMock, patch
+import pandas as pd
+import cache_manager.cache_facade as cf
+
+
+def test_get_uncached_empty_repolist():
+    """Verify get_uncached returns empty list immediately for empty or None repolist."""
+    with patch("cache_manager.cache_facade.pg.connect") as mock_connect:
+        assert cf.get_uncached("test_func", []) == []
+        assert cf.get_uncached("test_func", None) == []
+        mock_connect.assert_not_called()
+
+
+def test_get_uncached_success():
+    """Verify get_uncached queries bookkeeping and closes connection properly."""
+    mock_conn = MagicMock()
+    mock_cur = MagicMock()
+    mock_conn.cursor.return_value.__enter__.return_value = mock_cur
+    mock_cur.fetchall.return_value = [(101,)]
+
+    with patch("cache_manager.cache_facade.pg.connect", return_value=mock_conn):
+        uncached = cf.get_uncached("test_func", [101, 102])
+
+    assert uncached == [102]
+    mock_conn.close.assert_called_once()
+    mock_cur.execute.assert_called_once()
+
+
+def test_retrieve_from_cache_empty_repolist():
+    """Verify retrieve_from_cache returns empty DataFrame immediately for empty or None repolist."""
+    with patch("cache_manager.cache_facade.pg.connect") as mock_connect:
+        df_empty = cf.retrieve_from_cache("test_table", [])
+        assert isinstance(df_empty, pd.DataFrame)
+        assert df_empty.empty
+
+        df_none = cf.retrieve_from_cache("test_table", None)
+        assert isinstance(df_none, pd.DataFrame)
+        assert df_none.empty
+
+        mock_connect.assert_not_called()
+
+
+def test_retrieve_from_cache_success():
+    """Verify retrieve_from_cache formats query with Identifier, extracts data, and closes connection."""
+    mock_conn = MagicMock()
+    mock_cur = MagicMock()
+    mock_conn.cursor.return_value.__enter__.return_value = mock_cur
+    mock_cur.description = [("repo_id",), ("commit_count",)]
+    mock_cur.fetchall.return_value = [(101, 42), (102, 17)]
+
+    with patch("cache_manager.cache_facade.pg.connect", return_value=mock_conn):
+        df = cf.retrieve_from_cache("test_table", [101, 102])
+
+    assert isinstance(df, pd.DataFrame)
+    assert len(df) == 2
+    assert list(df.columns) == ["repo_id", "commit_count"]
+    mock_conn.close.assert_called_once()
+
+
+def test_wait_for_cache_empty_repolist():
+    """Verify wait_for_cache returns immediately for empty or None repolist."""
+    with patch("cache_manager.cache_facade.get_uncached") as mock_get_uncached:
+        cf.wait_for_cache("test_func", [])
+        cf.wait_for_cache("test_func", None)
+        mock_get_uncached.assert_not_called()
+
+
+def test_wait_for_cache_success():
+    """Verify wait_for_cache loops until get_uncached returns empty list."""
+    with patch("cache_manager.cache_facade.get_uncached", side_effect=[[101], []]):
+        with patch("time.sleep") as mock_sleep:
+            cf.wait_for_cache("test_func", [101], timeout=10)
+            mock_sleep.assert_called_once()
+
+
+def test_wait_for_cache_timeout():
+    """Verify wait_for_cache raises TimeoutError when deadline is exceeded."""
+    with patch("cache_manager.cache_facade.get_uncached", return_value=[101]):
+        with patch("time.sleep"):
+            with pytest.raises(TimeoutError, match="timed out"):
+                cf.wait_for_cache("test_func", [101], timeout=0.01)
+
+
+def test_caching_wrapper_empty_repolist():
+    """Verify caching_wrapper returns early without executing queries when repolist is empty."""
+    with (
+        patch("cache_manager.cache_facade.get_uncached") as mock_get_uncached,
+        patch("cache_manager.cache_facade.cache_query_results") as mock_cqr,
+    ):
+        cf.caching_wrapper("test_func", "SELECT 1;", [])
+        cf.caching_wrapper("test_func", "SELECT 1;", None)
+        mock_get_uncached.assert_not_called()
+        mock_cqr.assert_not_called()