LRO continuation_token

sdk/core/azure-core/azure/core/pipeline/__init__.py

@@ -69,8 +69,21 @@ def __init__(self, transport, **kwargs):  # pylint: disable=super-init-not-calle
         self.options = kwargs
         self._protected = ["transport", "options"]
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        # Remove the unpicklable entries.
+        del state['transport']
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        # Re-create the unpickable entries
+        self.transport = None
+
     def __setitem__(self, key, item):
-        if key in self._protected:
+        # If reloaded from pickle, _protected might not be here until restored by pickle
+        # this explains the hasattr test
+        if hasattr(self, '_protected') and key in self._protected:
             raise ValueError("Context value {} cannot be overwritten.".format(key))
         return super(PipelineContext, self).__setitem__(key, item)
 

sdk/core/azure-core/azure/core/pipeline/transport/_aiohttp.py

@@ -299,3 +299,14 @@ def stream_download(self, pipeline) -> AsyncIteratorType[bytes]:
         :type pipeline: azure.core.pipeline
         """
         return AioHttpStreamDownloadGenerator(pipeline, self)
+
+    def __getstate__(self):
+        # Be sure body is loaded in memory, otherwise not pickable and let it throw
+        self.body()
+
+        state = self.__dict__.copy()
+        # Remove the unpicklable entries.
+        state['internal_response'] = None  # aiohttp response are not pickable (see headers comments)
+        from multidict import MultiDict  # I know it's importable since aiohttp is loaded
+        state['headers'] = MultiDict(self.headers)  # MultiDictProxy is not pickable
+        return state

sdk/core/azure-core/azure/core/polling/__init__.py

@@ -31,5 +31,5 @@
 #pylint: disable=unused-import
 if sys.version_info >= (3, 5, 2):
     # Not executed on old Python, no syntax error
-    from ._async_poller import AsyncNoPolling, AsyncPollingMethod, async_poller
-    __all__ += ['AsyncNoPolling', 'AsyncPollingMethod', 'async_poller']
+    from ._async_poller import AsyncNoPolling, AsyncPollingMethod, async_poller, AsyncLROPoller
+    __all__ += ['AsyncNoPolling', 'AsyncPollingMethod', 'async_poller', 'AsyncLROPoller']

sdk/core/azure-core/azure/core/polling/_async_poller.py

@@ -23,6 +23,8 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
+from collections.abc import Awaitable
+
 from ._poller import NoPolling as _NoPolling
 
 class AsyncPollingMethod(object):
@@ -43,6 +45,15 @@ def finished(self):
     def resource(self):
         raise NotImplementedError("This method needs to be implemented")
 
+    def get_continuation_token(self):
+        # type() -> str
+        raise NotImplementedError("This polling method doesn't support get_continuation_token")
+
+    @classmethod
+    def from_continuation_token(cls, continuation_token, **kwargs):
+        # type(str, Any) -> PollingMethod
+        raise NotImplementedError("This polling method doesn't support from_continuation_token")
+
 
 class AsyncNoPolling(_NoPolling):
     """An empty async poller that returns the deserialized initial response.
@@ -56,6 +67,9 @@ async def run(self):
 async def async_poller(client, initial_response, deserialization_callback, polling_method):
     """Async Poller for long running operations.
 
+    .. deprecated:: 1.5.0
+       Use :class:`AsyncLROPoller` instead.
+
     :param client: A pipeline service client.
     :type client: ~azure.core.PipelineClient
     :param initial_response: The initial call response
@@ -66,15 +80,59 @@ async def async_poller(client, initial_response, deserialization_callback, polli
     :param polling_method: The polling strategy to adopt
     :type polling_method: ~azure.core.polling.PollingMethod
     """
+    poller = AsyncLROPoller(client, initial_response, deserialization_callback, polling_method)
+    return await poller
 
-    # This implicit test avoids bringing in an explicit dependency on Model directly
-    try:
-        deserialization_callback = deserialization_callback.deserialize
-    except AttributeError:
-        pass
+class AsyncLROPoller(Awaitable):
+    """Async poller for long running operations.
 
-    # Might raise a CloudError
-    polling_method.initialize(client, initial_response, deserialization_callback)
+    :param client: A pipeline service client
+    :type client: ~azure.core.PipelineClient
+    :param initial_response: The initial call response
+    :type initial_response:
+     ~azure.core.pipeline.transport.HttpResponse or ~azure.core.pipeline.transport.AsyncHttpResponse
+    :param deserialization_callback: A callback that takes a Response and return a deserialized object.
+                                     If a subclass of Model is given, this passes "deserialize" as callback.
+    :type deserialization_callback: callable or msrest.serialization.Model
+    :param polling_method: The polling strategy to adopt
+    :type polling_method: ~azure.core.polling.PollingMethod
+    """
+
+    def __init__(self, client, initial_response, deserialization_callback, polling_method):
+        # type: (Any, HttpResponse, DeserializationCallbackType, PollingMethod) -> None
+        self._polling_method = polling_method
+
+        # This implicit test avoids bringing in an explicit dependency on Model directly
+        try:
+            deserialization_callback = deserialization_callback.deserialize # type: ignore
+        except AttributeError:
+            pass
+
+        # Might raise a CloudError
+        self._polling_method.initialize(client, initial_response, deserialization_callback)
+
+    async def _start(self):
+        """Start the long running operation.
+        """
+        await self._polling_method.run()
+        return self._polling_method.resource()
+
+    def __await__(self):
+        return self._start().__await__()
+
+    def continuation_token(self):
+        # type: () -> str
+        """Return a continuation token that allows to restart the poller later.
+
+        :returns: An opaque continuation token
+        :rtype: str
+        """
+        return self._polling_method.get_continuation_token()
 
