diff --git a/cognite/client/_api/annotations.py b/cognite/client/_api/annotations.py index 953f03ced..ef24ca333 100644 --- a/cognite/client/_api/annotations.py +++ b/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) diff --git a/cognite/client/_api_client.py b/cognite/client/_api_client.py index 74f9d4c8e..b33350593 100644 --- a/cognite/client/_api_client.py +++ b/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} diff --git a/cognite/client/data_classes/_base.py b/cognite/client/data_classes/_base.py index f35617043..cc9cffbcf 100644 --- a/cognite/client/data_classes/_base.py +++ b/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) diff --git a/cognite/client/data_classes/hosted_extractors/sources.py b/cognite/client/data_classes/hosted_extractors/sources.py index 29b5fb676..1a273db27 100644 --- a/cognite/client/data_classes/hosted_extractors/sources.py +++ b/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__() diff --git a/tests/tests_integration/test_api/test_hosted_extractors/test_sources.py b/tests/tests_integration/test_api/test_hosted_extractors/test_sources.py index 611cad3ae..d7c742f74 100644 --- a/tests/tests_integration/test_api/test_hosted_extractors/test_sources.py +++ b/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)