Merge branch 'main' into fix/issue-1664-stuck-drop-overlay

Author: giswqs

backend/geolibre_server/geolibre_server/sedona_ops.py

@@ -17,10 +17,13 @@
 
 import concurrent.futures
 import json
+import logging
 import math
 import re
 from typing import Any, Optional
 
+logger = logging.getLogger(__name__)
+
 WGS84 = "EPSG:4326"
 
 # Wall-clock budget for a single SQL statement (execute + materialise). Mirrors
@@ -140,6 +143,7 @@ def run_sql(sql: str, layers: Optional[list[dict]] = None) -> dict:
     import geopandas as gpd  # noqa: PLC0415
 
     connection = sedona_db.connect()
+    future = None
     try:
         for layer in layers or []:
             name = str(layer.get("name") or "").strip()
@@ -171,12 +175,12 @@ def _execute():
             try:
                 frame = future.result(timeout=timeout_secs)
             except concurrent.futures.TimeoutError:
-                close = getattr(connection, "close", None)
-                if callable(close):
-                    try:
-                        close()
-                    except Exception:  # noqa: BLE001
-                        pass
+                # No cancellation attempt: the pool has a single worker and a
+                # single task, so by the time this fires ``_execute`` is already
+                # running and ``future.cancel()`` would return False. The Rust
+                # engine exposes no way to interrupt an in-flight query, so the
+                # statement runs to completion and the ``finally`` block below
+                # defers closing the connection until it does.
                 raise SqlTimeout(
                     f"Spatial SQL timed out after {int(timeout_secs)} seconds"
                 ) from None
@@ -226,9 +230,17 @@ def _execute():
     finally:
         # SedonaDB connections are Rust-backed; release promptly rather than
         # waiting on GC. Tolerate bindings that expose no close().
-        close = getattr(connection, "close", None)
-        if callable(close):
-            try:
-                close()
-            except Exception:  # noqa: BLE001 - best-effort cleanup
-                pass
+        def _close_connection(*_args: object) -> None:
+            close = getattr(connection, "close", None)
+            if callable(close):
+                try:
+                    close()
+                except Exception:  # noqa: BLE001 - best-effort cleanup
+                    # A failed close can strand Rust-backed resources, so give
+                    # operators a signal. The SQL text is deliberately omitted.
+                    logger.warning("Failed to close the SedonaDB connection", exc_info=True)
+
+        if future is not None and not future.done():
+            future.add_done_callback(_close_connection)
+        else:
+            _close_connection()

backend/geolibre_server/tests/test_sedona_ops.py

@@ -0,0 +1,56 @@
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from geolibre_server.sedona_ops import SqlTimeout, run_sql
+
+
+@pytest.fixture
+def mock_sedona_db() -> MagicMock:
+    with patch("geolibre_server.sedona_ops._import_sedona") as mock_import:
+        mock_sedona = MagicMock()
+        mock_conn = MagicMock()
+        mock_sedona.connect.return_value = mock_conn
+        mock_import.return_value = mock_sedona
+        yield mock_conn
+
+
+def test_sql_timeout_graceful_shutdown(
+    mock_sedona_db: MagicMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """
+    Test that a SQL query exceeding the timeout raises SqlTimeout,
+    but does NOT immediately close the database connection.
+    Instead, it should attach a callback to close it when the future finishes.
+    """
+    # Lower the timeout to 0.1 seconds for the test
+    monkeypatch.setattr("geolibre_server.sedona_ops._STATEMENT_TIMEOUT_MS", 100)
+
+    # Mock the execute method to block longer than the timeout
+    # We must patch the connection's sql method to block
+    def _slow_sql(*_args: object, **_kwargs: object) -> MagicMock:
+        time.sleep(0.3)  # Blocks longer than the 0.1s timeout
+        mock_df = MagicMock()
+        mock_df.limit.return_value.to_pandas.return_value = MagicMock(columns=[])
+        return mock_df
+
+    mock_sedona_db.sql = _slow_sql
+
+    with pytest.raises(SqlTimeout, match="timed out"):
+        run_sql("SELECT 1")
+
+    # At this exact moment, the TimeoutError was caught and SqlTimeout raised.
+    # The background thread is still sleeping (for another 0.2 seconds).
+    # We must assert that the connection has NOT been closed yet!
+    mock_sedona_db.close.assert_not_called()
+
+    # Wait for the background thread to finish. Poll instead of sleeping a fixed
+    # interval: on a loaded CI host the 0.3s sleep can overrun any margin we
+    # would pick, which would make the assertion below flaky.
+    deadline = time.monotonic() + 5.0
+    while time.monotonic() < deadline and not mock_sedona_db.close.called:
+        time.sleep(0.01)
+
+    # Now the callback should have fired and closed the connection
+    mock_sedona_db.close.assert_called_once()