diff --git a/docs/src/advanced/duckdb.md b/docs/src/advanced/duckdb.md index 1105372..c3b6b8b 100644 --- a/docs/src/advanced/duckdb.md +++ b/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. diff --git a/tests/duck/test_timeout.py b/tests/duck/test_timeout.py new file mode 100644 index 0000000..226bed3 --- /dev/null +++ b/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 diff --git a/tipg/duck/collections.py b/tipg/duck/collections.py index 6cba9f9..7369332 100644 --- a/tipg/duck/collections.py +++ b/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, ) diff --git a/tipg/duck/settings.py b/tipg/duck/settings.py index 4f8525b..1034d43 100644 --- a/tipg/duck/settings.py +++ b/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_", diff --git a/tipg/errors.py b/tipg/errors.py index a6ef9e1..ebbe698 100644 --- a/tipg/errors.py +++ b/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, }