Adding flag for asset indexing.

Adding scripts for reindexing existing indexes.
Adding tests for serializers.

Author: rhysrevans3

scripts/reindex_elasticsearch.py

@@ -0,0 +1,61 @@
+import time
+
+from stac_fastapi.elasticsearch.config import AsyncElasticsearchSettings
+from stac_fastapi.elasticsearch.database_logic import create_index_templates
+from stac_fastapi.sfeos_helpers.mappings import COLLECTIONS_INDEX, ITEMS_INDEX_PREFIX
+
+
+async def reindex():
+    """Reindex all STAC indexes for mapping update"""
+    client = AsyncElasticsearchSettings().create_client
+
+    indexes = await client.indices.get_alias(name=COLLECTIONS_INDEX)
+    indexes.update(await client.indices.get_alias(name=f"{ITEMS_INDEX_PREFIX}*"))
+
+    create_index_templates()
+
+    for index, aliases in indexes.items():
+        name, version = index.rsplit("-", 1)
+        new_index = f"{name}-{str(int(version) + 1).zfill(6)}"
+        await client.options(ignore_status=400).indices.create(index=new_index)
+
+        reindex_resp = await client.reindex(
+            dest={"index": new_index}, source={"index": [index]}, wait_for_completion=False
+        )
+
+        task_id = reindex_resp["task"]
+
+        old_count = await client.count(index=index)
+        reindex_complete = False
+        while not reindex_complete:
+            task_resp = await client.tasks.get(task_id=task_id)
+
+            if "completed" in task_resp and task_resp["completed"]:
+                new_count = await client.count(index=new_index)
+
+                if new_count["count"] == old_count["count"]:
+                    reindex_complete = True
+
+                else:
+                    print(f"Reindex failed for {index} with error: mismatch count")
+
+            elif "error" in task_resp:
+                reindex_complete = True
+                print(f"Reindex failed for {index} with error: {task_resp['error']}")
+
+            time.sleep(60)
+
+        actions = []
+        for alias in aliases["aliases"]:
+            actions.extend(
+                [
+                    {"add": {"index": new_index, "alias": alias}},
+                    {"remove": {"index": index, "alias": alias}},
+                ]
+            )
+
+        await client.indices.update_aliases(actions=actions)
+
+
+if __name__ == "__main__":
+    reindex()

scripts/reindex_opensearch.py

@@ -0,0 +1,61 @@
+import time
+
+from stac_fastapi.opensearch.config import AsyncOpensearchSettings
+from stac_fastapi.opensearch.database_logic import create_index_templates
+from stac_fastapi.sfeos_helpers.mappings import COLLECTIONS_INDEX, ITEMS_INDEX_PREFIX
+
+
+async def reindex():
+    """Reindex all STAC indexes for mapping update"""
+    client = AsyncOpensearchSettings().create_client
+
+    indexes = await client.indices.get_alias(name=COLLECTIONS_INDEX)
+    indexes.update(await client.indices.get_alias(name=f"{ITEMS_INDEX_PREFIX}*"))
+
+    create_index_templates()
+
+    for index, aliases in indexes.items():
+        name, version = index.rsplit("-", 1)
+        new_index = f"{name}-{str(int(version) + 1).zfill(6)}"
+        await client.options(ignore_status=400).indices.create(index=new_index)
+
+        reindex_resp = await client.reindex(
+            dest={"index": new_index}, source={"index": [index]}, wait_for_completion=False
+        )
+
+        task_id = reindex_resp["task"]
+
+        old_count = await client.count(index=index)
+        reindex_complete = False
+        while not reindex_complete:
+            task_resp = await client.tasks.get(task_id=task_id)
+
+            if "completed" in task_resp and task_resp["completed"]:
+                new_count = await client.count(index=new_index)
+
+                if new_count["count"] == old_count["count"]:
+                    reindex_complete = True
+
+                else:
+                    print(f"Reindex failed for {index} with error: mismatch count")
+
+            elif "error" in task_resp:
+                reindex_complete = True
+                print(f"Reindex failed for {index} with error: {task_resp['error']}")
+
+            time.sleep(60)
+
+        actions = []
+        for alias in aliases["aliases"]:
+            actions.extend(
+                [
+                    {"add": {"index": new_index, "alias": alias}},
+                    {"remove": {"index": index, "alias": alias}},
+                ]
+            )
+
+        await client.indices.update_aliases(actions=actions)
+
+
+if __name__ == "__main__":
+    reindex()

stac_fastapi/core/stac_fastapi/core/serializers.py

@@ -9,6 +9,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 get_bool_env
 from stac_fastapi.types import stac as stac_types
 from stac_fastapi.types.links import ItemLinks, resolve_links
 
@@ -66,9 +67,10 @@ def stac_to_db(cls, stac_data: stac_types.Item, base_url: str) -> stac_types.Ite
         item_links = resolve_links(stac_data.get("links", []), base_url)
         stac_data["links"] = item_links
 
-        stac_data["assets"] = [
-            {"es_key": k, **v} for k, v in stac_data.get("assets", {}).items()
-        ]
+        if get_bool_env("STAC_INDEX_ASSETS"):
+            stac_data["assets"] = [
+                {"es_key": k, **v} for k, v in stac_data.get("assets", {}).items()
+            ]
 
         now = now_to_rfc3339_str()
         if "created" not in stac_data["properties"]:
@@ -97,6 +99,12 @@ def db_to_stac(cls, item: dict, base_url: str) -> stac_types.Item:
         if original_links:
             item_links += resolve_links(original_links, base_url)
 
