[Key Vault] Custom polling method for admin backup client (Azure#19204)

Author: mccoyp

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_backup_client.py

@@ -2,21 +2,30 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+import base64
 import functools
+import pickle
 from typing import TYPE_CHECKING
 
-from azure.core.polling.base_polling import LROBasePolling
+from six import raise_from
+from six.moves.urllib_parse import urlparse
 
 from ._models import KeyVaultBackupOperation
 from ._internal import KeyVaultClientBase, parse_folder_url
-from ._internal.polling import KeyVaultBackupClientPolling
+from ._internal.polling import KeyVaultBackupClientPolling, KeyVaultBackupClientPollingMethod
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import
     from typing import Any
     from azure.core.polling import LROPoller
 
 
+def _parse_status_url(url):
+    parsed = urlparse(url)
+    job_id = parsed.path.split("/")[2]
+    return job_id
+
+
 class KeyVaultBackupClient(KeyVaultClientBase):
     """Performs Key Vault backup and restore operations.
 
@@ -37,15 +46,48 @@ def begin_backup(self, blob_storage_url, sas_token, **kwargs):
         :returns: An :class:`~azure.core.polling.LROPoller` instance. Call `result()` on this object to wait for the
             operation to complete and get a :class:`KeyVaultBackupOperation`.
         :rtype: ~azure.core.polling.LROPoller[~azure.keyvault.administration.KeyVaultBackupOperation]
+
+        Example:
+            .. literalinclude:: ../tests/test_examples_administration.py
+                :start-after: [START begin_backup]
+                :end-before: [END begin_backup]
+                :language: python
+                :caption: Create a vault backup
+                :dedent: 8
         """
         polling_interval = kwargs.pop("_polling_interval", 5)
         sas_parameter = self._models.SASTokenParameter(storage_resource_uri=blob_storage_url, token=sas_token)
+
+        continuation_token = kwargs.pop("continuation_token", None)
+        status_response = None
+        if continuation_token:
+            status_url = base64.b64decode(continuation_token.encode()).decode("ascii")
+            try:
+                job_id = _parse_status_url(status_url)
+            except Exception as ex:  # pylint: disable=broad-except
+                raise_from(
+                    ValueError(
+                        "The provided continuation_token is malformed. A valid token can be obtained from the "
+                        + "operation poller's continuation_token() method"
+                    ),
+                    ex,
+                )
+
+            pipeline_response = self._client.full_backup_status(
+                vault_base_url=self._vault_url, job_id=job_id, cls=lambda pipeline_response, _, __: pipeline_response
+            )
+            if "azure-asyncoperation" not in pipeline_response.http_response.headers:
+                pipeline_response.http_response.headers["azure-asyncoperation"] = status_url
+            status_response = base64.b64encode(pickle.dumps(pipeline_response)).decode("ascii")
+
         return self._client.begin_full_backup(
             vault_base_url=self._vault_url,
             azure_storage_blob_container_uri=sas_parameter,
             cls=KeyVaultBackupOperation._from_generated,
-            continuation_token=kwargs.pop("continuation_token", None),
-            polling=LROBasePolling(lro_algorithms=[KeyVaultBackupClientPolling()], timeout=polling_interval, **kwargs),
+            continuation_token=status_response,
+            polling=KeyVaultBackupClientPollingMethod(
+                lro_algorithms=[KeyVaultBackupClientPolling()], timeout=polling_interval, **kwargs
+            ),
             **kwargs
         )
 
@@ -62,14 +104,50 @@ def begin_restore(self, folder_url, sas_token, **kwargs):
         :keyword str continuation_token: a continuation token to restart polling from a saved state
         :keyword str key_name: name of a single key in the backup. When set, only this key will be restored.
         :rtype: ~azure.core.polling.LROPoller
+
+        Examples:
+            .. literalinclude:: ../tests/test_examples_administration.py
+                :start-after: [START begin_restore]
+                :end-before: [END begin_restore]
+                :language: python
+                :caption: Restore a vault backup
+                :dedent: 8
+
+            .. literalinclude:: ../tests/test_examples_administration.py
+                :start-after: [START begin_selective_restore]
+                :end-before: [END begin_selective_restore]
+                :language: python
+                :caption: Restore a single key
+                :dedent: 8
         """
         # LROBasePolling passes its kwargs to pipeline.run(), so we remove unexpected args before constructing it
         continuation_token = kwargs.pop("continuation_token", None)
         key_name = kwargs.pop("key_name", None)
 
+        status_response = None
+        if continuation_token:
+            status_url = base64.b64decode(continuation_token.encode()).decode("ascii")
+            try:
+                job_id = _parse_status_url(status_url)
+            except Exception as ex:  # pylint: disable=broad-except
+                raise_from(
+                    ValueError(
+                        "The provided continuation_token is malformed. A valid token can be obtained from the "
+                        + "operation poller's continuation_token() method"
+                    ),
+                    ex,
+                )
+
+            pipeline_response = self._client.restore_status(
+                vault_base_url=self._vault_url, job_id=job_id, cls=lambda pipeline_response, _, __: pipeline_response
+            )
+            if "azure-asyncoperation" not in pipeline_response.http_response.headers:
+                pipeline_response.http_response.headers["azure-asyncoperation"] = status_url
+            status_response = base64.b64encode(pickle.dumps(pipeline_response)).decode("ascii")
+
         container_url, folder_name = parse_folder_url(folder_url)
         sas_parameter = self._models.SASTokenParameter(storage_resource_uri=container_url, token=sas_token)
-        polling = LROBasePolling(
+        polling = KeyVaultBackupClientPollingMethod(
             lro_algorithms=[KeyVaultBackupClientPolling()], timeout=kwargs.pop("_polling_interval", 5), **kwargs
         )
 
@@ -88,7 +166,7 @@ def begin_restore(self, folder_url, sas_token, **kwargs):
             vault_base_url=self._vault_url,
             restore_blob_details=restore_details,
             cls=lambda *_: None,  # poller.result() returns None
-            continuation_token=continuation_token,
+            continuation_token=status_response,
             polling=polling,
             **kwargs
         )

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/async_polling.py

@@ -0,0 +1,12 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import base64
+
+from azure.core.polling.async_base_polling import AsyncLROBasePolling
+
+
+class KeyVaultAsyncBackupClientPollingMethod(AsyncLROBasePolling):
+    def get_continuation_token(self) -> str:
+        return base64.b64encode(self._operation.get_polling_url().encode()).decode("ascii")

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/polling.py

@@ -2,12 +2,32 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-from azure.core.polling.base_polling import OperationResourcePolling
+import base64
+
+from azure.core.polling.base_polling import LROBasePolling, OperationFailed, OperationResourcePolling
 
 
 class KeyVaultBackupClientPolling(OperationResourcePolling):
     def __init__(self):
+        self._polling_url = None
         super(KeyVaultBackupClientPolling, self).__init__(operation_location_header="azure-asyncoperation")
 
+    def get_polling_url(self):
+        return self._polling_url
+
     def get_final_get_url(self, pipeline_response):
         return None
+
+    def set_initial_status(self, pipeline_response):
+        response = pipeline_response.http_response
+        self._polling_url = response.headers["azure-asyncoperation"]
+
+        if response.status_code in {200, 201, 202, 204}:
+            return self.get_status(pipeline_response)
+        raise OperationFailed("Operation failed or canceled")
+
+
+class KeyVaultBackupClientPollingMethod(LROBasePolling):
+    def get_continuation_token(self):
+        # type: () -> str
+        return base64.b64encode(self._operation.get_polling_url().encode()).decode("ascii")

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/aio/_backup_client.py

@@ -2,12 +2,14 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+import base64
 import functools
+import pickle
 from typing import TYPE_CHECKING
 
-from azure.core.polling.async_base_polling import AsyncLROBasePolling
-
+from .._backup_client import _parse_status_url
 from .._internal import AsyncKeyVaultClientBase, parse_folder_url
+from .._internal.async_polling import KeyVaultAsyncBackupClientPollingMethod
 from .._internal.polling import KeyVaultBackupClientPolling
 from .._models import KeyVaultBackupOperation
 
@@ -37,15 +39,43 @@ async def begin_backup(
         :keyword str continuation_token: a continuation token to restart polling from a saved state
         :returns: An AsyncLROPoller. Call `result()` on this object to get a :class:`KeyVaultBackupOperation`.
         :rtype: ~azure.core.polling.AsyncLROPoller[~azure.keyvault.administration.KeyVaultBackupOperation]
+
+        Example:
+            .. literalinclude:: ../tests/test_examples_administration_async.py
+                :start-after: [START begin_backup]
+                :end-before: [END begin_backup]
+                :language: python
+                :caption: Create a vault backup
+                :dedent: 8
         """
         polling_interval = kwargs.pop("_polling_interval", 5)
         sas_parameter = self._models.SASTokenParameter(storage_resource_uri=blob_storage_url, token=sas_token)
+
+        continuation_token = kwargs.pop("continuation_token", None)
+        status_response = None
+        if continuation_token:
+            status_url = base64.b64decode(continuation_token.encode()).decode("ascii")
+            try:
+                job_id = _parse_status_url(status_url)
+            except Exception as ex:  # pylint: disable=broad-except
+                raise ValueError(
+                    "The provided continuation_token is malformed. A valid token can be obtained from the operation "
+                    + "poller's continuation_token() method"
+                ) from ex
+
+            pipeline_response = await self._client.full_backup_status(
+                vault_base_url=self._vault_url, job_id=job_id, cls=lambda pipeline_response, _, __: pipeline_response
+            )
+            if "azure-asyncoperation" not in pipeline_response.http_response.headers:
+                pipeline_response.http_response.headers["azure-asyncoperation"] = status_url
+            status_response = base64.b64encode(pickle.dumps(pipeline_response)).decode("ascii")
+
         return await self._client.begin_full_backup(
             vault_base_url=self._vault_url,
             azure_storage_blob_container_uri=sas_parameter,
             cls=KeyVaultBackupOperation._from_generated,
-            continuation_token=kwargs.pop("continuation_token", None),
-            polling=AsyncLROBasePolling(
+            continuation_token=status_response,
+            polling=KeyVaultAsyncBackupClientPollingMethod(
                 lro_algorithms=[KeyVaultBackupClientPolling()], timeout=polling_interval, **kwargs
             ),
             **kwargs
@@ -64,14 +94,47 @@ async def begin_restore(self, folder_url: str, sas_token: str, **kwargs: "Any")
         :keyword str continuation_token: a continuation token to restart polling from a saved state
         :keyword str key_name: name of a single key in the backup. When set, only this key will be restored.
         :rtype: ~azure.core.polling.AsyncLROPoller
+
+        Examples:
+            .. literalinclude:: ../tests/test_examples_administration_async.py
+                :start-after: [START begin_restore]
+                :end-before: [END begin_restore]
+                :language: python
+                :caption: Restore a vault backup
+                :dedent: 8
+
+            .. literalinclude:: ../tests/test_examples_administration_async.py
+                :start-after: [START begin_selective_restore]
+                :end-before: [END begin_selective_restore]
+                :language: python
+                :caption: Restore a single key
+                :dedent: 8
         """
         # AsyncLROBasePolling passes its kwargs to pipeline.run(), so we remove unexpected args before constructing it
         continuation_token = kwargs.pop("continuation_token", None)
         key_name = kwargs.pop("key_name", None)
 
+        status_response = None
+        if continuation_token:
+            status_url = base64.b64decode(continuation_token.encode()).decode("ascii")
+            try:
+                job_id = _parse_status_url(status_url)
+            except Exception as ex:  # pylint: disable=broad-except
+                raise ValueError(
+                    "The provided continuation_token is malformed. A valid token can be obtained from the operation "
+                    + "poller's continuation_token() method"
+                ) from ex
+
+            pipeline_response = await self._client.restore_status(
+                vault_base_url=self._vault_url, job_id=job_id, cls=lambda pipeline_response, _, __: pipeline_response
+            )
+            if "azure-asyncoperation" not in pipeline_response.http_response.headers:
+                pipeline_response.http_response.headers["azure-asyncoperation"] = status_url
+            status_response = base64.b64encode(pickle.dumps(pipeline_response)).decode("ascii")
+
         container_url, folder_name = parse_folder_url(folder_url)
         sas_parameter = self._models.SASTokenParameter(storage_resource_uri=container_url, token=sas_token)
-        polling = AsyncLROBasePolling(
+        polling = KeyVaultAsyncBackupClientPollingMethod(
             lro_algorithms=[KeyVaultBackupClientPolling()], timeout=kwargs.pop("_polling_interval", 5), **kwargs
         )
 
@@ -90,7 +153,7 @@ async def begin_restore(self, folder_url: str, sas_token: str, **kwargs: "Any")
             vault_base_url=self._vault_url,
             restore_blob_details=restore_details,
             cls=lambda *_: None,  # poller.result() returns None
-            continuation_token=continuation_token,
+            continuation_token=status_response,
             polling=polling,
             **kwargs
         )

sdk/keyvault/azure-keyvault-administration/setup.py

@@ -80,7 +80,7 @@
             "azure.keyvault",
         ]
     ),
-    install_requires=["azure-common~=1.1", "azure-core<2.0.0,>=1.11.0", "msrest>=0.6.21"],
+    install_requires=["azure-common~=1.1", "azure-core<2.0.0,>=1.11.0", "msrest>=0.6.21", "six>=1.11.0"],
     extras_require={
         ":python_version<'3.0'": ["azure-keyvault-nspkg"],
         ":python_version<'3.4'": ["enum34>=1.0.4"],