From db32b6a247eb9093c225d6f418a1f65db4845c99 Mon Sep 17 00:00:00 2001 From: Karthikeyan Raghavan Date: Wed, 4 Nov 2026 15:18:55 +0530 Subject: [PATCH 1/6] Added rotate keys api endpoint support in sdk --- .../_api/transformations/externaldata.py | 69 +++++++++++++++++++ .../_sync_api/transformations/externaldata.py | 47 +++++++++++++ .../data_classes/transformations/__init__.py | 1 + .../transformations/externaldata.py | 6 ++ .../transformations_external_data_sources.rst | 30 +++++++- .../test_transformations_external_data.py | 55 +++++++++++++++ 6 files changed, 207 insertions(+), 1 deletion(-) diff --git a/cognite/client/_api/transformations/externaldata.py b/cognite/client/_api/transformations/externaldata.py index 82020c5c36..fffa665f00 100644 --- a/cognite/client/_api/transformations/externaldata.py +++ b/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 `_. + + 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("now") + 30 * 24 * 3600 * 1000 + >>> res = client.transformations.external_data_sources.rotate_keys( + ... ExternalDataSourceRotateKeys("snowflake-analytics-prod", expiry_time) + ... ) + >>> new_public_key = res.public_key + """ + self._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 `_. diff --git a/cognite/client/_sync_api/transformations/externaldata.py b/cognite/client/_sync_api/transformations/externaldata.py index 65e8a14db7..b501e3cedf 100644 --- a/cognite/client/_sync_api/transformations/externaldata.py +++ b/cognite/client/_sync_api/transformations/externaldata.py @@ -16,6 +16,9 @@ ExternalDataSource, ExternalDataSourceFormat, ExternalDataSourceList, + ExternalDataSourceRotatedKeys, + ExternalDataSourceRotatedKeysList, + ExternalDataSourceRotateKeys, ExternalDataSourceUsability, ExternalDataSourceWrite, ) @@ -257,6 +260,50 @@ def delete(self, external_id: str | SequenceNotStr[str]) -> None: """ return run_sync(self.__async_client.transformations.external_data_sources.delete(external_id=external_id)) + @overload + def rotate_keys(self, rotation: ExternalDataSourceRotateKeys) -> ExternalDataSourceRotatedKeys: ... + + @overload + def rotate_keys(self, rotation: Sequence[ExternalDataSourceRotateKeys]) -> ExternalDataSourceRotatedKeysList: ... + + def rotate_keys( + self, rotation: ExternalDataSourceRotateKeys | Sequence[ExternalDataSourceRotateKeys] + ) -> ExternalDataSourceRotatedKeys | ExternalDataSourceRotatedKeysList: + """ + `Rotate the keys of Snowflake external data sources `_. + + 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("now") + 30 * 24 * 3600 * 1000 + >>> res = client.transformations.external_data_sources.rotate_keys( + ... ExternalDataSourceRotateKeys("snowflake-analytics-prod", expiry_time) + ... ) + >>> new_public_key = res.public_key + """ + return run_sync(self.__async_client.transformations.external_data_sources.rotate_keys(rotation=rotation)) + def verify_usability(self, external_id: str) -> ExternalDataSourceUsability: """ `Verify external data source usability `_. diff --git a/cognite/client/data_classes/transformations/__init__.py b/cognite/client/data_classes/transformations/__init__.py index 35760c49a2..5ed93af6e4 100644 --- a/cognite/client/data_classes/transformations/__init__.py +++ b/cognite/client/data_classes/transformations/__init__.py @@ -30,6 +30,7 @@ ExternalDataSource, ExternalDataSourceList, ExternalDataSourceRotatedKeys, + ExternalDataSourceRotatedKeysList, ExternalDataSourceRotateKeys, ExternalDataSourceUsability, ExternalDataSourceWrite, diff --git a/cognite/client/data_classes/transformations/externaldata.py b/cognite/client/data_classes/transformations/externaldata.py index a0af3cf24d..e84884fa3d 100644 --- a/cognite/client/data_classes/transformations/externaldata.py +++ b/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__() } diff --git a/docs/source/transformations_external_data_sources.rst b/docs/source/transformations_external_data_sources.rst index 4c1e75acbc..4917d04899 100644 --- a/docs/source/transformations_external_data_sources.rst +++ b/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("now") + 30 * 24 * 3600 * 1000) + ) + public_key = rotated.public_key + + # Final step, outside the SDK: ALTER USER COGNITE_SVC SET RSA_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: diff --git a/tests/tests_unit/test_api/test_transformations_external_data.py b/tests/tests_unit/test_api/test_transformations_external_data.py index 00cc0cad27..92cc45ec34 100644 --- a/tests/tests_unit/test_api/test_transformations_external_data.py +++ b/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,55 @@ 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() == [] From 84dc9a5dac07b22e61ccabc974a32ad56e1a47e7 Mon Sep 17 00:00:00 2001 From: Karthikeyan Raghavan Date: Wed, 4 Nov 2026 15:27:12 +0530 Subject: [PATCH 2/6] Added alpha tag to rotate keys endpoint --- cognite/client/_api/transformations/externaldata.py | 1 + 1 file changed, 1 insertion(+) diff --git a/cognite/client/_api/transformations/externaldata.py b/cognite/client/_api/transformations/externaldata.py index fffa665f00..c55b169dad 100644 --- a/cognite/client/_api/transformations/externaldata.py +++ b/cognite/client/_api/transformations/externaldata.py @@ -347,6 +347,7 @@ async def rotate_keys( 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)}" ) + self._rotate_keys_warning.warn() res = await self._post( url_path=self._RESOURCE_PATH + "/rotatekeys", json={"items": [item.dump(camel_case=True) for item in items]}, From dec9e149ed266be01625c7770a2fa356ba8af300 Mon Sep 17 00:00:00 2001 From: Karthikeyan Raghavan Date: Wed, 4 Nov 2026 15:34:40 +0530 Subject: [PATCH 3/6] Added alpha tag to rotate keys endpoint --- cognite/client/_api/transformations/externaldata.py | 5 ++--- cognite/client/_sync_api/transformations/externaldata.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/cognite/client/_api/transformations/externaldata.py b/cognite/client/_api/transformations/externaldata.py index c55b169dad..8621dd540f 100644 --- a/cognite/client/_api/transformations/externaldata.py +++ b/cognite/client/_api/transformations/externaldata.py @@ -306,7 +306,7 @@ async def rotate_keys( async def rotate_keys( self, rotation: ExternalDataSourceRotateKeys | Sequence[ExternalDataSourceRotateKeys] ) -> ExternalDataSourceRotatedKeys | ExternalDataSourceRotatedKeysList: - """`Rotate the keys of Snowflake external data sources `_. + """`Rotate the keys of Snowflake external data sources `_. 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 @@ -338,7 +338,7 @@ async def rotate_keys( ... ) >>> new_public_key = res.public_key """ - self._warning.warn() + self._rotate_keys_warning.warn() items = [rotation] if isinstance(rotation, ExternalDataSourceRotateKeys) else list(rotation) if not items: return ExternalDataSourceRotatedKeysList([]) @@ -347,7 +347,6 @@ async def rotate_keys( 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)}" ) - self._rotate_keys_warning.warn() res = await self._post( url_path=self._RESOURCE_PATH + "/rotatekeys", json={"items": [item.dump(camel_case=True) for item in items]}, diff --git a/cognite/client/_sync_api/transformations/externaldata.py b/cognite/client/_sync_api/transformations/externaldata.py index b501e3cedf..91fd22fdfe 100644 --- a/cognite/client/_sync_api/transformations/externaldata.py +++ b/cognite/client/_sync_api/transformations/externaldata.py @@ -270,7 +270,7 @@ def rotate_keys( self, rotation: ExternalDataSourceRotateKeys | Sequence[ExternalDataSourceRotateKeys] ) -> ExternalDataSourceRotatedKeys | ExternalDataSourceRotatedKeysList: """ - `Rotate the keys of Snowflake external data sources `_. + `Rotate the keys of Snowflake external data sources `_. 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 From f95abdc0f93eed992029418eaaa924a1729393f5 Mon Sep 17 00:00:00 2001 From: Karthikeyan Raghavan Date: Wed, 4 Nov 2026 22:02:40 +0530 Subject: [PATCH 4/6] Added a unit test for rotate keys empty input --- .../test_api/test_transformations_external_data.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/tests_unit/test_api/test_transformations_external_data.py b/tests/tests_unit/test_api/test_transformations_external_data.py index 92cc45ec34..55b8c6a917 100644 --- a/tests/tests_unit/test_api/test_transformations_external_data.py +++ b/tests/tests_unit/test_api/test_transformations_external_data.py @@ -514,3 +514,12 @@ def test_rotate_keys_over_limit_raises_without_request( 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() == [] \ No newline at end of file From f5330deb839f53d95489680e42ea5f5d91b51976 Mon Sep 17 00:00:00 2001 From: Karthikeyan Raghavan Date: Wed, 4 Nov 2026 22:47:22 +0530 Subject: [PATCH 5/6] Reformatted files --- .../tests_unit/test_api/test_transformations_external_data.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/tests_unit/test_api/test_transformations_external_data.py b/tests/tests_unit/test_api/test_transformations_external_data.py index 55b8c6a917..b3a85a3d57 100644 --- a/tests/tests_unit/test_api/test_transformations_external_data.py +++ b/tests/tests_unit/test_api/test_transformations_external_data.py @@ -516,10 +516,10 @@ def test_rotate_keys_over_limit_raises_without_request( assert httpx2_mock.get_requests() == [] def test_rotate_keys_no_input( - self, cognite_client: CogniteClient, httpx2_mock: HTTPXMock, external_data_url: str + 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() == [] \ No newline at end of file + assert httpx2_mock.get_requests() == [] From f933c76d6352738112c248cbc6dcaefec4138928 Mon Sep 17 00:00:00 2001 From: Karthikeyan Raghavan Date: Thu, 5 Nov 2026 08:54:20 +0530 Subject: [PATCH 6/6] Refactored code to use timestamp_to_ms functionality correctly --- cognite/client/_api/transformations/externaldata.py | 2 +- cognite/client/_sync_api/transformations/externaldata.py | 2 +- docs/source/transformations_external_data_sources.rst | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cognite/client/_api/transformations/externaldata.py b/cognite/client/_api/transformations/externaldata.py index 8621dd540f..0b309905a5 100644 --- a/cognite/client/_api/transformations/externaldata.py +++ b/cognite/client/_api/transformations/externaldata.py @@ -332,7 +332,7 @@ async def rotate_keys( >>> from cognite.client.utils._time import timestamp_to_ms >>> client = CogniteClient() >>> # async_client = AsyncCogniteClient() # another option - >>> expiry_time = timestamp_to_ms("now") + 30 * 24 * 3600 * 1000 + >>> expiry_time = timestamp_to_ms("30d-ahead") >>> res = client.transformations.external_data_sources.rotate_keys( ... ExternalDataSourceRotateKeys("snowflake-analytics-prod", expiry_time) ... ) diff --git a/cognite/client/_sync_api/transformations/externaldata.py b/cognite/client/_sync_api/transformations/externaldata.py index 91fd22fdfe..408c9e8d87 100644 --- a/cognite/client/_sync_api/transformations/externaldata.py +++ b/cognite/client/_sync_api/transformations/externaldata.py @@ -296,7 +296,7 @@ def rotate_keys( >>> from cognite.client.utils._time import timestamp_to_ms >>> client = CogniteClient() >>> # async_client = AsyncCogniteClient() # another option - >>> expiry_time = timestamp_to_ms("now") + 30 * 24 * 3600 * 1000 + >>> expiry_time = timestamp_to_ms("30d-ahead") >>> res = client.transformations.external_data_sources.rotate_keys( ... ExternalDataSourceRotateKeys("snowflake-analytics-prod", expiry_time) ... ) diff --git a/docs/source/transformations_external_data_sources.rst b/docs/source/transformations_external_data_sources.rst index 4917d04899..75f4569de8 100644 --- a/docs/source/transformations_external_data_sources.rst +++ b/docs/source/transformations_external_data_sources.rst @@ -281,7 +281,7 @@ The rotation is complete only after you register the new public key with the Sno from cognite.client.utils._time import timestamp_to_ms rotated = client.transformations.external_data_sources.rotate_keys( - ExternalDataSourceRotateKeys("snowflake-analytics-prod", timestamp_to_ms("now") + 30 * 24 * 3600 * 1000) + ExternalDataSourceRotateKeys("snowflake-analytics-prod", timestamp_to_ms("30d-ahead")) ) public_key = rotated.public_key