+        if get_bool_env("STAC_INDEX_ASSETS"):
+            assets = {a.pop("es_key"): a for a in item.get("assets", [])}
+
+        else:
+            assets = item.get("assets", {})
+
         return stac_types.Item(
             type="Feature",
             stac_version=item.get("stac_version", ""),
@@ -107,7 +115,7 @@ def db_to_stac(cls, item: dict, base_url: str) -> stac_types.Item:
             bbox=item.get("bbox", []),
             properties=item.get("properties", {}),
             links=item_links,
-            assets={a.pop("es_key"): a for a in item.get("assets", [])},
+            assets=assets,
         )
 
 
@@ -132,9 +140,12 @@ def stac_to_db(
         collection["links"] = resolve_links(
             collection.get("links", []), str(request.base_url)
         )
-        collection["assets"] = [
-            {"es_key": k, **v} for k, v in collection.get("assets", {}).items()
-        ]
+
+        if get_bool_env("STAC_INDEX_ASSETS"):
+            collection["assets"] = [
+                {"es_key": k, **v} for k, v in collection.get("assets", {}).items()
+            ]
+
         return collection
 
     @classmethod
@@ -181,9 +192,13 @@ def db_to_stac(
             collection_links += resolve_links(original_links, str(request.base_url))
         collection["links"] = collection_links
 
-        collection["assets"] = {
-            a.pop("es_key"): a for a in collection.get("assets", [])
-        }
+        if get_bool_env("STAC_INDEX_ASSETS"):
+            collection["assets"] = {
+                a.pop("es_key"): a for a in collection.get("assets", [])
+            }
+
+        else:
+            collection["assets"] = collection.get("assets", {})
 
         # Return the stac_types.Collection object
         return stac_types.Collection(**collection)

stac_fastapi/sfeos_helpers/stac_fastapi/sfeos_helpers/mappings.py

@@ -28,6 +28,8 @@
 import os
 from typing import Any, Dict, Literal, Protocol
 
+from stac_fastapi.core.utilities import get_bool_env
+
 
 # stac_pydantic classes extend _GeometryBase, which doesn't have a type field,
 # So create our own Protocol for typing
@@ -134,7 +136,7 @@ class Geometry(Protocol):  # noqa
         "id": {"type": "keyword"},
         "collection": {"type": "keyword"},
         "geometry": {"type": "geo_shape"},
-        "assets": {"type": "object"},
+        "assets": {"type": "object", "enabled": get_bool_env("STAC_INDEX_ASSETS")},
         "links": {"type": "object", "enabled": False},
         "properties": {
             "type": "object",
@@ -162,7 +164,7 @@ class Geometry(Protocol):  # noqa
         "extent.temporal.interval": {"type": "date"},
         "providers": {"type": "object", "enabled": False},
         "links": {"type": "object", "enabled": False},
-        "item_assets": {"type": "object", "enabled": False},
+        "item_assets": {"type": "object", "enabled": get_bool_env("STAC_INDEX_ASSETS")},
     },
 }
 

stac_fastapi/tests/database/test_serializers.py

@@ -0,0 +1,76 @@
+import pytest
+
+from stac_fastapi.core.serializers import CollectionSerializer, ItemSerializer
+from stac_fastapi.types import stac as stac_types
+
+from ..conftest import MockRequest
+
+
+@pytest.mark.asyncio
+async def test_item_serializer(monkeypatch, load_test_data):
+    monkeypatch.setenv("STAC_INDEX_ASSETS", "false")
+
+    request = MockRequest
+
+    item_data = load_test_data("test_item.json")
+
+    item = stac_types.Item(**item_data)
+
+    serialized_item = ItemSerializer.stac_to_db(stac_data=item, base_url=str(request.base_url))
+
+    unserialized_item = ItemSerializer.db_to_stac(item=serialized_item, base_url=str(request.base_url))
+
+    assert unserialized_item == item
+
+
+@pytest.mark.asyncio
+async def test_item_serializer_with_asset_indexing(monkeypatch, load_test_data):
+    monkeypatch.setenv("STAC_INDEX_ASSETS", "true")
+
+    request = MockRequest
+
+    item_data = load_test_data("test_item.json")
+
+    item = stac_types.Item(**item_data)
+
+    serialized_item = ItemSerializer.stac_to_db(stac_data=item, base_url=str(request.base_url))
+
+    unserialized_item = ItemSerializer.db_to_stac(item=serialized_item, base_url=str(request.base_url))
+
+    assert unserialized_item == item
+
+
+@pytest.mark.asyncio
+async def test_collection_serializer(monkeypatch, load_test_data):
+    monkeypatch.setenv("STAC_INDEX_ASSETS", "false")
+
+    request = MockRequest
+    collection_data = load_test_data("test_collection.json")
+
+    collection = stac_types.Collection(**collection_data)
+
+    serialized_collection = CollectionSerializer.stac_to_db(collection=collection, request=request)
+
+    unserialized_collection = CollectionSerializer.db_to_stac(
+        collection=serialized_collection, request=request, extensions=serialized_collection["stac_extensions"]
+    )
+
+    assert unserialized_collection == collection
+
+
+@pytest.mark.asyncio
+async def test_collection_serializer_with_asset_indexing(monkeypatch, load_test_data):
+    monkeypatch.setenv("STAC_INDEX_ASSETS", "true")
+
+    request = MockRequest
+    collection_data = load_test_data("test_collection.json")
+
+    collection = stac_types.Collection(**collection_data)
+
+    serialized_collection = CollectionSerializer.stac_to_db(collection=collection, request=request)
+
+    unserialized_collection = CollectionSerializer.db_to_stac(
+        collection=serialized_collection, request=request, extensions=serialized_collection["stac_extensions"]
+    )
+
+    assert unserialized_collection == collection