move add bbox shape to fn

Author: jonhealy1

sfeos_tools/sfeos_tools/cli.py

@@ -14,15 +14,15 @@
 
 import click
 
-from stac_fastapi.core.utilities import bbox2polygon
+from stac_fastapi.sfeos_helpers.database import add_bbox_shape_to_collection
 from stac_fastapi.sfeos_helpers.mappings import COLLECTIONS_INDEX
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-async def add_bbox_shape_to_collection(client, collection_doc, backend):
-    """Add bbox_shape field to a single collection document.
+async def process_collection_bbox_shape(client, collection_doc, backend):
+    """Process a single collection document to add bbox_shape field.
 
     Args:
         client: Elasticsearch/OpenSearch client
@@ -35,56 +35,12 @@ async def add_bbox_shape_to_collection(client, collection_doc, backend):
     collection = collection_doc["_source"]
     collection_id = collection.get("id", collection_doc["_id"])
 
-    # Check if bbox_shape already exists
-    if "bbox_shape" in collection:
-        logger.info(
-            f"Collection '{collection_id}' already has bbox_shape field, skipping"
-        )
-        return False
-
-    # Check if collection has spatial extent
-    if "extent" not in collection or "spatial" not in collection["extent"]:
-        logger.warning(f"Collection '{collection_id}' has no spatial extent, skipping")
-        return False
-
-    spatial_extent = collection["extent"]["spatial"]
-    if "bbox" not in spatial_extent or not spatial_extent["bbox"]:
-        logger.warning(
-            f"Collection '{collection_id}' has no bbox in spatial extent, skipping"
-        )
-        return False
-
-    # Get the first bbox (collections can have multiple bboxes, but we use the first one)
-    bbox = (
-        spatial_extent["bbox"][0]
-        if isinstance(spatial_extent["bbox"][0], list)
-        else spatial_extent["bbox"]
-    )
+    # Use the shared function to add bbox_shape
+    was_added = add_bbox_shape_to_collection(collection)
 
-    if len(bbox) < 4:
-        logger.warning(
-            f"Collection '{collection_id}': bbox has insufficient coordinates (length={len(bbox)}), expected at least 4"
-        )
+    if not was_added:
         return False
 
