feat: Added update classes

cognite/client/_api/hosted_extractors/sources.py

@@ -5,7 +5,7 @@
 
 from cognite.client._api_client import APIClient
 from cognite.client._constants import DEFAULT_LIMIT_READ
-from cognite.client.data_classes.hosted_extractors.sources import Source, SourceList, SourceWrite
+from cognite.client.data_classes.hosted_extractors.sources import Source, SourceList, SourceUpdate, SourceWrite
 from cognite.client.utils._identifier import IdentifierSequence
 from cognite.client.utils.useful_types import SequenceNotStr
 
@@ -154,6 +154,38 @@ def create(self, items: SourceWrite | Sequence[SourceWrite]) -> Source | SourceL
             input_resource_cls=SourceWrite,  # type: ignore[arg-type]
         )
 
+    @overload
+    def update(self, items: SourceWrite | SourceUpdate) -> Source: ...
+
+    @overload
+    def update(self, items: Sequence[SourceWrite | SourceUpdate]) -> SourceList: ...
+
+    def update(self, items: SourceWrite | SourceUpdate | Sequence[SourceWrite | SourceUpdate]) -> Source | SourceList:
+        """`Update one or more sources. <https://developer.cognite.com/api#tag/Sources/operation/update_sources>`_
+
+        Args:
+            items (SourceWrite | SourceUpdate | Sequence[SourceWrite | SourceUpdate]): Space | Sequence[Space]): Source(s) to update.
+
+        Returns:
+            Source | SourceList: Updated source(s)
+
+        Examples:
+
+            Update source:
+
+                >>> from cognite.client import CogniteClient
+                >>> from cognite.client.data_classes.hosted_extractors import EventHubSourceUpdate
+                >>> client = CogniteClient()
+                >>> source = EventHubSourceUpdate('my_event_hub').event_hub_name.set("My Updated EventHub")
+                >>> res = client.hosted_extractors.sources.update(source)
+        """
+        return self._update_multiple(
+            items=items,  # type: ignore[arg-type]
+            list_cls=SourceList,
+            resource_cls=Source,  # type: ignore[type-abstract]
+            update_cls=SourceUpdate,  # type: ignore[type-abstract]
+        )
+
     def list(
         self,
         limit: int | None = DEFAULT_LIMIT_READ,

cognite/client/data_classes/_base.py

@@ -534,7 +534,7 @@ def __init__(self, update_object: T_CogniteUpdate, name: str) -> None:
         self._update_object = update_object
         self._name = name
 
-    def _set(self, value: None | str | int | bool) -> T_CogniteUpdate:
+    def _set(self, value: None | str | int | bool | dict) -> T_CogniteUpdate:
         if value is None:
             self._update_object._set_null(self._name)
         else:

cognite/client/data_classes/hosted_extractors/__init__.py

@@ -1,12 +1,16 @@
 from cognite.client.data_classes.hosted_extractors.sources import (
     EventHubSource,
+    EventHubSourceUpdate,
     EventHubSourceWrite,
     MQTT3Source,
+    MQTT3SourceUpdate,
     MQTT3SourceWrite,
     MQTT5Source,
+    MQTT5SourceUpdate,
     MQTT5SourceWrite,
     Source,
     SourceList,
+    SourceUpdate,
     SourceWrite,
     SourceWriteList,
 )
@@ -22,4 +26,8 @@
     "SourceList",
     "SourceWrite",
     "SourceWriteList",
+    "SourceUpdate",
+    "MQTT3SourceUpdate",
+    "MQTT5SourceUpdate",
+    "EventHubSourceUpdate",
 ]

cognite/client/data_classes/hosted_extractors/sources.py

