feat(duckdb): add query timeout

docs/src/advanced/duckdb.md

@@ -29,6 +29,7 @@ The DuckDB backend requires `duckdb>=2.0` (CRS aware `GEOMETRY` type).
 | `only_spatial_tables` | `TIPG_DUCKDB_ONLY_SPATIAL_TABLES` | `True` | Only expose tables/views with a geometry column |
 | `spatial_extent` | `TIPG_DUCKDB_SPATIAL_EXTENT` | `True` | Read the spatial extent and geometry type of parquet sources from their GeoParquet metadata (file footers only). For a glob, every matched file's footer is read on each catalog refresh: disable it for sources matching many remote files |
 | `metadata_threads` | `TIPG_DUCKDB_METADATA_THREADS` | `16` | Number of parquet sources whose metadata is read concurrently when building the catalog |
+| `query_timeout` | `TIPG_DUCKDB_QUERY_TIMEOUT` | | Maximum duration (in seconds) of items and tiles queries. Longer queries are interrupted and return a `504 Gateway Timeout` error. No timeout when not set |
 
 Table level configuration (`TIPG_TABLE_CONFIG__...`, fallback key names, column sorting) is shared with the PostgreSQL backend (see [Configuration](../user_guide/configuration.md)).
 
@@ -259,3 +260,5 @@ gdf.to_parquet("parcels_sorted.parquet", write_covering_bbox=True, row_group_siz
 #### Low zoom levels
 
 A good layout only helps when a tile covers a small part of the dataset: a tile at zoom level 5 over the whole dataset still reads every row group, and only returns the first `TIPG_MAX_FEATURES_PER_TILE` features. For large datasets, only request tiles at higher zoom levels (e.g with the `minZoom` option of your map library's source).
+
+Set `TIPG_DUCKDB_QUERY_TIMEOUT` (e.g `10`) so a single expensive request (a low zoom tile, an items request with a filter on a large dataset) can't keep a worker busy for minutes: the query is interrupted and the request returns a `504 Gateway Timeout` error.

tests/duck/test_timeout.py

@@ -0,0 +1,117 @@
+"""Test DuckDB query timeout."""
+
+import asyncio
+import datetime
+import time
+
+import morecantile
+import pytest
+from fastapi.testclient import TestClient
+
+from tipg.collections import Catalog
+from tipg.duck import DuckDBCatalogSettings, DuckDBSettings
+from tipg.duck.collections import _collection_index
+from tipg.duck.database import create_connection
+from tipg.errors import DEFAULT_STATUS_CODES, QueryTimeout
+
+from .conftest import create_duck_app
+
+duckdb = pytest.importorskip("duckdb")
+
+# Never ending scan: 1e11 rows, none of them in the requested bbox/tile
+SLOW_VIEW = """
+    CREATE VIEW memory.main.slow AS
+    SELECT x AS id, ST_Point(x % 10, x % 10) AS geom
+    FROM range(100000000000) r(x)
+"""
+
+
+def _slow_collection(conn, query_timeout):
+    conn.execute(SLOW_VIEW)
+    (col,) = _collection_index(
+        conn, DuckDBCatalogSettings(tables=["main.slow"], query_timeout=query_timeout)
+    )
+    return col
+
+
+@pytest.fixture
+def slow_collection():
+    """Collection over a never ending scan, with a 0.5s query timeout."""
+    conn = create_connection(DuckDBSettings())
+    yield _slow_collection(conn, 0.5)
+    conn.close()
+
+
+def test_features_timeout(slow_collection):
+    """Items queries are interrupted after the timeout."""
+    assert slow_collection.query_timeout == 0.5
+
+    start = time.perf_counter()
+    with pytest.raises(QueryTimeout):
+        asyncio.run(slow_collection.features(bbox_filter=[50, 50, 60, 60]))
+    assert time.perf_counter() - start < 5
+
+    # the connection is still usable
+    assert slow_collection.connection.execute("SELECT 1").fetchone() == (1,)
+
+
+def test_tile_timeout(slow_collection):
+    """Tile queries are interrupted after the timeout."""
+    tms = morecantile.tms.get("WebMercatorQuad")
+    start = time.perf_counter()
+    with pytest.raises(QueryTimeout):
+        asyncio.run(slow_collection.get_tile(tms=tms, tile=tms.tile(55, 55, 6)))
+    assert time.perf_counter() - start < 5
+
+
+def test_fast_query_not_interrupted():
+    """Queries finishing before the timeout are not affected."""
+    conn = create_connection(DuckDBSettings())
+    try:
+        conn.execute(
+            "CREATE VIEW memory.main.fast AS "
+            "SELECT x AS id, ST_Point(x, x) AS geom FROM range(100) r(x)"
+        )
+        (col,) = _collection_index(
+            conn, DuckDBCatalogSettings(tables=["main.fast"], query_timeout=0.5)
+        )
+        items = asyncio.run(col.features(limit=3))
+        assert len(items["items"]) == 3
+        assert items["matched"] == 100
+
+        # the timer of a finished query doesn't interrupt the next ones
+        time.sleep(0.7)
+        items = asyncio.run(col.features(limit=3))
+        assert len(items["items"]) == 3
+
+        tms = morecantile.tms.get("WebMercatorQuad")
+        assert asyncio.run(col.get_tile(tms=tms, tile=tms.tile(5, 5, 4)))
+    finally:
+        conn.close()
+
+
+def test_no_timeout_by_default():
+    """No timeout unless configured."""
+    assert DuckDBCatalogSettings().query_timeout is None
+
+
+def test_timeout_status_code():
+    """Timeouts are returned as `504 Gateway Timeout`."""
+    assert DEFAULT_STATUS_CODES[QueryTimeout] == 504
+
+    app = create_duck_app(DuckDBSettings(), DuckDBCatalogSettings())
+    with TestClient(app) as client:
+        conn = client.app.state.duckdb_connection
+        col = _slow_collection(conn, 0.5)
+        client.app.state.collection_catalog = Catalog(
+            collections={col.id: col}, last_updated=datetime.datetime.now()
+        )
+
+        response = client.get(
+            "/collections/main.slow/items", params={"bbox": "50,50,60,60"}
+        )
+        assert response.status_code == 504
+        assert "timeout" in response.json()["detail"].lower()
+
+        response = client.get("/collections/main.slow/tiles/WebMercatorQuad/6/41/28")
+        assert response.status_code == 504

tipg/duck/collections.py

@@ -3,8 +3,10 @@
 import functools
 import math
 import re
-from collections.abc import Iterator
+import threading
+from collections.abc import Generator, Iterator
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
 from typing import Any
 
 import duckdb
@@ -17,7 +19,12 @@
 from starlette.concurrency import run_in_threadpool
 
 from ..collections import Collection, Column, Feature, ItemList, Srid, _quote_ident
-from ..errors import InvalidGeometryColumnName, InvalidLimit, InvalidPropertyName
+from ..errors import (
+    InvalidGeometryColumnName,
+    InvalidLimit,
+    InvalidPropertyName,
+    QueryTimeout,
+)
 from ..logger import logger
 from ..settings import FeaturesSettings, MVTSettings, TableConfig, TableSettings
 from .settings import DuckDBCatalogSettings
@@ -141,6 +148,31 @@ def _crs_literal(srid: Srid) -> str:
     return _quote_literal(_crs(srid))
 
 
+@contextmanager
+def _interrupt_after(
+    cursor: duckdb.DuckDBPyConnection,
+    timeout: float | None,
+) -> Generator[None]:
+    """Interrupt the queries of ``cursor`` still running after ``timeout`` seconds.
+
+    Raises ``QueryTimeout`` for interrupted queries. Interrupting a cursor
+    after its query finished has no effect on its next queries.
+    """
+    if not timeout:
+        yield
+        return
+
+    timer = threading.Timer(timeout, cursor.interrupt)
+    timer.daemon = True
+    timer.start()
+    try:
+        yield
+    except duckdb.InterruptException as e:
+        raise QueryTimeout(f"Query timeout: interrupted after {timeout}s") from e
+    finally:
+        timer.cancel()
+
+
 class DuckCollection(Collection):
     """Model for DuckDB Table and View."""
 
@@ -155,6 +187,9 @@ class DuckCollection(Collection):
     # expression of its ``xmin``, ``ymin``, ``xmax`` and ``ymax`` columns.
     bbox_coverings: dict[str, dict[str, str]] = {}
 
+    # Maximum duration (seconds) of items and tiles queries (no limit if None)
+    query_timeout: float | None = None
+
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @property
@@ -433,30 +468,33 @@ def _run_features_queries(
         """
         cursor = self.connection.cursor()
         try:
-            matched = None
-            if count_sql:
-                logger.debug(count_sql)
-                matched = cursor.execute(count_sql).fetchone()[0]  # type: ignore [index]
-
-            logger.debug(sql)
-            cursor.execute(sql)
-            names = [d[0] for d in cursor.description]  # type: ignore [union-attr]
-            rows = cursor.fetchall()
-            features = [
-                self._row_to_feature(dict(zip(names, r)), geom_as_wkt, offset + i)
-                for i, r in enumerate(rows[:limit], start=1)
-            ]
+            with _interrupt_after(cursor, self.query_timeout):
+                matched = None
+                if count_sql:
+                    logger.debug(count_sql)
+                    matched = cursor.execute(count_sql).fetchone()[0]  # type: ignore [index]
+
+                logger.debug(sql)
+                cursor.execute(sql)
+                names = [d[0] for d in cursor.description]  # type: ignore [union-attr]
+                rows = cursor.fetchall()
         finally:
             cursor.close()
 
+        features = [
+            self._row_to_feature(dict(zip(names, r)), geom_as_wkt, offset + i)
+            for i, r in enumerate(rows[:limit], start=1)
+        ]
+
         return matched, features, len(rows) > limit
 
     def _run_tile_query(self, sql: str, layer: str) -> bytes:
         """Run the MVT query (blocking)."""
         cursor = self.connection.cursor()
         try:
-            logger.debug(sql)
-            tile = cursor.execute(sql, [layer]).fetchone()[0]  # type: ignore [index]
+            with _interrupt_after(cursor, self.query_timeout):
+                logger.debug(sql)
+                tile = cursor.execute(sql, [layer]).fetchone()[0]  # type: ignore [index]
         finally:
             cursor.close()
 
@@ -1168,6 +1206,7 @@ def _collection_index(  # noqa: C901
                 datetime_column=datetime_column,
                 geometry_column=geometry_column,
                 bbox_coverings=bbox_coverings,
+                query_timeout=settings.query_timeout,
                 connection=connection,
             )
 

tipg/duck/settings.py

@@ -1,5 +1,6 @@
 """tipg DuckDB settings."""
 
+from pydantic import Field
 from pydantic_settings import BaseSettings
 
 
@@ -48,6 +49,9 @@ class DuckDBCatalogSettings(BaseSettings):
             speed up spatial filters (`bbox`, tiles, cql2 spatial predicates).
         metadata_threads: number of parquet sources whose metadata is read
             concurrently when building the catalog.
+        query_timeout: maximum duration (in seconds) of the items and tiles
+            queries of the collections. Longer queries are interrupted and
+            return a `504 Gateway Timeout` error. No timeout when not set.
 
     """
 
@@ -57,6 +61,7 @@ class DuckDBCatalogSettings(BaseSettings):
     spatial_extent: bool = True
     metadata_threads: int = 16
     only_spatial_tables: bool = True
+    query_timeout: float | None = Field(default=None, gt=0)
 
     model_config = {
         "env_prefix": "TIPG_DUCKDB_",

tipg/errors.py

@@ -68,6 +68,10 @@ class MissingCollectionCatalog(TiPgError):
     """`collection_catalog` not registered in the application state."""
 
 
+class QueryTimeout(TiPgError):
+    """Query interrupted after the configured timeout."""
+
+
 DEFAULT_STATUS_CODES = {
     NotFound: status.HTTP_404_NOT_FOUND,
     InvalidBBox: status.HTTP_422_UNPROCESSABLE_CONTENT,
@@ -84,6 +88,7 @@ class MissingCollectionCatalog(TiPgError):
     MissingFunctionParameter: status.HTTP_422_UNPROCESSABLE_CONTENT,
     FunctionDirectoryDoesNotExist: status.HTTP_500_INTERNAL_SERVER_ERROR,
     MissingCollectionCatalog: status.HTTP_500_INTERNAL_SERVER_ERROR,
+    QueryTimeout: status.HTTP_504_GATEWAY_TIMEOUT,
 }