Merge pull request #496 from wri/cache_get_data_env

Cache _get_data_env results for 5 minutes

app/routes/datasets/queries.py

@@ -7,25 +7,23 @@
 from uuid import UUID, uuid4
 
 import httpx
+from async_lru import alru_cache
 from asyncpg import DataError, InsufficientPrivilegeError, SyntaxOrAccessError
 from fastapi import APIRouter, Depends, HTTPException, Query
 from fastapi import Request as FastApiRequest
 from fastapi import Response as FastApiResponse
 from fastapi.encoders import jsonable_encoder
 from fastapi.logger import logger
-
 # from fastapi.openapi.models import APIKey
 from fastapi.responses import RedirectResponse
 from pglast import printers  # noqa
 from pglast import Node, parse_sql
 from pglast.parser import ParseError
 from pglast.printer import RawStream
 from pydantic.tools import parse_obj_as
-from sqlalchemy.sql import and_
 
 from ...authentication.token import is_gfwpro_admin_for_query
 from ...application import db
-
 # from ...authentication.api_keys import get_api_key
 from ...crud import assets
 from ...models.enum.assets import AssetType
@@ -62,7 +60,6 @@
 from ...models.enum.queries import QueryFormat, QueryType
 from ...models.orm.assets import Asset as AssetORM
 from ...models.orm.queries.raster_assets import latest_raster_tile_sets
-from ...models.orm.versions import Version as VersionORM
 from ...models.pydantic.asset_metadata import RasterTable, RasterTableRow
 from ...models.pydantic.creation_options import NoDataType
 from ...models.pydantic.geostore import Geometry, GeostoreCommon
@@ -659,13 +656,14 @@ async def _query_raster_lambda(
 
 
 def _get_area_density_name(nm):
-    """Return empty string if nm doesn't not have an area-density suffix, else
+    """Return empty string if nm doesn't have an area-density suffix, else
     return nm with the area-density suffix removed."""
     for suffix in AREA_DENSITY_RASTER_SUFFIXES:
         if nm.endswith(suffix):
             return nm[:-len(suffix)]
     return ""
 
+
 def _get_default_layer(dataset, pixel_meaning):
     default_type = pixel_meaning
     area_density_name = _get_area_density_name(default_type)
@@ -683,6 +681,7 @@ def _get_default_layer(dataset, pixel_meaning):
         return f"{dataset}__{default_type}"
 
 
+@alru_cache(maxsize=16, ttl=300.0)
 async def _get_data_environment(grid: Grid) -> DataEnvironment:
     # get all raster tile set assets with the same grid.
     latest_tile_sets = await db.all(latest_raster_tile_sets, {"grid": grid})

tests_v2/unit/app/routes/analysis/test_analysis.py

@@ -81,6 +81,9 @@ async def test_analysis_with_huge_geostore(
 async def test_raster_analysis_payload_shape(
     generic_dataset, async_client: AsyncClient, monkeypatch: MonkeyPatch
 ):
+    """Note that this test depends on the output of _get_data_environment
+    which will likely have cached values from other tests, so we clear it."""
+
     dataset_name, _ = generic_dataset
     pixel_meaning: str = "date_conf"
     no_data_value = 0
@@ -103,6 +106,9 @@ async def test_raster_analysis_payload_shape(
         )
         monkeypatch.setattr(geostore.rw_api, "get_geostore", mock_rw_get_geostore)
 
+        # The other tests will have polluted the data env cache. Clear it.
+        queries._get_data_environment.cache_clear()
+
         _ = await async_client.get(
             f"/analysis/zonal/17076d5ea9f214a5bdb68cc40433addb?geostore_origin=rw&group_by=umd_tree_cover_loss__year&filters=is__umd_regional_primary_forest_2001&filters=umd_tree_cover_density_2000__30&sum=area__ha&start_date=2001"
         )