feat(transformations): Added rotate keys api endpoint support in sdk

cognite/client/_api/transformations/externaldata.py

@@ -9,6 +9,9 @@
     ExternalDataSource,
     ExternalDataSourceFormat,
     ExternalDataSourceList,
+    ExternalDataSourceRotatedKeys,
+    ExternalDataSourceRotatedKeysList,
+    ExternalDataSourceRotateKeys,
     ExternalDataSourceUsability,
     ExternalDataSourceWrite,
 )
@@ -23,6 +26,7 @@
 
 class TransformationExternalDataSourcesAPI(APIClient):
     _RESOURCE_PATH = "/transformations/externaldata"
+    _ROTATE_KEYS_LIMIT = 10
 
     def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: AsyncCogniteClient) -> None:
         super().__init__(config, api_version, cognite_client)
@@ -34,6 +38,11 @@ def __init__(self, config: ClientConfig, api_version: str | None, cognite_client
             sdk_maturity="alpha",
             feature_name="Retrieving Transformation External Data Sources by external ID",
         )
+        self._rotate_keys_warning = FeaturePreviewWarning(
+            api_maturity="alpha",
+            sdk_maturity="alpha",
+            feature_name="Rotating keys of Transformation External Data Sources",
+        )
         self._RETRIEVE_LIMIT = 100  # ExternalDataSource.ByIds.maxItemsPerRequest
 
     @overload
@@ -286,6 +295,66 @@ async def delete(self, external_id: str | SequenceNotStr[str]) -> None:
             wrap_ids=True,
         )
 
