diff --git a/sdk/core/azure-core/CHANGELOG.md b/sdk/core/azure-core/CHANGELOG.md index c6d82e9509f5..fb28a1917d86 100644 --- a/sdk/core/azure-core/CHANGELOG.md +++ b/sdk/core/azure-core/CHANGELOG.md @@ -4,6 +4,8 @@ ### Features Added +- Add support for `final-state-via` LRO option in core. #22713 + ### Breaking Changes ### Bugs Fixed diff --git a/sdk/core/azure-core/azure/core/polling/base_polling.py b/sdk/core/azure-core/azure/core/polling/base_polling.py index 6b6fd4ff7c93..01b8e4789768 100644 --- a/sdk/core/azure-core/azure/core/polling/base_polling.py +++ b/sdk/core/azure-core/azure/core/polling/base_polling.py @@ -26,6 +26,7 @@ import abc import base64 import json +from enum import Enum from typing import TYPE_CHECKING, Optional, Any, Union from ..exceptions import HttpResponseError, DecodeError @@ -174,20 +175,38 @@ def get_final_get_url(self, pipeline_response): """ raise NotImplementedError() +class _LroOption(str, Enum): + """Known LRO options from Swagger.""" + + FINAL_STATE_VIA = "final-state-via" + + +class _FinalStateViaOption(str, Enum): + """Possible final-state-via options.""" + + AZURE_ASYNC_OPERATION_FINAL_STATE = "azure-async-operation" + LOCATION_FINAL_STATE = "location" + OPERATION_LOCATION_FINAL_STATE = "operation-location" + class OperationResourcePolling(LongRunningOperation): """Implements a operation resource polling, typically from Operation-Location. :param str operation_location_header: Name of the header to return operation format (default 'operation-location') + :keyword dict[str, any] lro_options: Additional options for LRO. For more information, see + https://aka.ms/azsdk/autorest/openapi/lro-options """ - def __init__(self, operation_location_header="operation-location"): + def __init__( + self, operation_location_header="operation-location", **kwargs + ): self._operation_location_header = operation_location_header # Store the initial URLs self._async_url = None self._location_url = None self._request = None + self._lro_options = kwargs.pop("lro_options", {}) or {} def can_poll(self, pipeline_response): """Answer if this polling method could be used. @@ -207,6 +226,19 @@ def get_final_get_url(self, pipeline_response): :rtype: str """ + if ( + self._lro_options.get(_LroOption.FINAL_STATE_VIA) == _FinalStateViaOption.LOCATION_FINAL_STATE + and self._location_url + ): + return self._location_url + if ( + self._lro_options.get(_LroOption.FINAL_STATE_VIA) + in [ + _FinalStateViaOption.AZURE_ASYNC_OPERATION_FINAL_STATE, + _FinalStateViaOption.OPERATION_LOCATION_FINAL_STATE + ] + ): + return None response = pipeline_response.http_response if not _is_empty(response): body = _as_json(response) @@ -381,7 +413,7 @@ def __init__( **operation_config ): self._lro_algorithms = lro_algorithms or [ - OperationResourcePolling(), + OperationResourcePolling(lro_options=lro_options), LocationPolling(), StatusCheckPolling(), ] diff --git a/sdk/core/azure-core/tests/async_tests/test_base_polling_async.py b/sdk/core/azure-core/tests/async_tests/test_base_polling_async.py index 82a5d21bb4d4..ed16f4712701 100644 --- a/sdk/core/azure-core/tests/async_tests/test_base_polling_async.py +++ b/sdk/core/azure-core/tests/async_tests/test_base_polling_async.py @@ -42,7 +42,7 @@ from msrest import Deserializer -from azure.core.polling import async_poller +from azure.core.polling import async_poller, AsyncLROPoller from azure.core.exceptions import DecodeError, HttpResponseError from azure.core import AsyncPipelineClient from azure.core.pipeline import PipelineResponse, AsyncPipeline, PipelineContext @@ -52,6 +52,7 @@ AsyncLROBasePolling, ) from utils import ASYNCIO_REQUESTS_TRANSPORT_RESPONSES, request_and_responses_product, create_transport_response +from rest_client_async import AsyncTestRestClient class SimpleResource: """An implementation of Python 3 SimpleNamespace. @@ -748,3 +749,125 @@ async def test_long_running_negative(http_request, http_response): LOCATION_BODY = json.dumps({ 'name': TEST_NAME }) POLLING_STATUS = 200 +@pytest.mark.asyncio +@pytest.mark.parametrize("http_request,http_response", request_and_responses_product(ASYNCIO_REQUESTS_TRANSPORT_RESPONSES)) +async def test_post_final_state_via(async_pipeline_client_builder, deserialization_cb, http_request, http_response): + # Test POST LRO with both Location and Operation-Location + CLIENT.http_request_type = http_request + CLIENT.http_response_type = http_response + # The initial response contains both Location and Operation-Location, a 202 and no Body + initial_response = TestBasePolling.mock_send( + http_request, + http_response, + 'POST', + 202, + { + 'location': 'http://example.org/location', + 'operation-location': 'http://example.org/async_monitor', + }, + '' + ) + + async def send(request, **kwargs): + assert request.method == 'GET' + + if request.url == 'http://example.org/location': + return TestBasePolling.mock_send( + http_request, + http_response, + 'GET', + 200, + body={'location_result': True} + ).http_response + elif request.url == 'http://example.org/async_monitor': + return TestBasePolling.mock_send( + http_request, + http_response, + 'GET', + 200, + body={'status': 'Succeeded'} + ).http_response + else: + pytest.fail("No other query allowed") + + client = async_pipeline_client_builder(send) + + # Test 1, LRO options with Location final state + poll = async_poller( + client, + initial_response, + deserialization_cb, + AsyncLROBasePolling(0, lro_options={"final-state-via": "location"})) + result = await poll + assert result['location_result'] == True + + # Test 2, LRO options with Operation-Location final state + poll = async_poller( + client, + initial_response, + deserialization_cb, + AsyncLROBasePolling(0, lro_options={"final-state-via": "operation-location"})) + result = await poll + assert result['status'] == 'Succeeded' + + # Test 3, "do the right thing" and use Location by default + poll = async_poller( + client, + initial_response, + deserialization_cb, + AsyncLROBasePolling(0)) + result = await poll + assert result['location_result'] == True + + # Test 4, location has no body + + async def send(request, **kwargs): + assert request.method == 'GET' + + if request.url == 'http://example.org/location': + return TestBasePolling.mock_send( + http_request, + http_response, + 'GET', + 200, + body=None + ).http_response + elif request.url == 'http://example.org/async_monitor': + return TestBasePolling.mock_send( + http_request, + http_response, + 'GET', + 200, + body={'status': 'Succeeded'} + ).http_response + else: + pytest.fail("No other query allowed") + + client = async_pipeline_client_builder(send) + + poll = async_poller( + client, + initial_response, + deserialization_cb, + AsyncLROBasePolling(0, lro_options={"final-state-via": "location"})) + result = await poll + assert result is None + +@pytest.mark.asyncio +@pytest.mark.parametrize("http_request", HTTP_REQUESTS) +async def test_final_get_via_location(port, http_request, deserialization_cb): + client = AsyncTestRestClient(port) + request = http_request( + "PUT", + "http://localhost:{}/polling/polling-with-options".format(port), + ) + request.set_json_body({"hello": "world!"}) + initial_response = await client._client._pipeline.run(request) + poller = AsyncLROPoller( + client._client, + initial_response, + deserialization_cb, + AsyncLROBasePolling(0, lro_options={"final-state-via": "location"}), + ) + result = await poller.result() + assert result == {"returnedFrom": "locationHeaderUrl"} diff --git a/sdk/core/azure-core/tests/test_base_polling.py b/sdk/core/azure-core/tests/test_base_polling.py index ad030099c6eb..47f947ef63cb 100644 --- a/sdk/core/azure-core/tests/test_base_polling.py +++ b/sdk/core/azure-core/tests/test_base_polling.py @@ -31,6 +31,7 @@ import pickle import platform import six + try: from unittest import mock except ImportError: @@ -48,8 +49,9 @@ from azure.core.polling.base_polling import LROBasePolling from azure.core.pipeline.policies._utils import _FixedOffset -from utils import request_and_responses_product, REQUESTS_TRANSPORT_RESPONSES, create_transport_response +from utils import request_and_responses_product, REQUESTS_TRANSPORT_RESPONSES, create_transport_response, HTTP_REQUESTS from azure.core.pipeline._tools import is_rest +from rest_client import TestRestClient class SimpleResource: """An implementation of Python 3 SimpleNamespace. @@ -756,3 +758,124 @@ def test_long_running_negative(self, http_request, http_response): LOCATION_BODY = json.dumps({ 'name': TEST_NAME }) POLLING_STATUS = 200 + + @pytest.mark.parametrize("http_request,http_response", request_and_responses_product(REQUESTS_TRANSPORT_RESPONSES)) + def test_post_final_state_via(self, pipeline_client_builder, deserialization_cb, http_request, http_response): + # Test POST LRO with both Location and Operation-Location + CLIENT.http_request_type = http_request + CLIENT.http_response_type = http_response + # The initial response contains both Location and Operation-Location, a 202 and no Body + initial_response = TestBasePolling.mock_send( + http_request, + http_response, + 'POST', + 202, + { + 'location': 'http://example.org/location', + 'operation-location': 'http://example.org/async_monitor', + }, + '' + ) + + def send(request, **kwargs): + assert request.method == 'GET' + + if request.url == 'http://example.org/location': + return TestBasePolling.mock_send( + http_request, + http_response, + 'GET', + 200, + body={'location_result': True} + ).http_response + elif request.url == 'http://example.org/async_monitor': + return TestBasePolling.mock_send( + http_request, + http_response, + 'GET', + 200, + body={'status': 'Succeeded'} + ).http_response + else: + pytest.fail("No other query allowed") + + client = pipeline_client_builder(send) + + # Test 1, LRO options with Location final state + poll = LROPoller( + client, + initial_response, + deserialization_cb, + LROBasePolling(0, lro_options={"final-state-via": "location"})) + result = poll.result() + assert result['location_result'] == True + + # Test 2, LRO options with Operation-Location final state + poll = LROPoller( + client, + initial_response, + deserialization_cb, + LROBasePolling(0, lro_options={"final-state-via": "operation-location"})) + result = poll.result() + assert result['status'] == 'Succeeded' + + # Test 3, "do the right thing" and use Location by default + poll = LROPoller( + client, + initial_response, + deserialization_cb, + LROBasePolling(0)) + result = poll.result() + assert result['location_result'] == True + + # Test 4, location has no body + + def send(request, **kwargs): + assert request.method == 'GET' + + if request.url == 'http://example.org/location': + return TestBasePolling.mock_send( + http_request, + http_response, + 'GET', + 200, + body=None + ).http_response + elif request.url == 'http://example.org/async_monitor': + return TestBasePolling.mock_send( + http_request, + http_response, + 'GET', + 200, + body={'status': 'Succeeded'} + ).http_response + else: + pytest.fail("No other query allowed") + + client = pipeline_client_builder(send) + + poll = LROPoller( + client, + initial_response, + deserialization_cb, + LROBasePolling(0, lro_options={"final-state-via": "location"})) + result = poll.result() + assert result is None + +@pytest.mark.parametrize("http_request", HTTP_REQUESTS) +def test_final_get_via_location(port, http_request, deserialization_cb): + client = TestRestClient(port) + request = http_request( + "PUT", + "http://localhost:{}/polling/polling-with-options".format(port), + ) + request.set_json_body({"hello": "world!"}) + initial_response = client._client._pipeline.run(request) + poller = LROPoller( + client._client, + initial_response, + deserialization_cb, + LROBasePolling(0, lro_options={"final-state-via": "location"}), + ) + result = poller.result() + assert result == {"returnedFrom": "locationHeaderUrl"} diff --git a/sdk/core/azure-core/tests/testserver_tests/coretestserver/coretestserver/test_routes/polling.py b/sdk/core/azure-core/tests/testserver_tests/coretestserver/coretestserver/test_routes/polling.py index c433bfb4d04d..326f2dce63dc 100644 --- a/sdk/core/azure-core/tests/testserver_tests/coretestserver/coretestserver/test_routes/polling.py +++ b/sdk/core/azure-core/tests/testserver_tests/coretestserver/coretestserver/test_routes/polling.py @@ -151,3 +151,22 @@ def request_id_location(): '{"status": "Succeeded"}', status=200 ) + +@polling_api.route('/polling-with-options', methods=["PUT"]) +def polling_with_options_first(): + base_url = get_base_url(request) + return Response( + '{"properties":{"provisioningState": "InProgress"}}', + headers={ + 'location': '{}/polling/final-get-with-location'.format(base_url), + 'operation-location': '{}/polling/post/resource-location/operation-location-url'.format(base_url), + }, + status=202 + ) + +@polling_api.route('/final-get-with-location', methods=["GET"]) +def polling_with_options_final_get_with_location(): + return Response( + '{"returnedFrom": "locationHeaderUrl"}', + status=200 + )