fix(LAB-3532): handle external_ids arg for add_metadata and set_metadata (#1891)

Co-authored-by: paulruelle <[email protected]>

Author: RuellePaul

src/kili/entrypoints/mutations/asset/__init__.py

@@ -10,7 +10,7 @@
 from kili.adapters.kili_api_gateway.helpers.queries import QueryOptions
 from kili.core.helpers import is_empty_list_with_warning
 from kili.core.utils.pagination import mutate_from_paginated_call
-from kili.domain.asset import AssetFilters, AssetId
+from kili.domain.asset import AssetExternalId, AssetFilters, AssetId
 from kili.domain.project import ProjectId
 from kili.entrypoints.base import BaseOperationEntrypointMixin
 from kili.entrypoints.mutations.asset.helpers import (
@@ -412,16 +412,18 @@ def generate_variables(batch: Dict) -> Dict:
     def add_metadata(
         self,
         json_metadata: List[Dict[str, Union[str, int, float]]],
-        asset_ids: List[str],
         project_id: str,
+        asset_ids: Optional[List[str]] = None,
+        external_ids: Optional[List[str]] = None,
     ) -> List[Dict[Literal["id"], str]]:
         """Add metadata to assets without overriding existing metadata.
 
         Args:
             json_metadata: List of metadata dictionaries to add to each asset.
                 Each dictionary contains key/value pairs to be added to the asset's metadata.
-            asset_ids: The asset IDs to modify.
             project_id: The project ID.
+            asset_ids: The asset IDs to modify.
+            external_ids: The external asset IDs to modify (if `asset_ids` is not already provided).
 
         Returns:
             A list of dictionaries with the asset ids.
@@ -432,21 +434,39 @@ def add_metadata(
                         {"key1": "value1", "key2": "value2"},
                         {"key3": "value3"}
                     ],
-                    asset_ids=["ckg22d81r0jrg0885unmuswj8", "ckg22d81s0jrh0885pdxfd03n"],
-                    project_id="cm92to3cx012u7l0w6kij9qvx"
+                    project_id="cm92to3cx012u7l0w6kij9qvx",
+                    asset_ids=["ckg22d81r0jrg0885unmuswj8", "ckg22d81s0jrh0885pdxfd03n"]
+                )
+
+                # Or using external IDs
+            >>> kili.add_metadata(
+                    json_metadata=[
+                        {"key1": "value1", "key2": "value2"},
+                        {"key3": "value3"}
+                    ],
+                    project_id="cm92to3cx012u7l0w6kij9qvx",
+                    external_ids=["asset1", "asset2"]
                 )
         """
         if is_empty_list_with_warning("add_metadata", "json_metadata", json_metadata):
             return []
 
+        if (asset_ids is not None and external_ids is not None) or (
+            asset_ids is None and external_ids is None
+        ):
+            raise MissingArgumentError("Please provide either `asset_ids` or `external_ids`.")
+
         assets = self.kili_api_gateway.list_assets(
             AssetFilters(
-                project_id=ProjectId(project_id), asset_id_in=cast(List[AssetId], asset_ids)
+                project_id=ProjectId(project_id),
+                asset_id_in=cast(List[AssetId], asset_ids),
+                external_id_in=cast(List[AssetExternalId], external_ids),
             ),
             ["id", "jsonMetadata"],
             QueryOptions(disable_tqdm=True),
         )
 
+        resolved_asset_ids = []
         json_metadatas = []
         for i, asset in enumerate(assets):
             current_metadata = asset.get("jsonMetadata", {}) if asset.get("jsonMetadata") else {}
@@ -455,26 +475,29 @@ def add_metadata(
             current_metadata.update(new_metadata)
 
             json_metadatas.append(current_metadata)
+            resolved_asset_ids.append(asset["id"])
 
         return self.update_properties_in_assets(
-            asset_ids=asset_ids,
+            asset_ids=cast(List[str], resolved_asset_ids),
             json_metadatas=json_metadatas,
         )
 
     @typechecked
     def set_metadata(
         self,
         json_metadata: List[Dict[str, Union[str, int, float]]],
-        asset_ids: List[str],
         project_id: str,
+        asset_ids: Optional[List[str]] = None,
+        external_ids: Optional[List[str]] = None,
     ) -> List[Dict[Literal["id"], str]]:
         """Set metadata on assets, replacing any existing metadata.
 
         Args:
             json_metadata: List of metadata dictionaries to set on each asset.
                 Each dictionary contains key/value pairs to be set as the asset's metadata.
-            asset_ids: The asset IDs to modify.
             project_id: The project ID.
+            asset_ids: The asset IDs to modify (if `external_ids` is not already provided).
+            external_ids: The external asset IDs to modify (if `asset_ids` is not already provided).
 
         Returns:
             A list of dictionaries with the asset ids.
@@ -485,21 +508,39 @@ def set_metadata(
                         {"key1": "value1", "key2": "value2"},
                         {"key3": "value3"}
                     ],
-                    asset_ids=["ckg22d81r0jrg0885unmuswj8", "ckg22d81s0jrh0885pdxfd03n"],
                     project_id="cm92to3cx012u7l0w6kij9qvx"
+                    asset_ids=["ckg22d81r0jrg0885unmuswj8", "ckg22d81s0jrh0885pdxfd03n"]
+                )
+
+                # Or using external IDs
+            >>> kili.set_metadata(
+                    json_metadata=[
+                        {"key1": "value1", "key2": "value2"},
+                        {"key3": "value3"}
+                    ],
+                    project_id="cm92to3cx012u7l0w6kij9qvx",
+                    external_ids=["asset1", "asset2"]
                 )
         """
         if is_empty_list_with_warning("set_metadata", "json_metadata", json_metadata):
             return []
 
+        if (asset_ids is not None and external_ids is not None) or (
+            asset_ids is None and external_ids is None
+        ):
+            raise MissingArgumentError("Please provide either `asset_ids` or `external_ids`.")
+
         assets = self.kili_api_gateway.list_assets(
             AssetFilters(
-                project_id=ProjectId(project_id), asset_id_in=cast(List[AssetId], asset_ids)
+                project_id=ProjectId(project_id),
+                asset_id_in=cast(List[AssetId], asset_ids),
+                external_id_in=cast(List[AssetExternalId], external_ids),
             ),
             ["id", "jsonMetadata"],
             QueryOptions(disable_tqdm=True),
         )
 
+        resolved_asset_ids = []
         json_metadatas = []
         for i, asset in enumerate(assets):
             current_metadata = asset.get("jsonMetadata", {}) if asset.get("jsonMetadata") else {}
@@ -513,9 +554,10 @@ def set_metadata(
             preserved_metadata.update(new_metadata)
 
             json_metadatas.append(preserved_metadata)
+            resolved_asset_ids.append(asset["id"])
 
         return self.update_properties_in_assets(
-            asset_ids=asset_ids,
+            asset_ids=cast(List[str], resolved_asset_ids),
             json_metadatas=json_metadatas,
         )
 

tests/integration/presentation/test_metadata.py

@@ -6,7 +6,7 @@
 
 from kili.adapters.kili_api_gateway.helpers.queries import QueryOptions
 from kili.adapters.kili_api_gateway.kili_api_gateway import KiliAPIGateway
-from kili.domain.asset import AssetFilters, AssetId
+from kili.domain.asset import AssetExternalId, AssetFilters, AssetId
 from kili.domain.project import ProjectId
 from kili.entrypoints.mutations.asset import MutationsAsset
 
@@ -229,3 +229,123 @@ def test_multiple_assets_with_different_metadata_structures(mocker: pytest_mock.
 
     assert add_metadata_call_args["json_metadatas"] == expected_add_metadata
     assert set_metadata_call_args["json_metadatas"] == expected_set_metadata
+
+
+def test_add_metadata_with_external_ids(mocker: pytest_mock.MockerFixture):
+    """Test that add_metadata works correctly with external IDs."""
+    kili_api_gateway = mocker.Mock(spec=KiliAPIGateway)
+
+    existing_assets = [
+        {
+            "id": "asset1",
+            "jsonMetadata": {"text": "Some text", "existing1": "value1"},
+        },
+        {
+            "id": "asset2",
+            "jsonMetadata": {"imageUrl": "http://example.com/image.jpg", "existing2": "value2"},
+        },
+    ]
+
+    kili_api_gateway.list_assets.return_value = existing_assets
+
+    update_mock = mocker.patch.object(MutationsAsset, "update_properties_in_assets")
+    update_mock.return_value = [{"id": "asset1"}, {"id": "asset2"}]
+
+    mutations_asset = MutationsAsset()
+    mutations_asset.kili_api_gateway = kili_api_gateway
+    mutations_asset.graphql_client = mocker.Mock()
+
+    project_id = "project1"
+    external_ids = ["ext1", "ext2"]
+    new_metadata: List[Dict[str, Union[str, float, int]]] = [
+        {"new_key1": "new_value1"},
+        {"new_key2": "new_value2"},
+    ]
+
+    result = mutations_asset.add_metadata(
+        json_metadata=new_metadata, external_ids=external_ids, project_id=project_id
+    )
+
+    kili_api_gateway.list_assets.assert_called_once_with(
+        AssetFilters(
+            project_id=ProjectId(project_id),
+            external_id_in=cast(List[AssetExternalId], external_ids),
+        ),
+        ["id", "jsonMetadata"],
+        QueryOptions(disable_tqdm=True),
+    )
+
+    update_mock.assert_called_once()
+    call_args = update_mock.call_args[1]
+
+    expected_metadata = [
+        {"text": "Some text", "existing1": "value1", "new_key1": "new_value1"},
+        {
+            "imageUrl": "http://example.com/image.jpg",
+            "existing2": "value2",
+            "new_key2": "new_value2",
+        },
+    ]
+
+    assert call_args["json_metadatas"] == expected_metadata
+    assert result == [{"id": "asset1"}, {"id": "asset2"}]
+
+
+def test_set_metadata_with_external_ids(mocker: pytest_mock.MockerFixture):
+    """Test that set_metadata works correctly with external IDs."""
+    kili_api_gateway = mocker.Mock(spec=KiliAPIGateway)
+
+    existing_assets = [
+        {
+            "id": "asset1",
+            "jsonMetadata": {"text": "Some text", "existing1": "value1", "should_be_removed": True},
+        },
+        {
+            "id": "asset2",
+            "jsonMetadata": {
+                "imageUrl": "http://example.com/image.jpg",
+                "existing2": "value2",
+                "also_remove": "yes",
+            },
+        },
+    ]
+
+    kili_api_gateway.list_assets.return_value = existing_assets
+
+    update_mock = mocker.patch.object(MutationsAsset, "update_properties_in_assets")
+    update_mock.return_value = [{"id": "asset1"}, {"id": "asset2"}]
+
+    mutations_asset = MutationsAsset()
+    mutations_asset.kili_api_gateway = kili_api_gateway
+    mutations_asset.graphql_client = mocker.Mock()
+
+    project_id = "project1"
+    external_ids = ["ext1", "ext2"]
+    new_metadata: List[Dict[str, Union[str, float, int]]] = [
+        {"new_key1": "new_value1"},
+        {"new_key2": "new_value2"},
+    ]
+
+    result = mutations_asset.set_metadata(
+        json_metadata=new_metadata, external_ids=external_ids, project_id=project_id
+    )
+
+    kili_api_gateway.list_assets.assert_called_once_with(
+        AssetFilters(
+            project_id=ProjectId(project_id),
+            external_id_in=cast(List[AssetExternalId], external_ids),
+        ),
+        ["id", "jsonMetadata"],
+        QueryOptions(disable_tqdm=True),
+    )
+
+    update_mock.assert_called_once()
+    call_args = update_mock.call_args[1]
+
+    expected_metadata = [
+        {"text": "Some text", "new_key1": "new_value1"},
+        {"imageUrl": "http://example.com/image.jpg", "new_key2": "new_value2"},
+    ]
+
+    assert call_args["json_metadatas"] == expected_metadata
+    assert result == [{"id": "asset1"}, {"id": "asset2"}]