+    @overload
+    async def rotate_keys(self, rotation: ExternalDataSourceRotateKeys) -> ExternalDataSourceRotatedKeys: ...
+
+    @overload
+    async def rotate_keys(
+        self, rotation: Sequence[ExternalDataSourceRotateKeys]
+    ) -> ExternalDataSourceRotatedKeysList: ...
+
+    async def rotate_keys(
+        self, rotation: ExternalDataSourceRotateKeys | Sequence[ExternalDataSourceRotateKeys]
+    ) -> ExternalDataSourceRotatedKeys | ExternalDataSourceRotatedKeysList:
+        """`Rotate the keys of Snowflake external data sources <https://api-docs.cognite.com/20230101-alpha/tag/Transformation-External-Data-Sources/operation/rotateExternalDataSourceKeys>`_.
+
+        The server generates a new key pair for each data source and returns the new public key, which you must
+        register with the Snowflake user. Only Snowflake data sources support this. The operation is atomic: if any
+        item fails, no keys are rotated. At most 10 data sources can be rotated per call, with unique external IDs.
+
+        Args:
+            rotation (ExternalDataSourceRotateKeys | Sequence[ExternalDataSourceRotateKeys]): The rotation(s) to perform. Each ``expiry_time`` must be in the future and at most one year from now.
+
+        Returns:
+            ExternalDataSourceRotatedKeys | ExternalDataSourceRotatedKeysList: The new public key and expiry time of each rotated data source.
+
+        Raises:
+            ValueError: If more than 10 rotations are passed.
+
+        Examples:
+
+            Rotate the keys of a Snowflake data source, valid for 30 days:
+
+                >>> from cognite.client import CogniteClient, AsyncCogniteClient
+                >>> from cognite.client.data_classes.transformations.externaldata import (
+                ...     ExternalDataSourceRotateKeys,
+                ... )
+                >>> from cognite.client.utils._time import timestamp_to_ms
+                >>> client = CogniteClient()
+                >>> # async_client = AsyncCogniteClient()  # another option
+                >>> expiry_time = timestamp_to_ms("30d-ahead")
+                >>> res = client.transformations.external_data_sources.rotate_keys(
+                ...     ExternalDataSourceRotateKeys("snowflake-analytics-prod", expiry_time)
+                ... )
+                >>> new_public_key = res.public_key
+        """
+        self._rotate_keys_warning.warn()
+        items = [rotation] if isinstance(rotation, ExternalDataSourceRotateKeys) else list(rotation)
+        if not items:
+            return ExternalDataSourceRotatedKeysList([])
+        if len(items) > self._ROTATE_KEYS_LIMIT:
+            raise ValueError(
+                f"At most {self._ROTATE_KEYS_LIMIT} data sources can be rotated per call (the operation is atomic "
+                f"and cannot be split), got {len(items)}"
+            )
+        res = await self._post(
+            url_path=self._RESOURCE_PATH + "/rotatekeys",
+            json={"items": [item.dump(camel_case=True) for item in items]},
+            semaphore=self._get_semaphore("write"),
+        )
+        rotated = ExternalDataSourceRotatedKeysList._load(res.json()["items"])
+        return rotated[0] if isinstance(rotation, ExternalDataSourceRotateKeys) else rotated
+
     async def verify_usability(self, external_id: str) -> ExternalDataSourceUsability:
         """`Verify external data source usability <https://api-docs.cognite.com/20230101-beta/tag/Transformation-External-Data-Sources/operation/verifyExternalDataSourceUsability>`_.
 

cognite/client/data_classes/transformations/__init__.py

@@ -30,6 +30,7 @@
     ExternalDataSource,
     ExternalDataSourceList,
     ExternalDataSourceRotatedKeys,
+    ExternalDataSourceRotatedKeysList,
     ExternalDataSourceRotateKeys,
     ExternalDataSourceUsability,
     ExternalDataSourceWrite,

cognite/client/data_classes/transformations/externaldata.py

@@ -726,6 +726,12 @@ def _load(cls, resource: dict[str, Any]) -> Self:
         )
 
 
+class ExternalDataSourceRotatedKeysList(CogniteResourceList[ExternalDataSourceRotatedKeys], ExternalIDTransformerMixin):
+    """A list of ExternalDataSourceRotatedKeys objects."""
+
+    _RESOURCE = ExternalDataSourceRotatedKeys
+
+
 _EXTERNAL_DATA_SOURCE_WRITE_CLASS_BY_FORMAT: dict[str, type[ExternalDataSourceWrite]] = {
     subclass._format: subclass for subclass in ExternalDataSourceWrite.__subclasses__()
 }

docs/source/transformations_external_data_sources.rst

@@ -56,6 +56,7 @@ Operation                                 Action
 ``retrieve()``                            ``READ``
 ``create()``                              ``WRITE``
 ``delete()``                              ``WRITE``
+``rotate_keys()``                         ``WRITE``
 ``verify_usability()``                    ``USE``
 ========================================= ==========
 
@@ -260,12 +261,39 @@ To iterate without holding every data source in memory, call the API directly, o
     for chunk in client.transformations.external_data_sources(chunk_size=25):
         ...  # do something with the chunk
 
+Rotate Snowflake keys
+---------------------
+
+Snowflake data sources have a dedicated endpoint that generates a new key pair on the server. This operation
+is Snowflake only: the API returns a 400 if you pass the external ID of a OneLake data source (for OneLake, see
+`Change OneLake credentials`_). Pass up to
+10 data sources per call; the operation is atomic, so either all keys are rotated or none. Each
+``expiry_time`` must be in the future and at most one year from now.
+
+The rotation is complete only after you register the new public key with the Snowflake user:
+
+1. Call ``rotate_keys()`` and read the returned ``public_key``.
+2. Run ``ALTER USER ... SET RSA_PUBLIC_KEY`` on the Snowflake user with that key.
+
+.. code-block:: python
+
+    from cognite.client.data_classes.transformations.externaldata import ExternalDataSourceRotateKeys
+    from cognite.client.utils._time import timestamp_to_ms
+
+    rotated = client.transformations.external_data_sources.rotate_keys(
+        ExternalDataSourceRotateKeys("snowflake-analytics-prod", timestamp_to_ms("30d-ahead"))
+    )
+    public_key = rotated.public_key
+
+    # Final step, outside the SDK: ALTER USER COGNITE_SVC SET RSA_PUBLIC_KEY = '<public_key>'
+
 Change OneLake credentials
 --------------------------
 
 OneLake credentials are managed by you, not by Cognite. The Entra ID application and its client secret live
 in your Azure tenant, and the Cognite REST APIs do not generate or rotate them. There is no update or rotate
-endpoint for OneLake data sources.
+endpoint for OneLake data sources. Snowflake data sources do not use this flow; rotate their keys with
+`Rotate Snowflake keys`_ instead.
 
 To change the credentials, or any other setting, delete the data source and create a replacement with the
 same external ID:

tests/tests_unit/test_api/test_transformations_external_data.py

@@ -9,6 +9,9 @@
 from cognite.client import AsyncCogniteClient, CogniteClient
 from cognite.client.data_classes.transformations.externaldata import (
     ExternalDataSourceList,
+    ExternalDataSourceRotatedKeys,
+    ExternalDataSourceRotatedKeysList,
+    ExternalDataSourceRotateKeys,
     OneLakeCredentialsWrite,
     OneLakeExternalDataSource,
     OneLakeExternalDataSourceWrite,
@@ -459,3 +462,64 @@ def test_create_snowflake(
         item = jsgz_load(httpx2_mock.get_requests()[-1].content)["items"][0]
         assert item["format"] == "snowflake"
         assert item["expiryTime"] == 1700000000000
+
+    def test_rotate_keys_single(
+        self, cognite_client: CogniteClient, httpx2_mock: HTTPXMock, external_data_url: str
+    ) -> None:
+        httpx2_mock.add_response(
+            method="POST",
+            url=external_data_url + "/rotatekeys",
+            status_code=200,
+            json={"items": [{"externalId": "src", "publicKey": "pk", "expiryTime": 123}]},
+        )
+
+        result = cognite_client.transformations.external_data_sources.rotate_keys(
+            ExternalDataSourceRotateKeys("src", 123)
+        )
+
+        assert isinstance(result, ExternalDataSourceRotatedKeys)
+        assert result.public_key == "pk"
+        assert jsgz_load(httpx2_mock.get_requests()[-1].content) == {
+            "items": [{"externalId": "src", "expiryTime": 123}]
+        }
+
+    def test_rotate_keys_multiple(
+        self, cognite_client: CogniteClient, httpx2_mock: HTTPXMock, external_data_url: str
+    ) -> None:
+        httpx2_mock.add_response(
+            method="POST",
+            url=external_data_url + "/rotatekeys",
+            status_code=200,
+            json={
+                "items": [
+                    {"externalId": "a", "publicKey": "pk-a", "expiryTime": 1},
+                    {"externalId": "b", "publicKey": "pk-b", "expiryTime": 2},
+                ]
+            },
+        )
+
+        result = cognite_client.transformations.external_data_sources.rotate_keys(
+            [ExternalDataSourceRotateKeys("a", 1), ExternalDataSourceRotateKeys("b", 2)]
+        )
+
+        assert isinstance(result, ExternalDataSourceRotatedKeysList)
+        assert result.as_external_ids() == ["a", "b"]
+
+    def test_rotate_keys_over_limit_raises_without_request(
+        self, cognite_client: CogniteClient, httpx2_mock: HTTPXMock
+    ) -> None:
+        too_many = [ExternalDataSourceRotateKeys(f"src-{no}", 1) for no in range(11)]
+
+        with pytest.raises(ValueError, match="At most 10"):
+            cognite_client.transformations.external_data_sources.rotate_keys(too_many)
+
+        assert httpx2_mock.get_requests() == []
+
+    def test_rotate_keys_no_input(
+        self, cognite_client: CogniteClient, httpx2_mock: HTTPXMock, external_data_url: str
+    ) -> None:
+        result = cognite_client.transformations.external_data_sources.rotate_keys([])
+
+        assert isinstance(result, ExternalDataSourceRotatedKeysList)
+        assert len(result) == 0
+        assert httpx2_mock.get_requests() == []