-    await polling_method.run()
-    return polling_method.resource()
+    @classmethod
+    def from_continuation_token(cls, polling_method, continuation_token, **kwargs):
+        # type: (str, Any) -> LROPoller
+        client, initial_response, deserialization_callback = polling_method.from_continuation_token(
+            continuation_token, **kwargs
+        )
+        return cls(client, initial_response, deserialization_callback, polling_method)

sdk/core/azure-core/azure/core/polling/_poller.py

@@ -30,7 +30,7 @@
 except ImportError:
     from urllib.parse import urlparse
 
-from typing import Any, Callable, Union, List, Optional, TYPE_CHECKING
+from typing import Any, Callable, Union, List, Optional, Tuple, TYPE_CHECKING
 from azure.core.pipeline.transport._base import HttpResponse  # type: ignore
 from azure.core.tracing.decorator import distributed_trace
 from azure.core.tracing.common import with_current_context
@@ -64,6 +64,16 @@ def resource(self):
         # type: () -> Any
         raise NotImplementedError("This method needs to be implemented")
 
+    def get_continuation_token(self):
+        # type() -> str
+        raise NotImplementedError("This polling method doesn't support get_continuation_token")
+
+    @classmethod
+    def from_continuation_token(cls, continuation_token, **kwargs):
+        # type(str, Any) -> PollingMethod
+        raise NotImplementedError("This polling method doesn't support from_continuation_token")
+
+
 class NoPolling(PollingMethod):
     """An empty poller that returns the deserialized initial response.
     """
@@ -101,6 +111,22 @@ def resource(self):
         # type: () -> Any
         return self._deserialization_callback(self._initial_response)
 
+    def get_continuation_token(self):
+        # type() -> str
+        import pickle
+        return pickle.dumps(self._initial_response)
+
+    @classmethod
+    def from_continuation_token(cls, continuation_token, **kwargs):
+        # type(str, Any) -> Tuple
+        try:
+            deserialization_callback = kwargs["deserialization_callback"]
+        except KeyError:
+            raise ValueError("Need kwarg 'deserialization_callback' to be recreated from continuation_token")
+        import pickle
+        initial_response = pickle.loads(continuation_token)
+        return None, initial_response, deserialization_callback
+
 
 class LROPoller(object):
     """Poller for long running operations.
@@ -119,8 +145,6 @@ class LROPoller(object):
 
     def __init__(self, client, initial_response, deserialization_callback, polling_method):
         # type: (Any, HttpResponse, DeserializationCallbackType, PollingMethod) -> None
-        self._client = client
-        self._response = initial_response
         self._callbacks = []  # type: List[Callable]
         self._polling_method = polling_method
 
@@ -131,7 +155,7 @@ def __init__(self, client, initial_response, deserialization_callback, polling_m
             pass
 
         # Might raise a CloudError
-        self._polling_method.initialize(self._client, self._response, deserialization_callback)
+        self._polling_method.initialize(client, initial_response, deserialization_callback)
 
         # Prepare thread execution
         self._thread = None
@@ -166,6 +190,23 @@ def _start(self):
                 call(self._polling_method)
             callbacks, self._callbacks = self._callbacks, []
 
+    def continuation_token(self):
+        # type: () -> str
+        """Return a continuation token that allows to restart the poller later.
+
+        :returns: An opaque continuation token
+        :rtype: str
+        """
+        return self._polling_method.get_continuation_token()
+
+    @classmethod
+    def from_continuation_token(cls, polling_method, continuation_token, **kwargs):
+        # type: (str, Any) -> LROPoller
+        client, initial_response, deserialization_callback = polling_method.from_continuation_token(
+            continuation_token, **kwargs
+        )
+        return cls(client, initial_response, deserialization_callback, polling_method)
+
     def status(self):
         # type: () -> str
         """Returns the current status string.

sdk/core/azure-core/azure/core/polling/base_polling.py

@@ -447,6 +447,30 @@ def initialize(self, client, initial_response, deserialization_callback):
         except OperationFailed as err:
             raise HttpResponseError(response=initial_response.http_response, error=err)
 
+    def get_continuation_token(self):
+        # type() -> str
+        import pickle
+        return pickle.dumps(self._initial_response)
+
+    @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")
+
+        try:
+            deserialization_callback = kwargs["deserialization_callback"]
+        except KeyError:
+            raise ValueError("Need kwarg 'deserialization_callback' to be recreated from continuation_token")
+
+        import pickle
+        initial_response = pickle.loads(continuation_token)
+        # Restore the transport in the context
+        initial_response.context.transport = client._pipeline._transport  # pylint: disable=protected-access
+        return client, initial_response, deserialization_callback
+
     def run(self):
         try:
             self._poll()