Improve fields extension implementation

Correctly supports deeply nested property keys in both include and
exclude, as well as improves variable naming, comments, and test cases.

stac_fastapi/pgstac/stac_fastapi/pgstac/utils.py

@@ -19,25 +19,44 @@ def filter_fields(
     if not include and not exclude:
         return item
 
-    # Build a shallow copy of included fields on item
+    # Build a shallow copy of included fields on an item, or a sub-tree of an item
     def include_fields(
-        full_item: Dict[str, Any], fields: Optional[Set[str]]
+        source: Dict[str, Any], fields: Optional[Set[str]]
     ) -> Dict[str, Any]:
         if not fields:
-            return full_item
+            return source
 
         clean_item: Dict[str, Any] = {}
         for key_path in fields or []:
-            keys = key_path.split(".")
-            root = keys[0]
-            if root in full_item:
-                if isinstance(full_item[root], dict) and len(keys) > 1:
-                    # Recurse on "includes" key paths notation for sub-keys
-                    clean_item[root] = include_fields(
-                        full_item[root], fields=set([".".join(keys[1:])])
+            key_path_parts = key_path.split(".")
+            key_root = key_path_parts[0]
+            if key_root in source:
+                if isinstance(source[key_root], dict) and len(key_path_parts) > 1:
+                    # The root of this key path on the item is a dict, and the
+                    # key path indicates a sub-key to be included. Walk the dict
+                    # from the root key and get the full nested value to include.
+                    value = include_fields(
+                        source[key_root], fields=set([".".join(key_path_parts[1:])])
                     )
+
+                    if isinstance(clean_item.get(key_root), dict):
+                        # A previously specified key and sub-keys may have been included
+                        # already, so do a deep merge update if the root key already exists.
+                        dict_deep_update(clean_item[key_root], value)
+                    else:
+                        # The root key does not exist, so add it. Fields
+                        # extension only allows nested referencing on dicts, so
+                        # this won't overwrite anything.
+                        clean_item[key_root] = value
                 else:
-                    clean_item[root] = full_item[root]
+                    # The item value to include is not a dict, or, it is a dict but the
+                    # key path is for the whole value, not a sub-key. Include the entire
+                    # value in the cleaned item.
+                    clean_item[key_root] = source[key_root]
+            else:
+                # The key, or root key of a multi-part key, is not present in the item,
+                # so it is ignored
+                pass
         return clean_item
 
     # For an item built up for included fields, remove excluded fields
@@ -70,3 +89,20 @@ def exclude_fields(
     clean_item = include_fields(item, include)
     clean_item = exclude_fields(clean_item, exclude)
     return Item(**clean_item)
+
+
+def dict_deep_update(merge_to: Dict[str, Any], merge_from: Dict[str, Any]) -> None:
+    """Perform a deep update of two dicts.
+
+    merge_to is updated in-place with the values from merge_from.
+    merge_from values take precedence over existing values in merge_to.
+    """
+    for k, v in merge_from.items():
+        if (
+            k in merge_to
+            and isinstance(merge_to[k], dict)
+            and isinstance(merge_from[k], dict)
+        ):
+            dict_deep_update(merge_to[k], merge_from[k])
+        else:
+            merge_to[k] = v

stac_fastapi/pgstac/tests/resources/test_item.py

@@ -1087,6 +1087,84 @@ async def test_field_extension_exclude_default_includes(
     assert "geometry" not in resp_json["features"][0]
 
 
+async def test_field_extension_include_multiple_subkeys(
+    app_client, load_test_item, load_test_collection
+):
+    """Test that multiple subkeys of an object field are included"""
+    body = {"fields": {"include": ["properties.width", "properties.height"]}}
+
+    resp = await app_client.post("/search", json=body)
+    assert resp.status_code == 200
+    resp_json = resp.json()
+
+    resp_prop_keys = resp_json["features"][0]["properties"].keys()
+    assert set(resp_prop_keys) == set(["width", "height"])
+
+
+async def test_field_extension_include_multiple_deeply_nested_subkeys(
+    app_client, load_test_item, load_test_collection
+):
+    """Test that multiple deeply nested subkeys of an object field are included"""
+    body = {"fields": {"include": ["assets.ANG.type", "assets.ANG.href"]}}
+
+    resp = await app_client.post("/search", json=body)
+    assert resp.status_code == 200
+    resp_json = resp.json()
+
+    resp_assets = resp_json["features"][0]["assets"]
+    assert set(resp_assets.keys()) == set(["ANG"])
+    assert set(resp_assets["ANG"].keys()) == set(["type", "href"])
+
+
+async def test_field_extension_exclude_multiple_deeply_nested_subkeys(
+    app_client, load_test_item, load_test_collection
+):
+    """Test that multiple deeply nested subkeys of an object field are excluded"""
+    body = {"fields": {"exclude": ["assets.ANG.type", "assets.ANG.href"]}}
+
+    resp = await app_client.post("/search", json=body)
+    assert resp.status_code == 200
+    resp_json = resp.json()
+
+    resp_assets = resp_json["features"][0]["assets"]
+    assert len(resp_assets.keys()) > 0
+    assert "type" not in resp_assets["ANG"]
+    assert "href" not in resp_assets["ANG"]
+
+
+async def test_field_extension_exclude_deeply_nested_included_subkeys(
+    app_client, load_test_item, load_test_collection
+):
+    """Test that deeply nested keys of a nested object that was included are excluded"""
+    body = {
+        "fields": {
+            "include": ["assets.ANG.type", "assets.ANG.href"],
+            "exclude": ["assets.ANG.href"],
+        }
+    }
+
+    resp = await app_client.post("/search", json=body)
+    assert resp.status_code == 200
+    resp_json = resp.json()
+
+    resp_assets = resp_json["features"][0]["assets"]
+    assert "type" in resp_assets["ANG"]
+    assert "href" not in resp_assets["ANG"]
+
+
+async def test_field_extension_exclude_links(
+    app_client, load_test_item, load_test_collection
+):
+    """Links have special injection behavior, ensure they can be excluded with the fields extension"""
+    body = {"fields": {"exclude": ["links"]}}
+
+    resp = await app_client.post("/search", json=body)
+    assert resp.status_code == 200
+    resp_json = resp.json()
+
+    assert "links" not in resp_json["features"][0]
+
+
 async def test_search_intersects_and_bbox(app_client):
     """Test POST search intersects and bbox are mutually exclusive (core)"""
     bbox = [-118, 34, -117, 35]