@@ -9,9 +9,12 @@
 
 from cognite.client.data_classes._base import (
     CogniteObject,
+    CognitePrimitiveUpdate,
     CogniteResource,
     CogniteResourceList,
+    CogniteUpdate,
     ExternalIDTransformerMixin,
+    PropertySpec,
     T_WriteClass,
     UnknownCogniteObject,
     WriteableCogniteResource,
@@ -101,6 +104,18 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
         return output
 
 
+class SourceUpdate(CogniteUpdate, ABC):
+    _type: ClassVar[str]
+
+    def __init__(self, external_id: str) -> None:
+        super().__init__(external_id=external_id)
+
+    def dump(self, camel_case: Literal[True] = True) -> dict[str, Any]:
+        output = super().dump(camel_case)
+        output["type"] = self._type
+        return output
+
+
 class EventHubSourceWrite(SourceWrite):
     """A hosted extractor source represents an external source system on the internet.
     The source resource in CDF contains all the information the extractor needs to
@@ -193,6 +208,48 @@ def _load_source(cls, resource: dict[str, Any]) -> Self:
         return fast_dict_load(cls, resource, None)
 
 
+class EventHubSourceUpdate(SourceUpdate):
+    _type = "eventHub"
+
+    class _PrimitiveEventHubSourceUpdate(CognitePrimitiveUpdate):
+        def set(self, value: str) -> EventHubSourceUpdate:
+            return self._set(value)
+
+    class _PrimitiveNullableEventHubSourceUpdate(CognitePrimitiveUpdate):
+        def set(self, value: str | None) -> EventHubSourceUpdate:
+            return self._set(value)
+
+    @property
+    def host(self) -> _PrimitiveEventHubSourceUpdate:
+        return EventHubSourceUpdate._PrimitiveEventHubSourceUpdate(self, "host")
+
+    @property
+    def event_hub_name(self) -> _PrimitiveEventHubSourceUpdate:
+        return EventHubSourceUpdate._PrimitiveEventHubSourceUpdate(self, "eventHubName")
+
+    @property
+    def key_name(self) -> _PrimitiveEventHubSourceUpdate:
+        return EventHubSourceUpdate._PrimitiveEventHubSourceUpdate(self, "keyName")
+
+    @property
+    def key_value(self) -> _PrimitiveEventHubSourceUpdate:
+        return EventHubSourceUpdate._PrimitiveEventHubSourceUpdate(self, "keyValue")
+
+    @property
+    def consumer_group(self) -> _PrimitiveNullableEventHubSourceUpdate:
+        return EventHubSourceUpdate._PrimitiveNullableEventHubSourceUpdate(self, "consumerGroup")
+
+    @classmethod
+    def _get_update_properties(cls, item: CogniteResource | None = None) -> list[PropertySpec]:
+        return [
+            PropertySpec("host", is_nullable=False),
+            PropertySpec("event_hub_name", is_nullable=False),
+            PropertySpec("key_name", is_nullable=False),
+            PropertySpec("key_value", is_nullable=False),
+            PropertySpec("consumer_group", is_nullable=True),
+        ]
+
+
 @dataclass
 class MQTTAuthenticationWrite(CogniteObject, ABC):
     _type: ClassVar[str]
@@ -417,6 +474,67 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
         return output
 
 
+class _MQTTUpdate(SourceUpdate, ABC):
+    class _HostUpdate(CognitePrimitiveUpdate):
+        def set(self, value: str) -> _MQTTUpdate:
+            return self._set(value)
+
+    class _PortUpdate(CognitePrimitiveUpdate):
+        def set(self, value: int | None) -> _MQTTUpdate:
+            return self._set(value)
+
+    class _AuthenticationUpdate(CognitePrimitiveUpdate):
+        def set(self, value: MQTTAuthentication | None) -> _MQTTUpdate:
+            return self._set(value.dump() if value else None)
+
+    class _UseTlsUpdate(CognitePrimitiveUpdate):
+        def set(self, value: bool) -> _MQTTUpdate:
+            return self._set(value)
+
+    class _CACertificateUpdate(CognitePrimitiveUpdate):
+        def set(self, value: CACertificate | None) -> _MQTTUpdate:
+            return self._set(value.dump() if value else None)
+
+    class _AuthCertificateUpdate(CognitePrimitiveUpdate):
+        def set(self, value: AuthCertificate | None) -> _MQTTUpdate:
+            return self._set(value.dump() if value else None)
+
+    @property
+    def host(self) -> _HostUpdate:
+        return _MQTTUpdate._HostUpdate(self, "host")
+
+    @property
+    def port(self) -> _PortUpdate:
+        return _MQTTUpdate._PortUpdate(self, "port")
+
+    @property
+    def authentication(self) -> _AuthenticationUpdate:
+        return _MQTTUpdate._AuthenticationUpdate(self, "authentication")
+
+    @property
+    def useTls(self) -> _UseTlsUpdate:
+        return _MQTTUpdate._UseTlsUpdate(self, "useTls")
+
+    @property
+    def ca_certificate(self) -> _CACertificateUpdate:
+        return _MQTTUpdate._CACertificateUpdate(self, "caCertificate")
+
+    @property
+    def auth_certificate(self) -> _AuthCertificateUpdate:
+        return _MQTTUpdate._AuthCertificateUpdate(self, "authCertificate")
+
+    @classmethod
+    def _get_update_properties(cls, item: CogniteResource | None = None) -> list[PropertySpec]:
+        return [
+            PropertySpec("host", is_nullable=False),
+            PropertySpec("port", is_nullable=True),
+            PropertySpec("authentication", is_nullable=True, is_container=True),
+            PropertySpec("useTls", is_nullable=False),
+            PropertySpec("ca_certificate", is_nullable=True, is_container=True),
+            PropertySpec("auth_certificate", is_nullable=True, is_container=True),
+        ]
+
+
 class MQTT3SourceWrite(_MQTTSourceWrite):
     _type = "mqtt3"
 
@@ -433,6 +551,14 @@ class MQTT5Source(_MQTTSource):
     _type = "mqtt5"
 
 
+class MQTT3SourceUpdate(_MQTTUpdate):
+    _type = "mqtt3"
+
+
+class MQTT5SourceUpdate(_MQTTUpdate):
+    _type = "mqtt5"
+
+
 class SourceWriteList(CogniteResourceList[SourceWrite], ExternalIDTransformerMixin):
     _RESOURCE = SourceWrite