From 50bbcf8d29e4fe8eb8f2d1ff23df040229c1668b Mon Sep 17 00:00:00 2001
From: Clare72 <cp390@cam.ac.uk>
Date: Fri, 23 Oct 2026 15:08:38 +0100
Subject: [PATCH] FlyBase: speed up resolve_entity broad search and apply
 statement timeouts

resolve_entity's broad '%name%' search filtered feature types via a join on
cvterm.name, so Postgres scanned all ~98M feature rows: ~95s for a name with
no match, beyond chado.flybase.org's 120s limit on CI. Look up the type ids
first and filter on the type_id index instead (~1s, same results).

get_connection passed statement_timeout as a startup option, which the
server ignores (sessions kept the chado_public role's 120s). Apply it with
set_config after connecting, and commit so a rollback can't revert it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/test/test_flybase_stocks.py | 16 ++++++++++++++++
 src/vfbquery/flybase_db.py      | 19 +++++++++++++++----
 src/vfbquery/flybase_stocks.py  | 13 +++++++++++--
 3 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/test/test_flybase_stocks.py b/src/test/test_flybase_stocks.py
index 7bae36b..edbdef7 100644
--- a/src/test/test_flybase_stocks.py
+++ b/src/test/test_flybase_stocks.py
@@ -111,6 +111,22 @@ def test_nonexistent_name(self):
         assert result["match_type"] == "NOT_FOUND"
 
 
+class TestConnectionStatementTimeout:
+    @pytest.mark.integration
+    def test_requested_timeout_is_applied(self):
+        # chado.flybase.org ignores a startup "-c statement_timeout" option and
+        # keeps the role's 120s, so check the session actually got ours, and
+        # that it survives a rollback.
+        from vfbquery.flybase_db import get_connection
+        conn = get_connection(statement_timeout_ms=45000)
+        try:
+            assert conn.execute("SHOW statement_timeout").fetchone()[0] == "45s"
+            conn.rollback()
+            assert conn.execute("SHOW statement_timeout").fetchone()[0] == "45s"
+        finally:
+            conn.close()
+
+
 # ---------------------------------------------------------------------------
 # find_stocks tests
 # ---------------------------------------------------------------------------
diff --git a/src/vfbquery/flybase_db.py b/src/vfbquery/flybase_db.py
index a6a0734..813c1ad 100644
--- a/src/vfbquery/flybase_db.py
+++ b/src/vfbquery/flybase_db.py
@@ -14,13 +14,24 @@
 def get_connection(statement_timeout_ms=60000):
     """Get a connection to FlyBase Chado with statement timeout.
 
+    The timeout is applied with ``set_config`` after connecting. Passing it as a
+    startup option (``options="-c statement_timeout=..."``) has no effect on
+    chado.flybase.org: sessions keep the ``chado_public`` role's 120s default,
+    so every caller's requested timeout was silently ignored (2026-10). The
+    setting is committed straight away so a later rollback can't revert it.
+
     :param statement_timeout_ms: SQL statement timeout in milliseconds (default 60s)
     :return: psycopg connection
     """
-    return psycopg.connect(
-        **FLYBASE_DB,
-        options=f"-c statement_timeout={statement_timeout_ms}",
-    )
+    conn = psycopg.connect(**FLYBASE_DB)
+    try:
+        conn.execute("SELECT set_config('statement_timeout', %s, false)",
+                     (str(int(statement_timeout_ms)),))
+        conn.commit()
+    except BaseException:
+        conn.close()
+        raise
+    return conn
 
 
 def is_connection_error(exc):
diff --git a/src/vfbquery/flybase_stocks.py b/src/vfbquery/flybase_stocks.py
index bb4fa63..e6ec1a8 100644
--- a/src/vfbquery/flybase_stocks.py
+++ b/src/vfbquery/flybase_stocks.py
@@ -114,16 +114,25 @@ def _resolve_entity_impl(conn, name_or_id):
         }
 
     # 3. Broad ILIKE search
+    # A '%...%' pattern can't use the feature.name index, so filter on the
+    # type_id index instead. Passing the type ids as a literal array lets the
+    # planner use it (~1.5M rows); filtering on c.name via the join made it
+    # scan all ~98M feature rows — ~95s for a name with no match, past the
+    # server's 120s statement limit on CI (2026-10).
+    type_ids = _run_query(conn, """
+        SELECT cvterm_id FROM cvterm WHERE name = ANY(%(types)s)
+    """, {"types": feature_types})
     df = _run_query(conn, """
         SELECT DISTINCT f.name, f.uniquename, c.name AS type
         FROM feature f
         JOIN cvterm c ON f.type_id = c.cvterm_id
         WHERE f.is_obsolete = false
-          AND c.name = ANY(%(types)s)
+          AND f.type_id = ANY(%(type_ids)s)
           AND (f.name ILIKE %(pattern)s)
         ORDER BY c.name, f.name
         LIMIT 20
-    """, {"pattern": "%" + name_or_id + "%", "types": feature_types})
+    """, {"pattern": "%" + name_or_id + "%",
+          "type_ids": [] if type_ids.empty else type_ids["cvterm_id"].tolist()})
 
     if len(df) > 0:
         return {