Repository navigation
feat(transformations): Added rotate keys api endpoint support in sdk - #2921
Conversation
karthikeyanraghavan
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new rotate_keys endpoint to rotate Snowflake external data source keys, adding the corresponding async/sync API methods, data classes, documentation, and unit tests. The review feedback suggests enhancing rotate_keys to support dictionary inputs in addition to typed objects, ensuring consistency with other SDK methods and preventing runtime errors.
| 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)}" | ||
| ) | ||
| 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]}, | ||
| semaphore=self._get_semaphore("write"), | ||
| ) | ||
| rotated = ExternalDataSourceRotatedKeysList._load(res.json()["items"]) | ||
| return rotated[0] if isinstance(rotation, ExternalDataSourceRotateKeys) else rotated |
There was a problem hiding this comment.
The current implementation of rotate_keys does not support passing a dictionary or a list of dictionaries as input, which is inconsistent with other SDK methods (e.g., create, retrieve). Furthermore, if a dictionary is passed, isinstance(rotation, ExternalDataSourceRotateKeys) evaluates to False, and list(rotation) extracts the dictionary keys (['externalId', 'expiryTime']), leading to a cryptic AttributeError: 'str' object has no attribute 'dump' at runtime.
To align with the SDK's established patterns and improve robustness, we should support both ExternalDataSourceRotateKeys and dict inputs, loading them properly using ExternalDataSourceRotateKeys.load.
is_single = isinstance(rotation, (ExternalDataSourceRotateKeys, dict))
items = [rotation] if is_single 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)}"
)
processed_items: list[ExternalDataSourceRotateKeys] = []
for item in items:
if isinstance(item, dict):
processed_items.append(ExternalDataSourceRotateKeys.load(item))
elif isinstance(item, ExternalDataSourceRotateKeys):
processed_items.append(item)
else:
raise TypeError(f"Expected ExternalDataSourceRotateKeys or dict, got {type(item)}")
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 processed_items]},
semaphore=self._get_semaphore("write"),
)
rotated = ExternalDataSourceRotatedKeysList._load(res.json()["items"])
return rotated[0] if is_single else rotatedThere was a problem hiding this comment.
This is not true. Create and retrieve doesn't support dict or list of dicts. Will leave for reviewer in case they think otherwise
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #2921 +/- ##
==========================================
- Coverage 93.69% 93.68% -0.02%
==========================================
Files 540 540
Lines 57203 57248 +45
==========================================
+ Hits 53598 53633 +35
- Misses 3605 3615 +10
🚀 New features to boost your workflow:
|
HMEiding
left a comment
HMEiding
left a comment
There was a problem hiding this comment.
Looks good to me, just one nit.
| >>> 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 |
There was a problem hiding this comment.
I have commented on this a number of times now, I would prefer if you used the “days ahead” syntax
There was a problem hiding this comment.
Oh my bad. I will fix this.
| 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) |
karthikeyanraghavan
added
the
waiting-for-risk-review
|
🦄 |
haakonvt
added
risk-review-ongoing
Description
Checklist:
If a new method has been added it should be referenced in cognite.rst in order to generate docs based on its docstring.