fix: update with write object

cognite/client/_api/annotations.py

@@ -101,6 +101,7 @@ def _convert_resource_to_patch_object(
         resource: CogniteResource,
         update_attributes: list[PropertySpec],
         mode: Literal["replace_ignore_null", "patch", "replace"] = "replace_ignore_null",
+        identifying_properties: dict[str, Any] | None = None,
     ) -> dict[str, dict[str, dict]]:
         if not isinstance(resource, Annotation):
             return APIClient._convert_resource_to_patch_object(resource, update_attributes)

cognite/client/_api_client.py

@@ -1038,7 +1038,12 @@ def _update_multiple(
         for index, item in enumerate(item_list):
             if isinstance(item, CogniteResource):
                 patch_objects.append(
-                    self._convert_resource_to_patch_object(item, update_cls._get_update_properties(item), mode)
+                    self._convert_resource_to_patch_object(
+                        item,
+                        update_cls._get_update_properties(item),
+                        mode,
+                        update_cls._get_extra_identifying_properties(item),
+                    )
                 )
             elif isinstance(item, CogniteUpdate):
                 patch_objects.append(item.dump(camel_case=True))
@@ -1218,6 +1223,7 @@ def _convert_resource_to_patch_object(
         resource: CogniteResource,
         update_attributes: list[PropertySpec],
         mode: Literal["replace_ignore_null", "patch", "replace"] = "replace_ignore_null",
+        identifying_properties: dict[str, Any] | None = None,
     ) -> dict[str, dict[str, dict]]:
         dumped_resource = resource.dump(camel_case=True)
         has_id = "id" in dumped_resource
@@ -1233,6 +1239,9 @@ def _convert_resource_to_patch_object(
         elif has_external_id:
             patch_object["externalId"] = dumped_resource.pop("externalId")
 
+        if identifying_properties:
+            patch_object.update(identifying_properties)
+
         update: dict[str, dict] = cls._clear_all_attributes(update_attributes) if mode == "replace" else {}
 
         update_attribute_by_name = {prop.name: prop for prop in update_attributes}

cognite/client/data_classes/_base.py

@@ -525,6 +525,12 @@ def dump(self, camel_case: Literal[True] = True) -> dict[str, Any]:
     def _get_update_properties(cls, item: CogniteResource | None = None) -> list[PropertySpec]:
         raise NotImplementedError
 
+    @classmethod
+    def _get_extra_identifying_properties(cls, item: CogniteResource | None = None) -> dict[str, Any]:
+        # This method is used to provide additional identifying properties for the update object.
+        # It is intended to be overridden by subclasses that need to provide additional identifying properties.
+        return {}
+
 
 T_CogniteUpdate = TypeVar("T_CogniteUpdate", bound=CogniteUpdate)
 

cognite/client/data_classes/hosted_extractors/sources.py

@@ -114,6 +114,18 @@ def dump(self, camel_case: Literal[True] = True) -> dict[str, Any]:
         output["type"] = self._type
         return output
 
+    @classmethod
+    def _get_update_properties(cls, item: CogniteResource | None = None) -> list[PropertySpec]:
+        if item is None or not isinstance(item, SourceWrite):
+            return []
+        return _SOURCE_UPDATE_BY_TYPE[item._type]._get_update_properties(item)
+
+    @classmethod
+    def _get_extra_identifying_properties(cls, item: CogniteResource | None = None) -> dict[str, Any]:
+        if not isinstance(item, SourceWrite):
+            return {}
+        return {"type": item._type}
+
 
 class EventHubSourceWrite(SourceWrite):
     """A hosted extractor source represents an external source system on the internet.
@@ -605,6 +617,12 @@ def as_write(
     if hasattr(subclass, "_type")
 }
 
+_SOURCE_UPDATE_BY_TYPE: dict[str, type[SourceUpdate]] = {
+    subclass._type: subclass  # type: ignore[type-abstract, misc]
+    for subclass in itertools.chain(SourceUpdate.__subclasses__(), _MQTTUpdate.__subclasses__())
+    if hasattr(subclass, "_type")
+}
+
 _MQTTAUTHENTICATION_WRITE_CLASS_BY_TYPE: dict[str, type[MQTTAuthenticationWrite]] = {
     subclass._type: subclass  # type: ignore[type-abstract]
     for subclass in MQTTAuthenticationWrite.__subclasses__()

tests/tests_integration/test_api/test_hosted_extractors/test_sources.py

@@ -66,3 +66,30 @@ def test_list(self, cognite_client: CogniteClient) -> None:
         res = cognite_client.hosted_extractors.sources.list(limit=1)
         assert len(res) == 1
         assert isinstance(res, SourceList)
+
+    def test_update_using_write_object(self, cognite_client: CogniteClient) -> None:
+        my_hub = EventHubSourceWrite(
+            external_id=f"toupdatate-{random_string(10)}",
+            host="myHost",
+            key_name="myKeyName",
+            key_value="myKey",
+            event_hub_name="myEventHub",
+        )
+        created: EventHubSource | None = None
+        try:
+            created = cognite_client.hosted_extractors.sources.create(my_hub)
+
+            my_new_hub = EventHubSourceWrite(
+                external_id=created.external_id,
+                host="updatedHost",
+                key_name="updatedKeyName",
+                key_value="updatedKey",
+                event_hub_name="updatedEventHub",
+            )
+
+            updated = cognite_client.hosted_extractors.sources.update(my_new_hub)
+
+            assert updated.host == my_new_hub.host
+        finally:
+            if created:
+                cognite_client.hosted_extractors.sources.delete(created.external_id, ignore_unknown_ids=True)