diff --git a/sdk/keyvault/azure-keyvault-administration/CHANGELOG.md b/sdk/keyvault/azure-keyvault-administration/CHANGELOG.md index 3b902caa0f26..ec1385d2867c 100644 --- a/sdk/keyvault/azure-keyvault-administration/CHANGELOG.md +++ b/sdk/keyvault/azure-keyvault-administration/CHANGELOG.md @@ -6,10 +6,15 @@ ### Breaking Changes -### Bugs Fixed +- Changed the continuation token format. Continuation tokens generated by previous versions of + `azure-keyvault-administration` are not compatible with this version. Similarly, continuation tokens generated by + previous versions of this library are not compatible with versions of `azure-core>=1.38.0`. +### Bugs Fixed ### Other Changes +- Updated minimum `azure-core` version to 1.38.0 + ## 4.6.0 (2025-06-16) ### Features Added diff --git a/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_backup_client.py b/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_backup_client.py index 714907cdf182..bdf2c71c4fcc 100644 --- a/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_backup_client.py +++ b/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_backup_client.py @@ -4,12 +4,13 @@ # ------------------------------------ import base64 import functools -import pickle +import json from typing import Any, Callable, Optional, overload from urllib.parse import urlparse from typing_extensions import Literal +from azure.core.pipeline import PipelineResponse from azure.core.polling import LROPoller from azure.core.tracing.decorator import distributed_trace @@ -19,12 +20,61 @@ from ._internal.polling import KeyVaultBackupClientPolling, KeyVaultBackupClientPollingMethod -def _parse_status_url(url): +def _parse_status_url(url: str) -> str: parsed = urlparse(url) job_id = parsed.path.split("/")[2] return job_id +def _get_continuation_token(pipeline_response: PipelineResponse) -> str: + """Returns an opaque token which can be used by the user to rehydrate/restart the LRO. + + Saves the state of the LRO based on a status response so that polling can be resumed from that context. Because + the service has different operations for backup/restore starting vs. status checking, the caller is expected to + first use the status URL from the initial response to make a status request and then pass that status response to + this function to be serialized into a continuation token. + + :param pipeline_response: The pipeline response of the operation status request. + :type pipeline_response: ~azure.core.pipeline.PipelineResponse + :returns: An opaque continuation token that can be provided to Core to rehydrate the LRO. + :rtype: str + """ + # Headers needed for LRO rehydration - use an allowlist approach for security + lro_headers = {"azure-asyncoperation", "operation-location", "location", "content-type", "retry-after"} + response = pipeline_response.http_response + filtered_headers = {k: v for k, v in response.headers.items() if k.lower() in lro_headers} + + request = response.request + # Serialize the essential parts of the PipelineResponse to JSON. + if request: + request_headers = {} + # Preserve x-ms-client-request-id for request correlation + if "x-ms-client-request-id" in request.headers: + request_headers["x-ms-client-request-id"] = request.headers["x-ms-client-request-id"] + request_state = { + "method": request.method, + "url": request.url, + "headers": request_headers, + } + else: + request_state = None + + # Use a versioned token schema: {"version": , "data": } + # This allows for future compatibility checking when deserializing + token = { + "version": 1, + "data": { + "request": request_state, + "response": { + "status_code": response.status_code, + "headers": filtered_headers, + "content": base64.b64encode(response.content).decode("ascii"), + }, + }, + } + return base64.b64encode(json.dumps(token).encode("utf-8")).decode("ascii") + + class KeyVaultBackupClient(KeyVaultClientBase): """Performs Key Vault backup and restore operations. @@ -56,7 +106,7 @@ def _use_continuation_token(self, continuation_token: str, status_method: Callab ) if "azure-asyncoperation" not in pipeline_response.http_response.headers: pipeline_response.http_response.headers["azure-asyncoperation"] = status_url - return base64.b64encode(pickle.dumps(pipeline_response)).decode("ascii") + return _get_continuation_token(pipeline_response) @overload def begin_backup( diff --git a/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/polling.py b/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/polling.py index eba971b80c47..955c6a9ea6a6 100644 --- a/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/polling.py +++ b/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/polling.py @@ -15,7 +15,7 @@ class KeyVaultBackupClientPolling(OperationResourcePolling): def __init__(self) -> None: self._polling_url = "" - super(KeyVaultBackupClientPolling, self).__init__(operation_location_header="azure-asyncoperation") + super().__init__(operation_location_header="azure-asyncoperation") def get_polling_url(self) -> str: return self._polling_url @@ -34,4 +34,13 @@ def set_initial_status(self, pipeline_response: "PipelineResponse") -> str: class KeyVaultBackupClientPollingMethod(LROBasePolling): def get_continuation_token(self) -> str: + """ + Get a continuation token to resume the polling later. + + :return: A continuation token. + :rtype: str + """ + # Because of the operation structure, we need to use a "continuation token" that is just the status URL. + # This URL can then be used to fetch the status of the operation when resuming, at which point a genuine + # continuation token will be created from the response and provided to Core. return base64.b64encode(self._operation.get_polling_url().encode()).decode("ascii") diff --git a/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/aio/_backup_client.py b/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/aio/_backup_client.py index cc27d245b00e..f3dc496572bb 100644 --- a/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/aio/_backup_client.py +++ b/sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/aio/_backup_client.py @@ -4,7 +4,6 @@ # ------------------------------------ import base64 import functools -import pickle from typing import Any, Callable, Optional, overload from typing_extensions import Literal @@ -13,7 +12,7 @@ from azure.core.tracing.decorator_async import distributed_trace_async from .._generated.models import PreBackupOperationParameters, PreRestoreOperationParameters, SASTokenParameter -from .._backup_client import _parse_status_url +from .._backup_client import _get_continuation_token, _parse_status_url from .._internal import AsyncKeyVaultClientBase, parse_folder_url from .._internal.async_polling import KeyVaultAsyncBackupClientPollingMethod from .._internal.polling import KeyVaultBackupClientPolling @@ -51,7 +50,7 @@ async def _use_continuation_token(self, continuation_token: str, status_method: ) if "azure-asyncoperation" not in pipeline_response.http_response.headers: pipeline_response.http_response.headers["azure-asyncoperation"] = status_url - return base64.b64encode(pickle.dumps(pipeline_response)).decode("ascii") + return _get_continuation_token(pipeline_response) @overload async def begin_backup( diff --git a/sdk/keyvault/azure-keyvault-administration/setup.py b/sdk/keyvault/azure-keyvault-administration/setup.py index 186dd18143f0..0edb70640803 100644 --- a/sdk/keyvault/azure-keyvault-administration/setup.py +++ b/sdk/keyvault/azure-keyvault-administration/setup.py @@ -62,7 +62,7 @@ ), install_requires=[ "isodate>=0.6.1", - "azure-core>=1.31.0", + "azure-core>=1.38.0", "typing-extensions>=4.6.0", ], python_requires=">=3.9",