-    # Extract 2D coordinates (bbox can be 2D [minx, miny, maxx, maxy] or 3D [minx, miny, minz, maxx, maxy, maxz])
-    # For 2D polygon, we only need the x,y coordinates and discard altitude (z) values
-    minx, miny = bbox[0], bbox[1]
-    if len(bbox) == 4:
-        # 2D bbox: [minx, miny, maxx, maxy]
-        maxx, maxy = bbox[2], bbox[3]
-    else:
-        # 3D bbox: [minx, miny, minz, maxx, maxy, maxz]
-        # Extract indices 3,4 for maxx,maxy - discarding altitude at indices 2 (minz) and 5 (maxz)
-        maxx, maxy = bbox[3], bbox[4]
-
-    # Convert bbox to GeoJSON polygon
-    bbox_polygon_coords = bbox2polygon(minx, miny, maxx, maxy)
-    collection["bbox_shape"] = {
-        "type": "Polygon",
-        "coordinates": bbox_polygon_coords,
-    }
-
     # Update the collection in the database
     if backend == "elasticsearch":
         await client.index(
@@ -154,7 +110,7 @@ async def run_add_bbox_shape(backend):
         skipped_count = 0
 
         for hit in response["hits"]["hits"]:
-            was_updated = await add_bbox_shape_to_collection(client, hit, backend)
+            was_updated = await process_collection_bbox_shape(client, hit, backend)
             if was_updated:
                 updated_count += 1
             else:

stac_fastapi/core/stac_fastapi/core/serializers.py

@@ -10,7 +10,7 @@
 
 from stac_fastapi.core.datetime_utils import now_to_rfc3339_str
 from stac_fastapi.core.models.links import CollectionLinks
-from stac_fastapi.core.utilities import bbox2polygon, get_bool_env
+from stac_fastapi.core.utilities import get_bool_env
 from stac_fastapi.types import stac as stac_types
 from stac_fastapi.types.links import ItemLinks, resolve_links
 
@@ -139,45 +139,15 @@ def stac_to_db(
         Returns:
             stac_types.Collection: The database-ready STAC Collection object.
         """
+        from stac_fastapi.sfeos_helpers.database import add_bbox_shape_to_collection
+
         collection = deepcopy(collection)
         collection["links"] = resolve_links(
             collection.get("links", []), str(request.base_url)
         )
 
-        # Convert bbox to bbox_shape for geospatial queries
-        if "extent" in collection and "spatial" in collection["extent"]:
-            spatial_extent = collection["extent"]["spatial"]
-            if "bbox" in spatial_extent and spatial_extent["bbox"]:
-                # Get the first bbox (collections can have multiple bboxes, but we use the first one)
-                bbox = (
-                    spatial_extent["bbox"][0]
-                    if isinstance(spatial_extent["bbox"][0], list)
-                    else spatial_extent["bbox"]
-                )
-                collection_id = collection.get("id", "unknown")
-
-                if len(bbox) >= 4:
-                    # Extract 2D coordinates (bbox can be 2D [minx, miny, maxx, maxy] or 3D [minx, miny, minz, maxx, maxy, maxz])
-                    # For 2D polygon, we only need the x,y coordinates and discard altitude (z) values
-                    minx, miny = bbox[0], bbox[1]
-                    if len(bbox) == 4:
-                        # 2D bbox: [minx, miny, maxx, maxy]
-                        maxx, maxy = bbox[2], bbox[3]
-                    else:
-                        # 3D bbox: [minx, miny, minz, maxx, maxy, maxz]
-                        # Extract indices 3,4 for maxx,maxy - discarding altitude at indices 2 (minz) and 5 (maxz)
-                        maxx, maxy = bbox[3], bbox[4]
-
-                    # Convert bbox to GeoJSON polygon
-                    bbox_polygon_coords = bbox2polygon(minx, miny, maxx, maxy)
-                    collection["bbox_shape"] = {
-                        "type": "Polygon",
-                        "coordinates": bbox_polygon_coords,
-                    }
-                else:
-                    logger.warning(
-                        f"Collection '{collection_id}': bbox has insufficient coordinates (length={len(bbox)}), expected at least 4"
-                    )
+        # Convert bbox to bbox_shape for geospatial queries (ES/OS specific)
+        add_bbox_shape_to_collection(collection)
 
         if get_bool_env("STAC_INDEX_ASSETS"):
             collection["assets"] = [

stac_fastapi/sfeos_helpers/stac_fastapi/sfeos_helpers/database/__init__.py

@@ -48,7 +48,7 @@
     apply_intersects_filter_shared,
     populate_sort_shared,
 )
-from .utils import get_bool_env, validate_refresh
+from .utils import add_bbox_shape_to_collection, get_bool_env, validate_refresh
 
 __all__ = [
     # Index operations
@@ -72,6 +72,7 @@
     # Utility functions
     "validate_refresh",
     "get_bool_env",
+    "add_bbox_shape_to_collection",
     # Datetime utilities
     "return_date",
     "extract_date",

stac_fastapi/sfeos_helpers/stac_fastapi/sfeos_helpers/database/utils.py

@@ -5,16 +5,94 @@
 """
 
 import logging
-from typing import Dict, List, Union
+from typing import Any, Dict, List, Union
 
-from stac_fastapi.core.utilities import get_bool_env
+from stac_fastapi.core.utilities import bbox2polygon, get_bool_env
 from stac_fastapi.extensions.core.transaction.request import (
     PatchAddReplaceTest,
     PatchOperation,
     PatchRemove,
 )
 from stac_fastapi.sfeos_helpers.models.patch import ElasticPath, ESCommandSet
 
+logger = logging.getLogger(__name__)
+
+
+def add_bbox_shape_to_collection(collection: Dict[str, Any]) -> bool:
+    """Add bbox_shape field to a collection document for spatial queries.
+
+    This function extracts the bounding box from a collection's spatial extent
+    and converts it to a GeoJSON polygon shape that can be used for geospatial
+    queries in Elasticsearch/OpenSearch.
+
+    Args:
+        collection: Collection document dictionary to modify in-place.
+
+    Returns:
+        bool: True if bbox_shape was added, False if it was skipped (already exists,
+            no spatial extent, or invalid bbox).
+
+    Notes:
+        - Modifies the collection dictionary in-place by adding a 'bbox_shape' field
+        - Handles both 2D [minx, miny, maxx, maxy] and 3D [minx, miny, minz, maxx, maxy, maxz] bboxes
+        - Uses the first bbox if multiple are present in the collection
+        - Logs warnings for collections with invalid or missing bbox data
+    """
+    collection_id = collection.get("id", "unknown")
+
+    # Check if bbox_shape already exists
+    if "bbox_shape" in collection:
+        logger.debug(
+            f"Collection '{collection_id}' already has bbox_shape field, skipping"
+        )
+        return False
+
+    # Check if collection has spatial extent
+    if "extent" not in collection or "spatial" not in collection["extent"]:
+        logger.warning(f"Collection '{collection_id}' has no spatial extent, skipping")
+        return False
+
+    spatial_extent = collection["extent"]["spatial"]
+    if "bbox" not in spatial_extent or not spatial_extent["bbox"]:
+        logger.warning(
+            f"Collection '{collection_id}' has no bbox in spatial extent, skipping"
+        )
+        return False
+
+    # Get the first bbox (collections can have multiple bboxes, but we use the first one)
+    bbox = (
+        spatial_extent["bbox"][0]
+        if isinstance(spatial_extent["bbox"][0], list)
+        else spatial_extent["bbox"]
+    )
+
+    if len(bbox) < 4:
+        logger.warning(
+            f"Collection '{collection_id}': bbox has insufficient coordinates (length={len(bbox)}), expected at least 4"
+        )
+        return False
+
+    # Extract 2D coordinates (bbox can be 2D [minx, miny, maxx, maxy] or 3D [minx, miny, minz, maxx, maxy, maxz])
+    # For 2D polygon, we only need the x,y coordinates and discard altitude (z) values
+    minx, miny = bbox[0], bbox[1]
+    if len(bbox) == 4:
+        # 2D bbox: [minx, miny, maxx, maxy]
+        maxx, maxy = bbox[2], bbox[3]
+    else:
+        # 3D bbox: [minx, miny, minz, maxx, maxy, maxz]
+        # Extract indices 3,4 for maxx,maxy - discarding altitude at indices 2 (minz) and 5 (maxz)
+        maxx, maxy = bbox[3], bbox[4]
+
+    # Convert bbox to GeoJSON polygon
+    bbox_polygon_coords = bbox2polygon(minx, miny, maxx, maxy)
+    collection["bbox_shape"] = {
+        "type": "Polygon",
+        "coordinates": bbox_polygon_coords,
+    }
+
+    logger.debug(f"Collection '{collection_id}': Added bbox_shape field")
+    return True
+
 
 def validate_refresh(value: Union[str, bool]) -> str:
     """