Operations that are started by an initial call, then need to be monitored for status until completion are often represented as long-running operations (LRO).
The azure-core library provides the LROPoller (and AsyncLROPoller) protocols
that expose methods to interact with the LRO such as waiting until the operation reaches a terminal state, checking its status, or
providing a callback to do work on the final result when it is ready. If the LRO follows the
Azure REST API guidelines,
it's likely that the generated client library code should just work.
In cases where the LRO diverges from the guidelines, you may need to customize your code to achieve the desired scenario.
There are 3 options to customize the logic for LROs.
- Polling strategy - OperationResourcePolling, LocationPolling, StatusCheckPolling
- You need to customize the polling strategy
- Polling method - LROBasePolling/AsyncLROBasePolling
- You need to customize the polling loop
- Poller API - LROPoller/AsyncLROPoller
- You need to customize the public interface of the poller
The "poller API" is what the user uses to interact with the LRO. Internally, the poller uses the "polling method" to run the polling loop which makes calls to the status monitor, controls delay, and determines when the LRO has reached a terminal state. The "polling method" uses a "polling strategy" to determine how to extract the status information from the responses returned by the status monitoring.
The azure.core.polling module provides three built-in strategies for polling -OperationResourcePolling,
LocationPolling, and StatusCheckPolling. The type of polling needed will be
determined automatically using the response structure, unless otherwise specified by the client library developer. If
the LRO is determined not to fit either OperationResourcePolling or LocationPolling, StatusCheckPolling serves as
a fallback strategy which will not perform polling, but instead return a successful response for a 2xx status code.
If you need to customize the polling strategy, choose a polling algorithm that closely represents what you need to do or create your own that inherits from azure.core.polling.base_polling.LongRunningOperation and implements the necessary methods.
For our example, let's say that OperationResourcePolling closely resembles what the service does, but we
need to account for a non-standard status.
Example: Raise exception for a non-standard status that is returned via the initial POST call - "ValidationFailed".
from azure.core.polling.base_polling import OperationResourcePolling, _as_json
from azure.core.exceptions import HttpResponseError
class CustomOperationResourcePolling(OperationResourcePolling):
"""Implements a operation resource polling, typically from the Operation-Location header.
Customized to raise an exception in the case of a "ValidationFailed" status returned
from the service.
"""
def get_status(self, pipeline_response: "PipelineResponseType") -> str:
"""This method is called on the response for each polling request
and is used to extract and return the LRO status from that response.
In the case that the operation has failed (i.e. a non-successful status),
an exception should be raised. This will bring polling to an end and raise
the failure to the listener.
"""
status = super().get_status(pipeline_response)
if status.lower() == "validationfailed":
response = pipeline_response.http_response
body = _as_json(response)
raise HttpResponseError(response=response, error=body["error"])
return statusYou can then wrap a client method that was generated as an LRO, and pass the additional polling keyword argument. The polling
keyword argument takes an implementation of azure.core.polling.PollingMethod
(e.g. azure.core.polling.base_polling.LROBasePolling) and allows for a custom strategy to be passed
in to the keyword argument lro_algorithms:
from typing import AnyStr, MutableMapping, Any
from azure.core.polling import LROPoller
from azure.core.polling.base_polling import LROBasePolling
JSON = MutableMapping[str, Any]
class ServiceOperations:
def begin_analyze(self, data: AnyStr, name: str, **kwargs) -> LROPoller[JSON]:
return self._generated_client.begin_analyze(
data,
name,
polling=LROBasePolling(
lro_algorithms=[
CustomOperationResourcePolling() # overrides other LRO strategies
]
),
**kwargs
)This example uses the default polling method - LROBasePolling and just overrides the strategy used for polling.
If you need to control the polling loop, then see the next section.
Built-in methods for polling are included in azure-core as both sync / async variants - LROBasePolling
and AsyncLROBasePolling. The polling method runs the polling loop and performs GET requests
to the status monitor to check if a terminal state is reached. In between polls it inserts delay based on
- the service sent
retry-afterheader, or 2) the givenpolling_intervalif no retry-after header is present.
You can also use azure.core.polling.NoPolling(or AsyncNoPolling) which will not
initiate polling and simply return the deserialized initial response when called with poller.result().
To use NoPolling, you can pass polling=False to an operation generated as an LRO:
from typing import AnyStr, MutableMapping, Any
from azure.core.polling import LROPoller
JSON = MutableMapping[str, Any]
class ServiceOperations:
def begin_analyze(self, data: AnyStr, name: str, **kwargs) -> LROPoller[JSON]:
return self._generated_client.begin_analyze(
data,
name,
polling=False,
**kwargs
)To customize parts of the polling method, you can create a subclass which uses LROBasePolling and overrides necessary methods. If significant customization is necessary, use azure.core.polling.PollingMethod (or AsyncPollingMethod)and implement all the necessary methods.
Example: Create an LRO method which will poll for when a file gets uploaded successfully (greatly simplified)
For this example, the customization necessary requires defining our own polling strategy and polling method. First, we'll define a simple polling strategy that the polling method will use to get the status information from the response.
from typing import Optional, MutableMapping, Any
from azure.core.polling.base_polling import LongRunningOperation, OperationFailed
from azure.core.pipeline import PipelineResponse
JSON = MutableMapping[str, Any]
class CustomPollingStrategy(LongRunningOperation):
"""CustomPollingStrategy which provides default logic
for interpreting operation responses and status updates.
"""
def can_poll(self, pipeline_response: PipelineResponse) -> bool:
"""Determine from the initial response that we can poll.
In this example, we need a file_id present to proceed with polling.
:param PipelineResponse pipeline_response: initial REST call response.
"""
response = pipeline_response.http_response.json()
if response.get("file_id", None) is None:
return False
return True
def get_polling_url(self) -> str:
"""Return the polling URL. This is the URL for the status monitor
and where the GET requests will be made during polling.
For this example, we don't need to extract the URL
from the initial response so it is not implemented.
"""
raise NotImplementedError("The polling strategy does not need to extract a polling URL.")
def set_initial_status(self, pipeline_response: PipelineResponse) -> str:
"""Process first response after initiating long running operation and set initial status.
:param PipelineResponse pipeline_response: initial REST call response.
"""
response = pipeline_response.http_response
if response.status_code == 200:
return "InProgress"
raise OperationFailed("Operation failed or canceled")
def get_status(self, response: JSON) -> str:
"""Return the status based on this response.
Typically, this method extracts a status string from the
response. In this example, we determine status based on whether our
result is populated or not.
"""
if response is None:
return "InProgress"
return "Succeeded"
def get_final_get_url(self, pipeline_response: PipelineResponse) -> Optional[str]:
"""If a final GET is needed when the LRO is complete, returns the URL.
:rtype: str
"""
return NoneNext, we'll define the custom polling method:
import functools
import time
import base64
from typing import Any, Tuple, Callable, MutableMapping
from azure.core.pipeline import PipelineResponse
from azure.core.polling import PollingMethod
from azure.core.polling.base_polling import BadResponse
from azure.core.exceptions import ResourceNotFoundError
JSON = MutableMapping[str, Any]
class CustomPollingMethod(PollingMethod):
def __init__(self, polling_interval: float = 30, **kwargs: Any) -> None:
"""Creates a custom polling method which polls until a file is uploaded.
For our example, the operation is considered to have reached a terminal state once a successful GET
is done on the file (e.g. no ResourceNotFoundError is raised).
:param polling_interval: The amount of time to wait between polls. This fictitious service does not
use retry-after so we will default to this value.
:param kwargs: Any operation-specific keyword arguments that should be passed into the GET call.
"""
self._polling_interval = polling_interval
self._kwargs = kwargs
def initialize(self, client: Any, initial_response: PipelineResponse, deserialization_callback: Callable) -> None:
"""Set the initial status of this LRO, verify that we can poll, and
initialize anything necessary for polling.
:param client: An instance of a client. In this example, the generated client.
:param initial_response: In this example, the PipelineResponse returned from the initial call.
:param deserialization_callback: A callable to transform the final result before returning to the end user.
"""
# verify we have the information to poll
if self._operation.can_poll(initial_response) is False:
raise BadResponse("No file_id in response.")
response = initial_response.http_response.json()
# initialize
self.client = client
self.file_id = response["file_id"]
self._initial_response = initial_response
self._deserialization_callback = deserialization_callback
self._resource = None
self._finished = False
# sets our strategy
self._operation = CustomPollingStrategy()
# create the command which will be polled against as the status monitor
self._command = functools.partial(self.client.get_upload_file, file_id=self.file_id, **self._kwargs)
# set initial status
self._status = self._operation.set_initial_status(initial_response)
def status(self) -> str:
"""Should return the current status as a string. The initial status is set by
the polling strategy with set_initial_status() and then subsequently set by
each call to get_status().
This is what is returned to the user when status() is called on the LROPoller.
:rtype: str
"""
return self._status
def finished(self) -> bool:
"""Is this polling finished?
Controls whether the polling loop should continue to poll.
:returns: Return True if the operation has reached a terminal state
or False if polling should continue.
:rtype: bool
"""
return True if self.status() == "Succeeded" else False
def resource(self) -> JSON:
"""Return the built resource.
This is what is returned when to the user when result() is called on the LROPoller.
This might include a deserialization callback (passed in initialize())
to transform or customize the final result, if necessary.
"""
return self._deserialization_callback(self._resource)
def run(self) -> None:
"""The polling loop.
The polling should call the status monitor, evaluate and set the current status,
insert delay between polls, and continue polling until a terminal state is reached.
"""
while not self.finished():
self.update_status()
if not self.finished():
# inserts delay if not done
time.sleep(self._polling_interval)
def update_status(self):
"""Update the current status of the LRO by calling the status monitor
and then using the polling strategy's get_status() to set the status."""
try:
self._resource = self._command()
except ResourceNotFoundError:
pass
self._status = self._operation.get_status(self._resource)
def get_continuation_token(self) -> str:
"""Returns an opaque token which can be used by the user to rehydrate/restart the LRO.
Saves the initial state of the LRO so that polling can be resumed from that context.
.. code-block:: python
initial_poller = client.begin_upload(data)
continuation_token = initial_poller.continuation_token()
poller: LROPoller = client.begin_upload(None, continuation_token=continuation_token)
poller.result()
In standard LROs, the PipelineResponse is serialized here, however, there may be a need to
customize this further depending on your scenario.
It's recommended to:
1. Include a version number for future compatibility.
2. Filter headers to only include those needed for LRO rehydration, avoiding sensitive data.
"""
import json
# Headers needed for LRO rehydration - use an allowlist approach for security
lro_headers = {"operation-location", "location", "content-type", "retry-after"}
response = self._initial_response.http_response
filtered_headers = {k: v for k, v in response.headers.items() if k.lower() in lro_headers}
# Use a versioned token schema: {"version": <int>, "data": <your state>}
# This allows for future compatibility checking when deserializing
token = {
"version": 1,
"data": {
"file_id": self.file_id,
"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")
@classmethod
def from_continuation_token(cls, continuation_token: str, **kwargs: Any) -> Tuple[Any, PipelineResponse, Callable]:
"""Deserializes the user-provided continuation_token to the initial response and returns
the context necessary to rebuild the LROPoller from its classmethod.
"""
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 json
token = json.loads(base64.b64decode(continuation_token).decode("utf-8"))
# Validate token schema and version for compatibility
if not isinstance(token, dict) or "version" not in token or "data" not in token:
raise ValueError("Invalid continuation token format.")
if token["version"] != 1:
raise ValueError(
"This continuation token is not compatible with this version. "
"It may have been generated by a different version."
)
# Extract the state from the "data" field
state = token["data"]
# The file_id and other state can be extracted and used to resume polling
return client, state, deserialization_callbackAnd now, to plug into the client code:
from typing import AnyStr, MutableMapping, Any
from azure.core.polling import LROPoller
JSON = MutableMapping[str, Any]
class ServiceOperations:
def begin_upload(self, data: AnyStr, **kwargs) -> LROPoller[JSON]:
continuation_token = kwargs.pop("continuation_token", None)
polling_method = CustomPollingMethod(**kwargs)
# if continuation_token is provided, we should rehydrate the LRO using the from_continuation_token method
# which calls our implementation on the CustomPollingMethod method
if continuation_token is not None:
return LROPoller.from_continuation_token(
continuation_token=continuation_token,
polling_method=polling_method,
deserialization_callback=lambda x: x,
client=self
)
# the initial call. We pass in `cls` to receive the pipeline_response.
pipeline_response = self._generated_client.create_upload(data, cls=lambda response, x, y: response, **kwargs)
return LROPoller(
client=self,
initial_response=pipeline_response,
deserialization_callback=lambda x: x, # returning the result as-is, but could be a callable to transform the final result
polling_method=polling_method,
)Note that we need to account for a continuation_token being passed by the user, in which case we should not make the
initial call again, but rather resume polling from the rehydrated state. Since passing continuation_token doesn't
require the user to provide the parameters for the initial call, it can be helpful to add overloads to the method to
clarify its usage, especially in cases where required parameters become non-required:
from typing import AnyStr, MutableMapping, Any, overload
from azure.core.polling import LROPoller
JSON = MutableMapping[str, Any]
class ServiceOperations:
@overload
def begin_upload(self, data: AnyStr, **kwargs: Any) -> LROPoller[JSON]:
...
@overload
def begin_upload(self, *, continuation_token: str, **kwargs: Any) -> LROPoller[JSON]:
...
def begin_upload(self, *args, **kwargs) -> LROPoller[JSON]:
continuation_token = kwargs.pop("continuation_token", None)
polling_method = CustomPollingMethod(**kwargs)
if continuation_token is not None:
return LROPoller.from_continuation_token(
continuation_token=continuation_token,
polling_method=polling_method,
deserialization_callback=lambda x: x,
client=self
)
data = kwargs.pop("data", None)
if data is None:
try:
data = args[0]
except IndexError:
raise TypeError("begin_upload() missing 1 required positional argument: 'data'")
pipeline_response = self._generated_client.create_upload(data, cls=lambda response, x, y: response, **kwargs)
return LROPoller(
client=self,
initial_response=pipeline_response,
deserialization_callback=lambda x: x,
polling_method=polling_method,
)The last option is if you need to customize the public interface of the LROPoller / AsyncLROPoller.
Reasons to do this might include exposing important attributes or metadata of the operation in progress,
or adding new features to interact with the operation, such as to pause/resume or cancel it.
This example builds off the previous example and uses the custom polling method defined above. The custom polling
method gives us access to the client and file_id needed to make the cancel call. If you support rehydration of
the LRO via continuation_token, you must override the from_continuation_token method so that the custom poller is used.
from typing import TypeVar
from azure.core.polling import LROPoller, PollingMethod
PollingReturnType_co = TypeVar("PollingReturnType_co", covariant=True)
class CustomLROPoller(LROPoller[PollingReturnType_co]):
def cancel(self, **kwargs) -> None:
"""Cancel the upload"""
return self.polling_method().client.cancel_upload_file(self.polling_method().file_id, **kwargs)
@classmethod
def from_continuation_token(
cls, polling_method: PollingMethod[PollingReturnType_co], continuation_token: str, **kwargs
) -> "CustomLROPoller[PollingReturnType_co]":
(
client,
initial_response,
deserialization_callback,
) = polling_method.from_continuation_token(continuation_token, **kwargs)
return cls(client, initial_response, deserialization_callback, polling_method)And now, to plug into the client code:
from typing import AnyStr, MutableMapping, Any
JSON = MutableMapping[str, Any]
class ServiceOperations:
def begin_upload(self, data: AnyStr, **kwargs) -> CustomLROPoller[JSON]:
continuation_token = kwargs.pop("continuation_token", None)
polling_method = CustomPollingMethod(**kwargs)
if continuation_token is not None:
return CustomLROPoller.from_continuation_token(
continuation_token=continuation_token,
deserialization_callback=lambda x: x,
polling_method=polling_method,
client=self
)
response = self._generated_client.create_upload(data, **kwargs)
return CustomLROPoller[JSON](
client=self,
initial_response=response,
deserialization_callback=lambda x: x,
polling_method=polling_method,
)Note that we updated the begin_upload return type to CustomLROPoller. You should only need to explicitly reference
the custom poller if a new public API has been added. The custom poller should additionally be added to the package
__init__.py so that the new public API will be properly documented.