[Key Vault] Remove direct uses of asyncio.sleep (#33819)

Author: mccoyp

sdk/keyvault/azure-keyvault-certificates/CHANGELOG.md

@@ -10,6 +10,8 @@
 
 ### Other Changes
 - Python 3.7 is no longer supported. Please use Python version 3.8 or later.
+- `asyncio` is no longer directly referenced by the library
+  ([#33819](https://github.com/Azure/azure-sdk-for-python/pull/33819))
 
 ## 4.8.0b3 (2023-11-03)
 

sdk/keyvault/azure-keyvault-certificates/azure/keyvault/certificates/_client.py

@@ -110,8 +110,12 @@ def begin_create_certificate(
             tags=kwargs.pop("tags", None),
         )
 
-        cert_bundle = self._client.create_certificate(
-            vault_base_url=self.vault_url, certificate_name=certificate_name, parameters=parameters, **kwargs
+        pipeline_response, cert_bundle = self._client.create_certificate(
+            vault_base_url=self.vault_url,
+            certificate_name=certificate_name,
+            parameters=parameters,
+            cls=lambda pipeline_response, deserialized, _: (pipeline_response, deserialized),
+            **kwargs,
         )
 
         create_certificate_operation = CertificateOperation._from_certificate_operation_bundle(cert_bundle)
@@ -121,7 +125,9 @@ def begin_create_certificate(
         get_certificate_command = partial(self.get_certificate, certificate_name=certificate_name, **kwargs)
 
         create_certificate_polling = CreateCertificatePoller(
-            get_certificate_command=get_certificate_command, interval=polling_interval
+            pipeline_response=pipeline_response,
+            get_certificate_command=get_certificate_command,
+            interval=polling_interval
         )
 
         def no_op(*_, **__) -> Any:  # The deserialization callback is ignored based on polling implementation
@@ -217,14 +223,18 @@ def begin_delete_certificate(self, certificate_name: str, **kwargs: Any) -> LROP
         polling_interval = kwargs.pop("_polling_interval", None)
         if polling_interval is None:
             polling_interval = 2
-        deleted_cert_bundle = self._client.delete_certificate(
-            vault_base_url=self.vault_url, certificate_name=certificate_name, **kwargs
+        pipeline_response, deleted_cert_bundle = self._client.delete_certificate(
+            vault_base_url=self.vault_url,
+            certificate_name=certificate_name,
+            cls=lambda pipeline_response, deserialized, _: (pipeline_response, deserialized),
+            **kwargs,
         )
         deleted_cert = DeletedCertificate._from_deleted_certificate_bundle(deleted_cert_bundle)
 
         polling_method = DeleteRecoverPollingMethod(
             # no recovery ID means soft-delete is disabled, in which case we initialize the poller as finished
             finished=deleted_cert.recovery_id is None,
+            pipeline_response=pipeline_response,
             command=partial(self.get_deleted_certificate, certificate_name=certificate_name, **kwargs),
             final_resource=deleted_cert,
             interval=polling_interval,
@@ -312,13 +322,20 @@ def begin_recover_deleted_certificate(self, certificate_name: str, **kwargs: Any
         if polling_interval is None:
             polling_interval = 2
 
-        recovered_cert_bundle = self._client.recover_deleted_certificate(
-            vault_base_url=self.vault_url, certificate_name=certificate_name, **kwargs
+        pipeline_response, recovered_cert_bundle = self._client.recover_deleted_certificate(
+            vault_base_url=self.vault_url,
+            certificate_name=certificate_name,
+            cls=lambda pipeline_response, deserialized, _: (pipeline_response, deserialized),
+            **kwargs,
         )
         recovered_certificate = KeyVaultCertificate._from_certificate_bundle(recovered_cert_bundle)
         command = partial(self.get_certificate, certificate_name=certificate_name, **kwargs)
         polling_method = DeleteRecoverPollingMethod(
-            finished=False, command=command, final_resource=recovered_certificate, interval=polling_interval
+            finished=False,
+            pipeline_response=pipeline_response,
+            command=command,
+            final_resource=recovered_certificate,
+            interval=polling_interval
         )
 
         return KeyVaultOperationPoller(polling_method)

sdk/keyvault/azure-keyvault-certificates/azure/keyvault/certificates/_polling.py

@@ -3,17 +3,21 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import logging
-import time
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, cast, Optional, Union
 
+from azure.core.pipeline import PipelineResponse
+from azure.core.pipeline.transport import HttpTransport
 from azure.core.polling import PollingMethod
 from azure.keyvault.certificates._models import KeyVaultCertificate, CertificateOperation
 
 logger = logging.getLogger(__name__)
 
 
 class CreateCertificatePoller(PollingMethod):
-    def __init__(self, get_certificate_command: Callable, interval: int = 5) -> None:
+    def __init__(
+            self, pipeline_response: PipelineResponse, get_certificate_command: Callable, interval: int = 5
+        ) -> None:
+        self._pipeline_response = pipeline_response
         self._command: Optional[Callable] = None
         self._resource: Optional[Union[CertificateOperation, KeyVaultCertificate]] = None
         self._pending_certificate_op: Optional[CertificateOperation] = None
@@ -32,7 +36,9 @@ def run(self) -> None:
             while not self.finished():
                 self._update_status()
                 if not self.finished():
-                    time.sleep(self._polling_interval)
+                    # We should always ask the client's transport to sleep, instead of sleeping directly
+                    transport: HttpTransport = cast(HttpTransport, self._pipeline_response.context.transport)
+                    transport.sleep(self._polling_interval)
             operation = self._pending_certificate_op
             if operation and operation.status and operation.status.lower() == "completed":
                 self._resource = self._get_certificate_command()

sdk/keyvault/azure-keyvault-certificates/azure/keyvault/certificates/_shared/_polling.py

@@ -3,21 +3,18 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import logging
-import time
 import threading
 import uuid
-from typing import TYPE_CHECKING
+from typing import Any, Callable, cast, Optional
 
-from azure.core.polling import PollingMethod, LROPoller, NoPolling
 from azure.core.exceptions import ResourceNotFoundError, HttpResponseError
+from azure.core.pipeline import PipelineResponse
+from azure.core.pipeline.transport import HttpTransport
+from azure.core.polling import PollingMethod, LROPoller, NoPolling
 
 from azure.core.tracing.decorator import distributed_trace
 from azure.core.tracing.common import with_current_context
 
-if TYPE_CHECKING:
-    # pylint: disable=ungrouped-imports
-    from typing import Any, Callable, Optional
-
 logger = logging.getLogger(__name__)
 
 
@@ -45,7 +42,7 @@ def result(self) -> "Any":  # type: ignore
         return self._polling_method.resource()
 
     @distributed_trace
-    def wait(self, timeout: "Optional[float]" = None) -> None:
+    def wait(self, timeout: Optional[float] = None) -> None:
         """Wait on the long running operation for a number of seconds.
 
         You can check if this call has ended with timeout with the "done()" method.
@@ -84,14 +81,24 @@ class DeleteRecoverPollingMethod(PollingMethod):
     Similarly, while recovering a deleted resource, Key Vault will respond 404 to GET requests for the non-deleted
     resource; when it responds 2xx, the resource exists in the non-deleted collection, i.e. its recovery is complete.
 
+    :param pipeline_response: The operation's original pipeline response.
+    :type pipeline_response: PipelineResponse
     :param command: A callable to invoke when polling.
     :type command: Callable
     :param final_resource: The final resource returned by the polling operation.
     :type final_resource: Any
     :param bool finished: Whether or not the polling operation is completed.
     :param int interval: The polling interval, in seconds.
     """
-    def __init__(self, command: "Callable", final_resource: "Any", finished: bool, interval: int = 2) -> None:
+    def __init__(
+            self,
+            pipeline_response: PipelineResponse,
+            command: Callable,
+            final_resource: Any,
+            finished: bool,
+            interval: int = 2
+        ) -> None:
+        self._pipeline_response = pipeline_response
         self._command = command
         self._resource = final_resource
         self._polling_interval = interval
@@ -111,23 +118,25 @@ def _update_status(self) -> None:
             else:
                 raise
 
-    def initialize(self, client: "Any", initial_response: "Any", deserialization_callback: "Callable") -> None:
+    def initialize(self, client: Any, initial_response: Any, deserialization_callback: Callable) -> None:
         pass
 
     def run(self) -> None:
         try:
             while not self.finished():
                 self._update_status()
                 if not self.finished():
-                    time.sleep(self._polling_interval)
+                    # We should always ask the client's transport to sleep, instead of sleeping directly
+                    transport: HttpTransport = cast(HttpTransport, self._pipeline_response.context.transport)
+                    transport.sleep(self._polling_interval)
         except Exception as e:
             logger.warning(str(e))
             raise
 
     def finished(self) -> bool:
         return self._finished
 
-    def resource(self) -> "Any":
+    def resource(self) -> Any:
         return self._resource
 
     def status(self) -> str:

sdk/keyvault/azure-keyvault-certificates/azure/keyvault/certificates/_shared/_polling_async.py

@@ -2,13 +2,13 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-
-import asyncio
 import logging
-from typing import Any, Callable
+from typing import Any, Callable, cast
 
-from azure.core.polling import AsyncPollingMethod
 from azure.core.exceptions import ResourceNotFoundError, HttpResponseError
+from azure.core.pipeline import PipelineResponse
+from azure.core.pipeline.transport import AsyncHttpTransport
+from azure.core.polling import AsyncPollingMethod
 
 logger = logging.getLogger(__name__)
 
@@ -24,6 +24,8 @@ class AsyncDeleteRecoverPollingMethod(AsyncPollingMethod):
     Similarly, while recovering a deleted resource, Key Vault will respond 404 to GET requests for the non-deleted
     resource; when it responds 2xx, the resource exists in the non-deleted collection, i.e. its recovery is complete.
 
+    :param pipeline_response: The operation's original pipeline response.
+    :type pipeline_response: PipelineResponse
     :param command: An awaitable to invoke when polling.
     :type command: Callable
     :param final_resource: The final resource returned by the polling operation.
@@ -33,8 +35,14 @@ class AsyncDeleteRecoverPollingMethod(AsyncPollingMethod):
     """
 
     def __init__(
-            self, command: Callable, final_resource: Any, finished: bool, interval: int = 2
+            self,
+            pipeline_response: PipelineResponse,
+            command: Callable,
+            final_resource: Any,
+            finished: bool,
+            interval: int = 2
         ) -> None:
+        self._pipeline_response = pipeline_response
         self._command = command
         self._resource = final_resource
         self._polling_interval = interval
@@ -62,7 +70,9 @@ async def run(self) -> None:
             while not self.finished():
                 await self._update_status()
                 if not self.finished():
-                    await asyncio.sleep(self._polling_interval)
+                    # We should always ask the client's transport to sleep, instead of sleeping directly
+                    transport: AsyncHttpTransport = cast(AsyncHttpTransport, self._pipeline_response.context.transport)
+                    await transport.sleep(self._polling_interval)
         except Exception as e:
             logger.warning(str(e))
             raise

sdk/keyvault/azure-keyvault-certificates/azure/keyvault/certificates/aio/_client.py

@@ -7,7 +7,7 @@
 from typing import Any, Optional, List, Union
 from functools import partial
 
-from azure.core.polling import async_poller
+from azure.core.polling import AsyncLROPoller
 from azure.core.tracing.decorator import distributed_trace
 from azure.core.tracing.decorator_async import distributed_trace_async
 from azure.core.async_paging import AsyncItemPaged
@@ -106,10 +106,11 @@ async def create_certificate(
             tags=kwargs.pop("tags", None),
         )
 
-        cert_bundle = await self._client.create_certificate(
+        pipeline_response, cert_bundle = await self._client.create_certificate(
             vault_base_url=self.vault_url,
             certificate_name=certificate_name,
             parameters=parameters,
+            cls=lambda pipeline_response, deserialized, _: (pipeline_response, deserialized),
             **kwargs
         )
 
@@ -120,11 +121,13 @@ async def create_certificate(
         get_certificate_command = partial(self.get_certificate, certificate_name=certificate_name, **kwargs)
 
         create_certificate_polling = CreateCertificatePollerAsync(
-            get_certificate_command=get_certificate_command, interval=polling_interval
+            pipeline_response=pipeline_response,
+            get_certificate_command=get_certificate_command,
+            interval=polling_interval,
         )
         def no_op(*_, **__) -> Any:  # The deserialization callback is ignored based on polling implementation
             pass
-        return await async_poller(command, create_certificate_operation, no_op, create_certificate_polling)
+        return await AsyncLROPoller(command, create_certificate_operation, no_op, create_certificate_polling)
 
     @distributed_trace_async
     async def get_certificate(self, certificate_name: str, **kwargs: Any) -> KeyVaultCertificate:
@@ -216,14 +219,18 @@ async def delete_certificate(self, certificate_name: str, **kwargs: Any) -> Dele
         polling_interval = kwargs.pop("_polling_interval", None)
         if polling_interval is None:
             polling_interval = 2
-        deleted_cert_bundle = await self._client.delete_certificate(
-            vault_base_url=self.vault_url, certificate_name=certificate_name, **kwargs
+        pipeline_response, deleted_cert_bundle = await self._client.delete_certificate(
+            vault_base_url=self.vault_url,
+            certificate_name=certificate_name,
+            cls=lambda pipeline_response, deserialized, _: (pipeline_response, deserialized),
+            **kwargs,
         )
         deleted_certificate = DeletedCertificate._from_deleted_certificate_bundle(deleted_cert_bundle)
 
         polling_method = AsyncDeleteRecoverPollingMethod(
             # no recovery ID means soft-delete is disabled, in which case we initialize the poller as finished
             finished=deleted_certificate.recovery_id is None,
+            pipeline_response=pipeline_response,
             command=partial(self.get_deleted_certificate, certificate_name=certificate_name, **kwargs),
             final_resource=deleted_certificate,
             interval=polling_interval,
@@ -307,14 +314,21 @@ async def recover_deleted_certificate(self, certificate_name: str, **kwargs: Any
         polling_interval = kwargs.pop("_polling_interval", None)
         if polling_interval is None:
             polling_interval = 2
-        recovered_cert_bundle = await self._client.recover_deleted_certificate(
-            vault_base_url=self.vault_url, certificate_name=certificate_name, **kwargs
+        pipeline_response, recovered_cert_bundle = await self._client.recover_deleted_certificate(
+            vault_base_url=self.vault_url,
+            certificate_name=certificate_name,
+            cls=lambda pipeline_response, deserialized, _: (pipeline_response, deserialized),
+            **kwargs,
         )
         recovered_certificate = KeyVaultCertificate._from_certificate_bundle(recovered_cert_bundle)
 
         command = partial(self.get_certificate, certificate_name=certificate_name, **kwargs)
         polling_method = AsyncDeleteRecoverPollingMethod(
-            command=command, final_resource=recovered_certificate, finished=False, interval=polling_interval
+            pipeline_response=pipeline_response,
+            command=command,
+            final_resource=recovered_certificate,
+            finished=False,
+            interval=polling_interval
         )
         await polling_method.run()
 

sdk/keyvault/azure-keyvault-certificates/azure/keyvault/certificates/aio/_polling_async.py

@@ -2,22 +2,27 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-
-import asyncio
 import logging
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, cast, Optional, Union
+
+from azure.core.pipeline import PipelineResponse
+from azure.core.pipeline.transport import AsyncHttpTransport
 from azure.core.polling import AsyncPollingMethod
+
 from .._models import KeyVaultCertificate, CertificateOperation
 
 
 logger = logging.getLogger(__name__)
 
 
 class CreateCertificatePollerAsync(AsyncPollingMethod):
-    def __init__(self, get_certificate_command: Callable, interval: int = 5) -> None:
-        self._command = None  # type: Optional[Callable]
-        self._resource = None  # type: Optional[Union[CertificateOperation, KeyVaultCertificate]]
-        self._pending_certificate_op = None  # type: Optional[CertificateOperation]
+    def __init__(
+            self, pipeline_response: PipelineResponse, get_certificate_command: Callable, interval: int = 5
+        ) -> None:
+        self._pipeline_response = pipeline_response
+        self._command: Optional[Callable] = None
+        self._resource: Optional[Union[CertificateOperation, KeyVaultCertificate]] = None
+        self._pending_certificate_op: Optional[CertificateOperation] = None
         self._get_certificate_command = get_certificate_command
         self._polling_interval = interval
 
@@ -33,7 +38,9 @@ async def run(self) -> None:
             while not self.finished():
                 await self._update_status()
                 if not self.finished():
-                    await asyncio.sleep(self._polling_interval)
+                    # We should always ask the client's transport to sleep, instead of sleeping directly
+                    transport: AsyncHttpTransport = cast(AsyncHttpTransport, self._pipeline_response.context.transport)
+                    await transport.sleep(self._polling_interval)
             operation = self._pending_certificate_op
             if operation and operation.status and operation.status.lower() == "completed":
                 self._resource = await self._get_certificate_command()

sdk/keyvault/azure-keyvault-keys/CHANGELOG.md

@@ -10,6 +10,8 @@
 
 ### Other Changes
 - Python 3.7 is no longer supported. Please use Python version 3.8 or later.
+- `asyncio` is no longer directly referenced by the library
+  ([#33819](https://github.com/Azure/azure-sdk-for-python/pull/33819))
 
 ## 4.9.0b3 (2023-11-03)