Azure Communication Services - Phone Number Admin - Implementing long running operations for phone number search

sdk/communication/azure-communication-administration/azure/communication/administration/__init__.py

@@ -6,6 +6,7 @@
 
 from ._communication_identity_client import CommunicationIdentityClient
 from ._phone_number_administration_client import PhoneNumberAdministrationClient
+from ._polling import PhoneNumberPolling
 
 from ._identity._generated.models import (
     CommunicationTokenRequest,
@@ -16,6 +17,7 @@
     AcquiredPhoneNumber,
     AcquiredPhoneNumbers,
     AreaCodes,
+    CreateSearchOptions,
     CreateSearchResponse,
     LocationOptionsQuery,
     LocationOptionsResponse,
@@ -31,7 +33,7 @@
     ReleaseResponse,
     UpdateNumberCapabilitiesResponse,
     UpdatePhoneNumberCapabilitiesResponse,
-    CreateSearchOptions
+    SearchStatus
 )
 
 from ._shared.models import (
@@ -43,6 +45,7 @@
 __all__ = [
     'CommunicationIdentityClient',
     'PhoneNumberAdministrationClient',
+    'PhoneNumberPolling',
 
     # from _identity
     'CommunicationTokenRequest',
@@ -52,22 +55,23 @@
     'AcquiredPhoneNumber',
     'AcquiredPhoneNumbers',
     'AreaCodes',
+    'CreateSearchOptions',
     'CreateSearchResponse',
     'LocationOptionsQuery',
     'LocationOptionsResponse',
     'NumberConfigurationResponse',
     'NumberUpdateCapabilities',
     'PhoneNumberCountries',
     'PhoneNumberEntities',
-    'PhoneNumberRelease',
     'PhoneNumberSearch',
+    'PhoneNumberRelease',
     'PhonePlanGroups',
     'PhonePlansResponse',
     'PstnConfiguration',
     'ReleaseResponse',
+    'SearchStatus',
     'UpdateNumberCapabilitiesResponse',
     'UpdatePhoneNumberCapabilitiesResponse',
-    'CreateSearchOptions',
 
     # from _shared
     'CommunicationUser',

sdk/communication/azure-communication-administration/azure/communication/administration/_phone_number_administration_client.py

@@ -6,14 +6,15 @@
 # ------------------------------------
 from azure.core.tracing.decorator import distributed_trace
 from azure.core.paging import ItemPaged
+from azure.core.polling import LROPoller
+from ._polling import PhoneNumberPolling
 
 from ._phonenumber._generated._phone_number_administration_service\
     import PhoneNumberAdministrationService as PhoneNumberAdministrationClientGen
 
 from ._phonenumber._generated.models import (
     AcquiredPhoneNumbers,
     AreaCodes,
-    CreateSearchResponse,
     LocationOptionsResponse,
     NumberConfigurationResponse,
     NumberUpdateCapabilities,
@@ -25,6 +26,7 @@
     PhonePlansResponse,
     PstnConfiguration,
     ReleaseResponse,
+    SearchStatus,
     UpdateNumberCapabilitiesResponse,
     UpdatePhoneNumberCapabilitiesResponse
 )
@@ -405,21 +407,51 @@ def get_search_by_id(
         )
 
     @distributed_trace
-    def create_search(
+    def begin_create_search(
         self,
         **kwargs  # type: Any
     ):
-        # type: (...) -> CreateSearchResponse
-        """Creates a phone number search.
+        # type: (...) -> LROPoller
+        """Begins creating a phone number search.
 
         :keyword azure.communication.administration.CreateSearchOptions body:
-        An optional parameter for defining the search options.
+        A parameter for defining the search options.
         The default is None.
-        :rtype: ~azure.communication.administration.CreateSearchResponse
+        :keyword str continuation_token: A continuation token to restart a poller from a saved state.
+        Caller must provide either body, or continuation_token keywords to use the method.
+        :rtype: ~azure.core.polling.LROPoller[~azure.communication.administration.PhoneNumberSearch]
         """
-        return self._phone_number_administration_client.phone_number_administration.create_search(
+        polling_interval = kwargs.pop("_polling_interval", 5)
+        cont_token = kwargs.pop('continuation_token', None)  # type: Optional[str]
+
+        search_polling = PhoneNumberPolling(
+            is_terminated=lambda status: status in [
+                SearchStatus.Reserved,
+                SearchStatus.Expired,
+                SearchStatus.Success,
+                SearchStatus.Cancelled,
+                SearchStatus.Error
+            ],
+            interval=polling_interval
+        )
+
+        if cont_token is not None:
+            return LROPoller.from_continuation_token(
+                polling_method=search_polling,
+                continuation_token=cont_token,
+                client=self._phone_number_administration_client.phone_number_administration
+            )
+
+        create_search_response = self._phone_number_administration_client.phone_number_administration.create_search(
             **kwargs
         )
+        initial_state = self._phone_number_administration_client.phone_number_administration.get_search_by_id(
+            search_id=create_search_response.search_id
+        )
+        return LROPoller(client=self._phone_number_administration_client.phone_number_administration,
+                         initial_response=initial_state,
+                         deserialization_callback=None,
+                         polling_method=search_polling)
 
     @distributed_trace
     def list_all_searches(
@@ -440,37 +472,96 @@ def list_all_searches(
         )
 
     @distributed_trace
-    def cancel_search(
+    def begin_cancel_search(
         self,
-        search_id,  # type: str
         **kwargs  # type: Any
     ):
-        # type: (...) -> None
-        """Cancels the search. This means existing numbers in the search will be made available.
+        # type: (...) -> LROPoller
+        """Begins the phone number search cancellation.
 
-        :param search_id: The search id to be canceled.
-        :type search_id: str
-        :rtype: None
+        :keyword str search_id: The search id to be canceled.
+        :keyword str continuation_token: An optional continuation token to restart a poller from a saved state.
+        Caller must provide either search_id, or continuation_token keywords to use the method.
+
+        :rtype: ~azure.core.polling.LROPoller[~azure.communication.administration.PhoneNumberSearch]
         """
-        return self._phone_number_administration_client.phone_number_administration.cancel_search(
+        polling_interval = kwargs.pop("_polling_interval", 5)
+        cont_token = kwargs.pop('continuation_token', None)  # type: Optional[str]
+
+        search_polling = PhoneNumberPolling(
+            is_terminated=lambda status: status in [
+                SearchStatus.Expired,
+                SearchStatus.Cancelled,
+                SearchStatus.Error
+            ],
+            interval=polling_interval
+        )
+
+        if cont_token is not None:
+            return LROPoller.from_continuation_token(
+                polling_method=search_polling,
+                continuation_token=cont_token,
+                client=self._phone_number_administration_client.phone_number_administration
+            )
+
+        search_id = kwargs.pop('search_id')  # type: str
+
+        self._phone_number_administration_client.phone_number_administration.cancel_search(
             search_id,
             **kwargs
         )
+        initial_state = self._phone_number_administration_client.phone_number_administration.get_search_by_id(
+            search_id=search_id
+        )
+        return LROPoller(client=self._phone_number_administration_client.phone_number_administration,
+                         initial_response=initial_state,
+                         deserialization_callback=None,
+                         polling_method=search_polling)
 
     @distributed_trace
-    def purchase_search(
-            self,
-            search_id,  # type: str
-            **kwargs  # type: Any
+    def begin_purchase_search(
+        self,
+        **kwargs  # type: Any
     ):
-        # type: (...) -> None
-        """Purchases the phone number search.
+        # type: (...) -> LROPoller
+        """Begins the phone number search purchase.
 
-        :param search_id: The search id to be purchased.
-        :type search_id: str
-        :rtype: None
+        :keyword str search_id: The search id to be purchased.
+        :keyword str continuation_token: An optional continuation token to restart a poller from a saved state.
+        Caller must provide either search_id, or continuation_token keywords to use the method.
+
+        :rtype: ~azure.core.polling.LROPoller[~azure.communication.administration.PhoneNumberSearch]
         """
-        return self._phone_number_administration_client.phone_number_administration.purchase_search(
+        polling_interval = kwargs.pop("_polling_interval", 5)
+        cont_token = kwargs.pop('continuation_token', None)  # type: Optional[str]
+
+        search_polling = PhoneNumberPolling(
+            is_terminated=lambda status: status in [
+                SearchStatus.Success,
+                SearchStatus.Expired,
+                SearchStatus.Cancelled,
+                SearchStatus.Error
+            ],
+            interval=polling_interval
+        )
+
+        if cont_token is not None:
+            return LROPoller.from_continuation_token(
+                polling_method=search_polling,
+                continuation_token=cont_token,
+                client=self._phone_number_administration_client.phone_number_administration
+            )
+
+        search_id = kwargs.pop('search_id')  # type: str
+
+        self._phone_number_administration_client.phone_number_administration.purchase_search(
             search_id,
             **kwargs
         )
+        initial_state = self._phone_number_administration_client.phone_number_administration.get_search_by_id(
+            search_id=search_id
+        )
+        return LROPoller(client=self._phone_number_administration_client.phone_number_administration,
+                         initial_response=initial_state,
+                         deserialization_callback=None,
+                         polling_method=search_polling)

sdk/communication/azure-communication-administration/azure/communication/administration/_polling.py

@@ -0,0 +1,76 @@
+# pylint: disable=W0231
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+import base64
+import time
+from typing import Union
+from functools import partial
+import pickle
+
+from azure.core.polling import (
+    PollingMethod
+)
+from ._phonenumber._generated.models import (
+    PhoneNumberSearch,
+    PhoneNumberRelease
+)
+
+class PhoneNumberPolling(PollingMethod):
+    def __init__(self, is_terminated, interval=5):
+        self._response = None
+        self._client = None
+        self._query_status = None
+        self._is_terminated = is_terminated
+        self._polling_interval = interval
+
+    def _update_status(self):
+        # type: () -> None
+        if self._query_status is None:
+            raise Exception("this poller has not been initialized")
+        self._response = self._query_status()
+
+    def initialize(self, client, initial_response, _):
+        # type: (Any, Any, Callable) -> None
+        self._client = client
+        self._response = initial_response
+        self._query_status = partial(self._client.get_search_by_id, search_id=initial_response.search_id)
+
+    def run(self):
+        # type: () -> None
+        while not self.finished():
+            self._update_status()
+            if not self.finished():
+                time.sleep(self._polling_interval)
+
+    def finished(self):
+        # type: () -> bool
+        if self._response.status is None:
+            return False
+        return self._is_terminated(self._response.status)
+
+    def resource(self):
+        # type: () -> Union[PhoneNumberSearch, PhoneNumberRelease]
+        if not self.finished():
+            return None
+        return self._response
+
+    def status(self):
+        # type: () -> str
+        return self._response.status
+
+    def get_continuation_token(self):
+        # type() -> str
+        return base64.b64encode(pickle.dumps(self._response)).decode('ascii')
+
+    @classmethod
+    def from_continuation_token(cls, continuation_token, **kwargs):
+        # type(str, Any) -> Tuple
+        try:
+            client = kwargs["client"]
+        except KeyError:
+            raise ValueError("Need kwarg 'client' to be recreated from continuation_token")
+        initial_response = pickle.loads(base64.b64decode(continuation_token))  # nosec
+        return client, initial_response, None

sdk/communication/azure-communication-administration/azure/communication/administration/aio/__init__.py

@@ -1,7 +1,9 @@
 from ._communication_identity_client_async import CommunicationIdentityClient
 from ._phone_number_administration_client_async import PhoneNumberAdministrationClient
+from ._polling_async import PhoneNumberPollingAsync
 
 __all__ = [
     'CommunicationIdentityClient',
-    'PhoneNumberAdministrationClient'
+    'PhoneNumberAdministrationClient',
+    'PhoneNumberPollingAsync'
 ]