Deeper fixes #1: FK indexes, session-factory invalidation, HTTP pooling

data/db/migrations/0013_fk_indexes.sql

@@ -0,0 +1,24 @@
+-- Migration 0013: index FK columns that aren't already covered
+--
+-- Without these, "given a source/reach/guidebook, find related rows"
+-- queries fall back to a full table scan. Tables are small today
+-- (gauge_source 224, reach_class 401, reach_guidebook 1093,
+-- latest_gauge_observation 461) so latency is fine, but adding the
+-- indexes now removes the future scaling cliff and matches the
+-- pattern already used for ix_reach_state_state_id.
+--
+-- gauge_source / reach_guidebook PKs are (left_id, right_id) composite.
+-- The PK index serves "given left, find right" but not the reverse.
+-- These supplemental indexes cover the reverse direction.
+
+CREATE INDEX IF NOT EXISTS ix_gauge_source_source_id
+    ON gauge_source (source_id);
+
+CREATE INDEX IF NOT EXISTS ix_reach_class_reach_id
+    ON reach_class (reach_id);
+
+CREATE INDEX IF NOT EXISTS ix_reach_guidebook_guidebook_id
+    ON reach_guidebook (guidebook_id);
+
+CREATE INDEX IF NOT EXISTS ix_latest_gauge_observation_source_id
+    ON latest_gauge_observation (source_id);

src/kayak/db/engine.py

@@ -24,12 +24,15 @@ def get_engine(url: str | None = None) -> Engine:
 
     When ``url`` is supplied, the prior engine (if any) is disposed before the
     new one is built — otherwise every ``get_engine(url=…)`` call would orphan
-    its connection pool.
+    its connection pool. The cached session factory is also invalidated so the
+    next ``get_session_factory()`` call binds to the new engine instead of the
+    disposed one.
     """
-    global _engine
+    global _engine, _session_factory
     if _engine is None or url is not None:
         if _engine is not None and url is not None:
             _engine.dispose()
+            _session_factory = None
         db_url = url or DATABASE_URL
         connect_args = {}
         if db_url.startswith("sqlite"):
@@ -41,7 +44,11 @@ def get_engine(url: str | None = None) -> Engine:
 
 
 def get_session_factory(url: str | None = None) -> sessionmaker[Session]:
-    """Return a sessionmaker bound to the engine."""
+    """Return a sessionmaker bound to the current engine.
+
+    Invariant: the returned factory is always bound to the engine that
+    ``get_engine()`` would currently return — never to a disposed engine.
+    """
     global _session_factory
     if _session_factory is None or url is not None:
         _session_factory = sessionmaker(bind=get_engine(url))

src/kayak/db/models.py

@@ -208,6 +208,10 @@ class GaugeSource(Base):
         ForeignKey("source.id", ondelete="CASCADE"), primary_key=True
     )
 
+    # The composite PK indexes (gauge_id, source_id) which serves
+    # "given a gauge, find its sources" but not the reverse direction.
+    __table_args__ = (Index("ix_gauge_source_source_id", "source_id"),)
+
 
 # ---------------------------------------------------------------------------
 # fetch_url
@@ -403,6 +407,8 @@ class LatestGaugeObservation(Base):
     # relationships
     gauge: Mapped[Gauge] = relationship()
 
+    __table_args__ = (Index("ix_latest_gauge_observation_source_id", "source_id"),)
+
 
 # ---------------------------------------------------------------------------
 # reach
@@ -550,6 +556,7 @@ class ReachClass(Base):
             "low IS NULL OR high IS NULL OR low <= high",
             name="ck_reach_class_low_le_high",
         ),
+        Index("ix_reach_class_reach_id", "reach_id"),
     )
 
 
@@ -611,6 +618,10 @@ class ReachGuidebook(Base):
     run: Mapped[str | None] = mapped_column(Text)
     url: Mapped[str | None] = mapped_column(Text)
 
+    # Mirror of ix_reach_state_state_id pattern: composite PK indexes
+    # (reach_id, guidebook_id), so the reverse direction needs its own.
+    __table_args__ = (Index("ix_reach_guidebook_guidebook_id", "guidebook_id"),)
+
 
 # ---------------------------------------------------------------------------
 # editor (Phase 1 — editor accounts for proposing changes)

src/kayak/utils/http_client.py

@@ -17,6 +17,31 @@
 logger = logging.getLogger(__name__)
 
 
+# Module-level pooled Session. A typical ``levels fetch`` run hits ~50 URLs
+# across a handful of hosts; without pooling each call burns a fresh TLS
+# handshake. requests.Session keeps a per-host connection pool
+# (HTTPAdapter default pool_connections=10, pool_maxsize=10) so the second
+# call to a host reuses the existing connection.
+_session: requests.Session | None = None
+
+
+def _get_session() -> requests.Session:
+    """Return the module-level pooled Session, creating it on first call."""
+    global _session
+    if _session is None:
+        _session = requests.Session()
+        _session.headers["User-Agent"] = FETCH_USER_AGENT
+    return _session
+
+
+def reset_session() -> None:
+    """Close the pooled session and drop the reference (test isolation)."""
+    global _session
+    if _session is not None:
+        _session.close()
+        _session = None
+
+
 # Hosts whose TLS chain or cipher suite can't be validated with a stock Debian
 # CA bundle. Every entry here is a MITM risk — only add hosts we've confirmed
 # need relaxed TLS *and* where the payload is non-sensitive public data (river
@@ -167,10 +192,11 @@ def fetch(url: str, timeout: int | None = None) -> FetchResult:
 
     verify = not _is_insecure_host(url)
 
+    session = _get_session()
     last_result: FetchResult | None = None
     for attempt in range(_MAX_RETRIES):
         try:
-            response = requests.get(
+            response = session.get(
                 url,
                 timeout=timeout,
                 headers={"User-Agent": FETCH_USER_AGENT},

tests/test_engine.py

@@ -44,6 +44,30 @@ def test_url_override_disposes_prior_engine(self):
             assert e1 is not e2
             spy.assert_called_once()
 
+    def test_url_override_invalidates_session_factory(self):
+        """Replacing the engine via get_engine(url=...) must drop the cached
+        session factory so a subsequent get_session_factory() rebuilds it
+        against the new engine.
+
+        Regression: previously the factory was only rebuilt when
+        get_session_factory() itself was called with a url. If a caller went
+        through get_engine() directly, the old factory stayed bound to the
+        now-disposed engine — sessions created from it would fail or, worse,
+        write to a stale connection.
+        """
+        # Prime both caches
+        get_engine("sqlite:///:memory:")
+        f1 = get_session_factory()
+
+        # Swap the engine via get_engine() directly
+        e2 = get_engine("sqlite:///:memory:")
+
+        # A subsequent get_session_factory() with no url must NOT return f1.
+        # It should rebuild against e2.
+        f2 = get_session_factory()
+        assert f2 is not f1
+        assert f2.kw["bind"] is e2
+
 
 class TestSQLitePragmas:
     def teardown_method(self) -> None:

tests/test_utils/test_http_client.py

@@ -129,74 +129,78 @@ def test_write_file_creates_parent_dirs(self, tmp_path):
 
 
 class TestFetch:
-    @patch("kayak.utils.http_client.requests.get")
-    def test_fetch_success(self, mock_get):
+    @pytest.fixture(autouse=True)
+    def mock_session(self):
+        """Patch the module-level Session factory with a MagicMock.
+
+        Tests assert against ``mock_session.get`` instead of patching
+        ``requests.get`` directly. This matches the production pattern
+        where every fetch goes through one pooled Session.
+        """
+        sess = MagicMock()
+        with patch("kayak.utils.http_client._get_session", return_value=sess):
+            yield sess
+
+    def test_fetch_success(self, mock_session):
         mock_resp = MagicMock(spec=requests.Response)
         mock_resp.status_code = 200
         mock_resp.text = "ok"
-        mock_get.return_value = mock_resp
+        mock_session.get.return_value = mock_resp
 
         result = fetch("http://example.com/data")
         assert result.ok is True
         assert result.url == "http://example.com/data"
 
-    @patch("kayak.utils.http_client.requests.get")
-    def test_fetch_verifies_tls_by_default(self, mock_get):
+    def test_fetch_verifies_tls_by_default(self, mock_session):
         mock_resp = MagicMock(spec=requests.Response)
         mock_resp.status_code = 200
-        mock_get.return_value = mock_resp
+        mock_session.get.return_value = mock_resp
         fetch("https://example.com/data")
-        _, kwargs = mock_get.call_args
+        _, kwargs = mock_session.get.call_args
         assert kwargs["verify"] is True
 
-    @patch("kayak.utils.http_client.requests.get")
-    def test_fetch_skips_verify_for_insecure_host(self, mock_get):
+    def test_fetch_skips_verify_for_insecure_host(self, mock_session):
         mock_resp = MagicMock(spec=requests.Response)
         mock_resp.status_code = 200
-        mock_get.return_value = mock_resp
+        mock_session.get.return_value = mock_resp
         fetch("https://www.nwd-wc.usace.army.mil/foo")
-        _, kwargs = mock_get.call_args
+        _, kwargs = mock_session.get.call_args
         assert kwargs["verify"] is False
 
-    @patch("kayak.utils.http_client.requests.get")
-    def test_fetch_passes_user_agent(self, mock_get):
+    def test_fetch_passes_user_agent(self, mock_session):
         mock_resp = MagicMock(spec=requests.Response)
         mock_resp.status_code = 200
-        mock_get.return_value = mock_resp
+        mock_session.get.return_value = mock_resp
         fetch("http://example.com/data")
-        _, kwargs = mock_get.call_args
+        _, kwargs = mock_session.get.call_args
         assert "User-Agent" in kwargs["headers"]
 
     @patch("kayak.utils.http_client.time.sleep")
-    @patch("kayak.utils.http_client.requests.get")
-    def test_fetch_exception_returns_error_result(self, mock_get, mock_sleep):
-        mock_get.side_effect = requests.ConnectionError("refused")
+    def test_fetch_exception_returns_error_result(self, mock_sleep, mock_session):
+        mock_session.get.side_effect = requests.ConnectionError("refused")
         result = fetch("http://example.com/data")
         assert result.ok is False
         assert result.error is not None
 
-    @patch("kayak.utils.http_client.requests.get")
-    def test_fetch_custom_timeout(self, mock_get):
+    def test_fetch_custom_timeout(self, mock_session):
         mock_resp = MagicMock(spec=requests.Response)
         mock_resp.status_code = 200
-        mock_get.return_value = mock_resp
+        mock_session.get.return_value = mock_resp
         fetch("http://example.com/data", timeout=10)
-        _, kwargs = mock_get.call_args
+        _, kwargs = mock_session.get.call_args
         assert kwargs["timeout"] == 10
 
-    @patch("kayak.utils.http_client.requests.get")
-    def test_fetch_disables_redirects(self, mock_get):
+    def test_fetch_disables_redirects(self, mock_session):
         """Sync fetch must not follow redirects, else a 3xx could bypass
         _validate_url (the redirect target wouldn't be re-validated)."""
         mock_resp = MagicMock(spec=requests.Response)
         mock_resp.status_code = 200
-        mock_get.return_value = mock_resp
+        mock_session.get.return_value = mock_resp
         fetch("http://example.com/data")
-        _, kwargs = mock_get.call_args
+        _, kwargs = mock_session.get.call_args
         assert kwargs["allow_redirects"] is False
 
-    @patch("kayak.utils.http_client.requests.get")
-    def test_fetch_rejects_ssrf_url(self, mock_get):
+    def test_fetch_rejects_ssrf_url(self, mock_session):
         """fetch() returns an error FetchResult (no HTTP call) when the URL
         resolves to an internal IP."""
         with patch(
@@ -207,7 +211,49 @@ def test_fetch_rejects_ssrf_url(self, mock_get):
         assert result.ok is False
         assert result.error is not None
         assert "blocked IP" in result.error
-        mock_get.assert_not_called()
+        mock_session.get.assert_not_called()
+
+
+class TestSessionPooling:
+    """Two fetches go through the same pooled Session (one TLS handshake)."""
+
+    def teardown_method(self) -> None:
+        from kayak.utils.http_client import reset_session
+
+        reset_session()
+
+    def test_get_session_is_singleton(self):
+        from kayak.utils.http_client import _get_session
+
+        s1 = _get_session()
+        s2 = _get_session()
+        assert s1 is s2
+
+    def test_reset_session_creates_a_new_one(self):
+        from kayak.utils.http_client import _get_session, reset_session
+
+        s1 = _get_session()
+        reset_session()
+        s2 = _get_session()
+        assert s1 is not s2
+
+    def test_user_agent_set_on_session(self):
+        from kayak.config import FETCH_USER_AGENT
+        from kayak.utils.http_client import _get_session
+
+        sess = _get_session()
+        assert sess.headers.get("User-Agent") == FETCH_USER_AGENT
+
+    def test_two_fetches_share_session(self):
+        """Both calls route through the same Session instance."""
+        sess_mock = MagicMock()
+        mock_resp = MagicMock(spec=requests.Response)
+        mock_resp.status_code = 200
+        sess_mock.get.return_value = mock_resp
+        with patch("kayak.utils.http_client._get_session", return_value=sess_mock):
+            fetch("http://example.com/a")
+            fetch("http://example.com/b")
+        assert sess_mock.get.call_count == 2
 
 
 # ---------------------------------------------------------------------------