This is an automated email from the ASF dual-hosted git repository. dabla pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 40a97a34ca70383de8f9315e912d6227ff2ee23d Author: David Blain <[email protected]> AuthorDate: Sun Oct 18 23:27:23 2026 +0200 Make the PowerBI refresh and MSGraph operators durable (#74675) * Make the PowerBI refresh and MSGraph operators resume after a retry A retry of the PowerBI refresh operator started a new dataset refresh while the first could still be running, and a triggerer restart before the first event did the same, because only the trigger knew the refresh id. A retry of the MSGraph operator sent its request again and fetched every page from the first one. Keeping the refresh id, the pages and the result of a non-GET request in the task state store lets a retry pick up where the failed try stopped. * Read the kept MSGraph pages only when the durable operator returns them An intermediate page read every page kept so far and then threw the list away, since it only adds itself to the task state store. A listing of N pages did about N^2/2 reads where N do. The kept pages are now read only on the last page, or when there is no response. * Warn once per run when a durable MSGraph task has no task state store The warning came from the store lookup, which every page does, so on Airflow < 3.3 a paginated request logged it once per page. It is now logged from execute only. --- .../microsoft/azure/docs/operators/msgraph.rst | 12 + .../microsoft/azure/docs/operators/powerbi.rst | 10 + providers/microsoft/azure/pyproject.toml | 2 +- .../providers/microsoft/azure/hooks/msgraph.py | 31 ++- .../providers/microsoft/azure/operators/msgraph.py | 158 ++++++++--- .../providers/microsoft/azure/operators/powerbi.py | 89 ++++++- .../providers/microsoft/azure/triggers/powerbi.py | 7 +- .../unit/microsoft/azure/hooks/test_msgraph.py | 29 +++ .../unit/microsoft/azure/operators/test_msgraph.py | 162 +++++++++++- .../unit/microsoft/azure/operators/test_powerbi.py | 288 +++++++++++++-------- .../unit/microsoft/azure/triggers/test_powerbi.py | 33 +++ 11 files changed, 666 insertions(+), 155 deletions(-) diff --git a/providers/microsoft/azure/docs/operators/msgraph.rst b/providers/microsoft/azure/docs/operators/msgraph.rst index 99815a959ad..25e2db76297 100644 --- a/providers/microsoft/azure/docs/operators/msgraph.rst +++ b/providers/microsoft/azure/docs/operators/msgraph.rst @@ -80,6 +80,18 @@ Below is an example of using this operator to create an item schedule in Fabric. :start-after: [START howto_operator_ms_fabric_create_item_schedule] :end-before: [END howto_operator_ms_fabric_create_item_schedule] +Durable execution +~~~~~~~~~~~~~~~~~ + +Set ``durable=True`` to make the operator :doc:`durable <apache-airflow:core-concepts/resumable-tasks>`. +It keeps every processed page and the link to the next one in the +:doc:`task state store <apache-airflow:core-concepts/task-state-store>`, so a retry continues at the page +it stopped at. For a request that is not a ``GET``, it also keeps the result, so a retry returns it instead +of sending the request again. + +Pages are stored as they are, so pages larger than ``[state_store] max_value_storage_bytes`` need a state +store backend that offloads large values. Durable execution requires Airflow 3.3 or newer. + Cross-References ---------------- diff --git a/providers/microsoft/azure/docs/operators/powerbi.rst b/providers/microsoft/azure/docs/operators/powerbi.rst index f02b7782493..eb03eff0c50 100644 --- a/providers/microsoft/azure/docs/operators/powerbi.rst +++ b/providers/microsoft/azure/docs/operators/powerbi.rst @@ -40,6 +40,16 @@ To trigger a refresh for the specified dataset from the specified workspace, use :start-after: [START howto_operator_powerbi_refresh_async] :end-before: [END howto_operator_powerbi_refresh_async] +Durable execution +~~~~~~~~~~~~~~~~~ + +The operator is :doc:`durable <apache-airflow:core-concepts/resumable-tasks>` by default: it starts the +refresh from the worker and keeps its id in the +:doc:`task state store <apache-airflow:core-concepts/task-state-store>` before it defers. A retry reconnects +to a refresh that is still in progress and does not start a new one once it has completed. A refresh that +failed, or that is no longer in the dataset's refresh history, is started again. Set ``durable=False`` to +start a new refresh on every try. Durable execution requires Airflow 3.3 or newer. + .. _howto/operator:PowerBIDatasetListOperator: PowerBIDatasetListOperator diff --git a/providers/microsoft/azure/pyproject.toml b/providers/microsoft/azure/pyproject.toml index 184b879ffc7..12d89dc4af0 100644 --- a/providers/microsoft/azure/pyproject.toml +++ b/providers/microsoft/azure/pyproject.toml @@ -59,7 +59,7 @@ requires-python = ">=3.11" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.21.0", + "apache-airflow-providers-common-compat>=1.21.0", # use next version "adlfs>=2026.5.0", "aiohttp>=3.14.0", "azure-batch>=15.0.0", diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py index 31abb53e1fc..adbf499e305 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -41,7 +41,7 @@ from io import BytesIO from json import JSONDecodeError from pathlib import Path from types import TracebackType -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, TypeVar, cast from urllib.parse import quote, urljoin, urlparse # Stays on httpx (not httpx2): the client below is handed to kiota_http, which builds and @@ -71,7 +71,7 @@ from airflow.providers.common.compat.connection import get_async_connection, get from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException, BaseHook, redact if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Coroutine, Iterable from azure.core.pipeline.transport._requests_basic import RequestsTransport from kiota_abstractions.authentication import BaseBearerTokenAuthenticationProvider @@ -86,6 +86,8 @@ if TYPE_CHECKING: PaginationCallable = Callable[..., tuple[str, dict[str, Any] | None]] +T = TypeVar("T") + def execute_callable(func: Callable, *args: Any, **kwargs: Any) -> Any: """Dynamically call a function by matching its signature to provided args/kwargs.""" @@ -771,6 +773,31 @@ class KiotaRequestAdapterHook(BaseHook): credential = cast("CachedAsyncTokenCredential", access_token_provider._credentials) await credential._credential.close() + def run_sync(self, coroutine: Coroutine[Any, Any, T]) -> T: + """ + Run a coroutine of this hook from synchronous code, such as an operator's ``execute``. + + :func:`asyncio.run` tears down the event loop that the cached request adapter is bound to, so + the adapter is closed while that loop still runs and the next call builds a new one. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + pass + else: + coroutine.close() + raise RuntimeError( + "run_sync cannot be called from a running event loop, await the coroutine instead." + ) + + async def run_and_close() -> T: + try: + return await coroutine + finally: + await self.close() + + return asyncio.run(run_and_close()) + def request_information( self, url: str, diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py index e60dd36bdc4..6b6d383e567 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py @@ -23,6 +23,8 @@ from contextlib import suppress from typing import ( TYPE_CHECKING, Any, + TypedDict, + cast, ) from airflow.exceptions import AirflowProviderDeprecationWarning @@ -39,6 +41,20 @@ if TYPE_CHECKING: from msgraph_core import APIVersion from airflow.sdk import Context + from airflow.sdk.execution_time.context import TaskStateStoreAccessor + +# Keys in the task state store of a durable MSGraphAsyncOperator. +CURSOR_KEY = "msgraph_cursor" +PAGE_KEY_PREFIX = "msgraph_page_" +COMPLETED_KEY = "msgraph_completed" + + +class PaginationCursor(TypedDict): + """Where a durable MSGraphAsyncOperator continues: the next link and the number of pages kept.""" + + url: str + query_parameters: dict[str, Any] | None + pages: int def default_event_handler(event: dict[Any, Any] | None = None, **context) -> Any: @@ -106,6 +122,13 @@ class MSGraphAsyncOperator(BaseOperator): the message from the event, otherwise the response from the event payload is returned. :param serializer: Class which handles response serialization (default is ResponseSerializer). Bytes will be base64 encoded into a string, so it can be stored as an XCom. + :param durable: Keep progress in the task state store, so that a retry, for instance after a worker + crash, resumes instead of starting over (default is False). Every processed page is kept with the + link to the next one, so the retry continues at the page it stopped at. The result of a request + that is not a GET is kept too, so the retry returns it instead of sending the request again. A + crash between a response and its write to the store still repeats that one request. Pages are + stored as they are, so pages larger than the store's ``max_value_storage_bytes`` need a state + store backend that offloads large values. Needs Airflow 3.3 or later. """ template_fields: Sequence[str] = ( @@ -140,6 +163,7 @@ class MSGraphAsyncOperator(BaseOperator): result_processor: Callable[[Any, Context], Any] = lambda result, **context: result, event_handler: Callable[[dict[Any, Any] | None, Context], Any] | None = None, serializer: type[ResponseSerializer] = ResponseSerializer, + durable: bool = False, **kwargs: Any, ): super().__init__(**kwargs) @@ -161,8 +185,26 @@ class MSGraphAsyncOperator(BaseOperator): self.result_processor = result_processor self.event_handler = event_handler or default_event_handler self.serializer: ResponseSerializer = serializer() + self.durable = durable + + def execute(self, context: Context) -> Any: + if task_state_store := self.get_task_state_store(context): + if completed := task_state_store.get(COMPLETED_KEY): + self.log.info("Returning the result of the request completed by an earlier try.") + return cast("dict[str, Any]", completed)["result"] + if cursor := self.get_cursor(task_state_store): + self.log.info("Resuming at page %s: %s", cursor["pages"] + 1, cursor["url"]) + self.query_parameters = cursor["query_parameters"] + self.defer_to_next_link( + url=cursor["url"], + query_parameters=cursor["query_parameters"], + method_name=self.execute_complete.__name__, + ) + elif self.durable: + self.log.warning( + "task_state_store not available in context, crash recovery is disabled for this run" + ) - def execute(self, context: Context) -> None: self.defer( trigger=MSGraphTrigger( url=self.url, @@ -215,7 +257,10 @@ class MSGraphAsyncOperator(BaseOperator): self.log.debug("response: %s", response) - results = self.pull_xcom(context=context) + task_state_store = self.get_task_state_store(context) + # A durable task reads its kept pages only once no next page defers: an intermediate + # page only adds itself to the store, see push_page. + results = [] if task_state_store else self.pull_xcom(context=context) if response: response = self.serializer.deserialize(response) @@ -236,21 +281,67 @@ class MSGraphAsyncOperator(BaseOperator): response=response, method_name=self.execute_complete.__name__, context=context ) except TaskDeferred as exception: - self.append_result( - results=results, - result=result, - append_result_as_list_if_absent=True, - ) - self.push_xcom(context=context, value=results) + if task_state_store: + self.push_page(task_state_store, result=result, deferred=exception) + else: + self.append_result( + results=results, + result=result, + append_result_as_list_if_absent=True, + ) + self.push_xcom(context=context, value=results) raise exception - if not results: - return result - - self.append_result(results=results, result=result) + if task_state_store: + results = self.pull_pages(task_state_store) + if results: + self.append_result(results=results, result=result) + else: + results = result + elif task_state_store: + results = self.pull_pages(task_state_store) + if task_state_store and self.method.upper() != "GET": + task_state_store.set(COMPLETED_KEY, {"result": results}) return results return None + def get_task_state_store(self, context: Context) -> TaskStateStoreAccessor | None: + """Return the task state store when this task is durable and the store is available.""" + return context.get("task_state_store") if self.durable else None + + @staticmethod + def get_cursor(task_state_store: TaskStateStoreAccessor) -> PaginationCursor | None: + return cast("PaginationCursor | None", task_state_store.get(CURSOR_KEY)) + + def pull_pages(self, task_state_store: TaskStateStoreAccessor) -> list[Any]: + """Return the results of the pages processed so far, kept in the task state store.""" + results: list[Any] = [] + cursor = self.get_cursor(task_state_store) + for page in range(cursor["pages"] if cursor else 0): + self.append_result(results=results, result=task_state_store.get(f"{PAGE_KEY_PREFIX}{page}")) + return results + + def push_page( + self, task_state_store: TaskStateStoreAccessor, result: Any, deferred: TaskDeferred + ) -> None: + """ + Keep the result of a processed page and the link to the next one in the task state store. + + The page goes in before the cursor that counts it, so a crash in between only fetches that page + again. + """ + cursor = self.get_cursor(task_state_store) + page = cursor["pages"] if cursor else 0 + task_state_store.set(f"{PAGE_KEY_PREFIX}{page}", result) + task_state_store.set( + CURSOR_KEY, + { + "url": deferred.trigger.url, + "query_parameters": (deferred.kwargs or {}).get("query_parameters"), + "pages": page + 1, + }, + ) + @classmethod def append_result( cls, @@ -344,23 +435,26 @@ class MSGraphAsyncOperator(BaseOperator): self.log.debug("query_parameters: %s", query_parameters) if url: - self.defer( - trigger=MSGraphTrigger( - url=url, - method=self.method, - path_parameters=self.path_parameters, - query_parameters=query_parameters, - headers=self.headers, - data=self.data, - response_type=self.response_type, - conn_id=self.conn_id, - timeout=self.timeout, - proxies=self.proxies, - scopes=self.scopes, - api_version=self.api_version, - serializer=type(self.serializer), - pagination_link=True, - ), - method_name=method_name, - kwargs={"query_parameters": query_parameters}, - ) + self.defer_to_next_link(url=url, query_parameters=query_parameters, method_name=method_name) + + def defer_to_next_link(self, url: str, query_parameters: dict[str, Any] | None, method_name: str) -> None: + self.defer( + trigger=MSGraphTrigger( + url=url, + method=self.method, + path_parameters=self.path_parameters, + query_parameters=query_parameters, + headers=self.headers, + data=self.data, + response_type=self.response_type, + conn_id=self.conn_id, + timeout=self.timeout, + proxies=self.proxies, + scopes=self.scopes, + api_version=self.api_version, + serializer=type(self.serializer), + pagination_link=True, + ), + method_name=method_name, + kwargs={"query_parameters": query_parameters}, + ) diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/powerbi.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/powerbi.py index ccbf025effe..01d416fe1bc 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/powerbi.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/powerbi.py @@ -24,9 +24,15 @@ from airflow.providers.common.compat.sdk import ( AirflowException, BaseOperator, BaseOperatorLink, + ResumableJobMixin, conf, ) -from airflow.providers.microsoft.azure.hooks.powerbi import PowerBIHook +from airflow.providers.microsoft.azure.hooks.powerbi import ( + PowerBIDatasetRefreshException, + PowerBIDatasetRefreshFields, + PowerBIDatasetRefreshStatus, + PowerBIHook, +) from airflow.providers.microsoft.azure.triggers.powerbi import ( PowerBIDatasetListTrigger, PowerBITrigger, @@ -35,10 +41,14 @@ from airflow.providers.microsoft.azure.triggers.powerbi import ( if TYPE_CHECKING: from msgraph_core import APIVersion + from pydantic import JsonValue from airflow.providers.common.compat.sdk import TaskInstanceKey from airflow.sdk import Context +# Status reported for a stored refresh id that is no longer in the dataset's refresh history. +_REFRESH_NOT_IN_HISTORY = "Not in refresh history" + class PowerBILink(BaseOperatorLink): """Construct a link to monitor a dataset in Power BI.""" @@ -55,10 +65,15 @@ class PowerBILink(BaseOperatorLink): return url -class PowerBIDatasetRefreshOperator(BaseOperator): +class PowerBIDatasetRefreshOperator(ResumableJobMixin, BaseOperator): """ Refreshes a Power BI dataset. + The refresh is started from the worker, and its id is kept in the task state store before the task + defers to wait for it. A retry, for instance after a worker crash, reconnects to that refresh while it + is still in progress, and does not start a new one once it has completed. This needs Airflow 3.3 or + later; on older versions every try starts a new refresh. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:PowerBIDatasetRefreshOperator` @@ -73,8 +88,12 @@ class PowerBIDatasetRefreshOperator(BaseOperator): :param wait_for_completion: If True, wait for the dataset refresh to complete. If False, trigger the refresh and return immediately without waiting. :param deferrable: This parameter is deprecated and no longer has any effect. The operator now always uses deferrable execution. + :param durable: Reconnect to the refresh started by an earlier try instead of starting a new one + (default is True). Needs Airflow 3.3 or later. """ + external_id_key = "powerbi_dataset_refresh_id" + template_fields: Sequence[str] = ( "dataset_id", "group_id", @@ -120,25 +139,85 @@ class PowerBIDatasetRefreshOperator(BaseOperator): def execute(self, context: Context): """Refresh the Power BI Dataset.""" + return self.execute_resumable(context) + + def submit_job(self, context: Context) -> str: + dataset_refresh_id = self.hook.run_sync( + self.hook.trigger_dataset_refresh( + dataset_id=self.dataset_id, + group_id=self.group_id, + request_body=self.request_body, + ) + ) + if not dataset_refresh_id: + raise PowerBIDatasetRefreshException("Failed to trigger the dataset refresh.") + self.log.info("Triggered dataset refresh %s", dataset_refresh_id) + return dataset_refresh_id + + def get_job_status(self, external_id: JsonValue, context: Context) -> str: + refresh_histories = self.hook.run_sync( + self.hook.get_refresh_history(dataset_id=self.dataset_id, group_id=self.group_id) + ) + for refresh_details in refresh_histories: + if refresh_details[PowerBIDatasetRefreshFields.REQUEST_ID.value] == external_id: + return refresh_details[PowerBIDatasetRefreshFields.STATUS.value] + # Power BI keeps a limited refresh history; a refresh that dropped out of it is not resumable. + return _REFRESH_NOT_IN_HISTORY + + def is_job_active(self, status: str) -> bool: + return status == PowerBIDatasetRefreshStatus.IN_PROGRESS + + def is_job_succeeded(self, status: str) -> bool: + return status == PowerBIDatasetRefreshStatus.COMPLETED + + def poll_until_complete(self, external_id: JsonValue, context: Context) -> None: + dataset_refresh_id = str(external_id) + self._push_refresh_xcom(context, dataset_refresh_id) + + if not self.wait_for_completion: + self.log.info("Triggered dataset refresh %s (fire-and-forget)", dataset_refresh_id) + return + self.defer( trigger=PowerBITrigger( conn_id=self.conn_id, group_id=self.group_id, dataset_id=self.dataset_id, + dataset_refresh_id=dataset_refresh_id, timeout=self.timeout, proxies=self.proxies, api_version=self.api_version, check_interval=self.check_interval, - wait_for_termination=False, - request_body=self.request_body, + wait_for_termination=True, + cancel_on_error=not self.durable, ), - method_name=self.handle_refresh.__name__, + method_name=self.execute_complete.__name__, ) + def get_job_result(self, external_id: JsonValue, context: Context) -> None: + # With wait_for_completion the mixin only gets here for a refresh an earlier try saw complete. + self._push_refresh_xcom( + context, + str(external_id), + PowerBIDatasetRefreshStatus.COMPLETED if self.wait_for_completion else None, + ) + + def _push_refresh_xcom( + self, context: Context, dataset_refresh_id: str, dataset_refresh_status: str | None = None + ) -> None: + context["ti"].xcom_push(key=f"{self.task_id}.powerbi_dataset_refresh_id", value=dataset_refresh_id) + if dataset_refresh_status: + context["ti"].xcom_push( + key=f"{self.task_id}.powerbi_dataset_refresh_status", value=dataset_refresh_status + ) + def handle_refresh(self, context: Context, event: dict[str, str] | None) -> None: """ Handle refresh-trigger event and optionally defer again to wait for refresh completion. + Only tasks deferred by an earlier version of this provider, whose trigger started the refresh, + still resume here. + :param context: Airflow context dictionary :param event: Event dict from trigger with status and dataset_refresh_id """ diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/powerbi.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/powerbi.py index 354e84b8553..241ee3993d8 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/powerbi.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/powerbi.py @@ -101,6 +101,8 @@ class PowerBITrigger(BasePowerBITrigger): :param check_interval: Time in seconds to wait between each poll. :param wait_for_termination: Wait for the dataset refresh to complete or fail. :param request_body: Additional arguments to pass to the request body, as described in https://learn.microsoft.com/en-us/rest/api/power-bi/datasets/refresh-dataset-in-group#request-body. + :param cancel_on_error: Cancel the dataset refresh when polling fails unexpectedly (default is True). + A durable operator passes False, so that its retry reconnects to the refresh instead. """ def __init__( @@ -115,6 +117,7 @@ class PowerBITrigger(BasePowerBITrigger): check_interval: int = 60, wait_for_termination: bool = True, request_body: dict[str, Any] | None = None, + cancel_on_error: bool = True, ): super().__init__(conn_id=conn_id, timeout=timeout, proxies=proxies, api_version=api_version) self.dataset_id = dataset_id @@ -123,6 +126,7 @@ class PowerBITrigger(BasePowerBITrigger): self.check_interval = check_interval self.wait_for_termination = wait_for_termination self.request_body = request_body + self.cancel_on_error = cancel_on_error def serialize(self): """Serialize the trigger instance.""" @@ -139,6 +143,7 @@ class PowerBITrigger(BasePowerBITrigger): "check_interval": self.check_interval, "wait_for_termination": self.wait_for_termination, "request_body": self.request_body, + "cancel_on_error": self.cancel_on_error, }, ) @@ -242,7 +247,7 @@ class PowerBITrigger(BasePowerBITrigger): ) return except Exception as error: - if self.dataset_refresh_id: + if self.dataset_refresh_id and self.cancel_on_error: try: self.log.info( "Unexpected error %s caught. Canceling dataset refresh %s", diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py index 4aa0fea5aca..27e4e0ca06e 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py @@ -827,6 +827,35 @@ class TestKiotaRequestAdapterHook: assert hook.conn_id not in hook.cached_request_adapters + def test_run_sync_returns_result_and_closes_cached_adapter(self): + """run_sync closes the adapter before asyncio.run tears down the loop it is bound to.""" + with patch_hook(): + hook = KiotaRequestAdapterHook(conn_id="msgraph_api") + + adapter = Mock(spec=HttpxRequestAdapter) + adapter._http_client = Mock(spec=AsyncClient, is_closed=False) + adapter._authentication_provider = mock_authentication_provider(closed=False) + hook.cached_request_adapters[hook.conn_id] = (hook.api_version, adapter) + + async def request(): + return "response" + + assert hook.run_sync(request()) == "response" + adapter._http_client.aclose.assert_awaited_once() + assert hook.conn_id not in hook.cached_request_adapters + + @pytest.mark.asyncio + async def test_run_sync_refuses_a_running_event_loop(self): + with patch_hook(): + hook = KiotaRequestAdapterHook(conn_id="msgraph_api") + request = AsyncMock(return_value="response") + coroutine = request() + + with pytest.raises(RuntimeError, match="cannot be called from a running event loop"): + hook.run_sync(coroutine) + + assert coroutine.cr_frame is None + def test_allowed_hosts_is_empty_list_when_not_configured(self): """An unset allowed_hosts/authority must yield [].""" actual = KiotaRequestAdapterHook.get_allowed_hosts(None, {}) diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py index e2d14a67ff9..772be3739df 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py @@ -27,8 +27,14 @@ from unittest import mock import pytest from airflow.exceptions import AirflowProviderDeprecationWarning -from airflow.providers.common.compat.sdk import AirflowException, Context, timezone -from airflow.providers.microsoft.azure.operators.msgraph import MSGraphAsyncOperator, execute_callable +from airflow.providers.common.compat.sdk import AirflowException, Context, TaskDeferred, timezone +from airflow.providers.microsoft.azure.operators.msgraph import ( + COMPLETED_KEY, + CURSOR_KEY, + PAGE_KEY_PREFIX, + MSGraphAsyncOperator, + execute_callable, +) from airflow.triggers.base import TriggerEvent from tests_common.test_utils.file_loading import load_file_from_resources, load_json_from_resources @@ -41,6 +47,36 @@ from unit.microsoft.azure.test_utils import ( ) +class TaskStateStore: + """In-memory stand-in for ``context["task_state_store"]``.""" + + def __init__(self, **values): + self.values = dict(values) + self.reads: list[str] = [] + + def get(self, key, default=None): + self.reads.append(key) + return self.values.get(key, default) + + def set(self, key, value, *, retention=None): + self.values[key] = value + + +def users_event(name: str) -> dict[str, Any]: + users = load_json_from_resources(dirname(__file__), "..", "resources", name) + return {"status": "success", "type": "builtins.dict", "response": json.dumps(users)} + + +def users_operator(**kwargs) -> MSGraphAsyncOperator: + return MSGraphAsyncOperator( + task_id="users_delta", + conn_id="msgraph_api", + url="users/delta()", + result_processor=lambda result, **context: result.get("value"), + **kwargs, + ) + + class TestMSGraphAsyncOperator: def test_execute_with_old_result_processor_signature(self): users = load_json_from_resources(dirname(__file__), "..", "resources", "users.json") @@ -510,3 +546,125 @@ class TestMSGraphAsyncOperator: == "response" ) assert len(recorded_warnings) == 0 + + +class TestDurableMSGraphAsyncOperator: + def test_keeps_the_page_and_the_link_to_the_next_one(self): + operator = users_operator(durable=True) + task_state_store = TaskStateStore() + context = {"ti": mock.MagicMock(), "task_state_store": task_state_store} + + with pytest.raises(TaskDeferred) as exc: + operator.execute_complete(context=context, event=users_event("users.json")) + + users = json.loads(users_event("users.json")["response"]) + assert task_state_store.values == { + f"{PAGE_KEY_PREFIX}0": users["value"], + CURSOR_KEY: {"url": users["@odata.nextLink"], "query_parameters": None, "pages": 1}, + } + assert exc.value.trigger.url == users["@odata.nextLink"] + context["ti"].xcom_push.assert_not_called() + + def test_intermediate_page_reads_only_the_cursor(self): + operator = users_operator(durable=True) + task_state_store = TaskStateStore( + **{f"{PAGE_KEY_PREFIX}{page}": [{"id": str(page)}] for page in range(3)}, + **{CURSOR_KEY: {"url": "users/delta()?$skiptoken=abc", "query_parameters": None, "pages": 3}}, + ) + context = {"ti": mock.MagicMock(), "task_state_store": task_state_store} + + with pytest.raises(TaskDeferred): + operator.execute_complete(context=context, event=users_event("users.json")) + + assert task_state_store.reads == [CURSOR_KEY] + assert task_state_store.values[CURSOR_KEY]["pages"] == 4 + + def test_execute_resumes_at_the_kept_link(self): + operator = users_operator(durable=True, query_parameters={"$top": 12}) + cursor = { + "url": "users/delta()?$skiptoken=abc", + "query_parameters": {"$top": 12, "$skip": 12}, + "pages": 1, + } + context = {"ti": mock.MagicMock(), "task_state_store": TaskStateStore(**{CURSOR_KEY: cursor})} + + with pytest.raises(TaskDeferred) as exc: + operator.execute(context) + + assert exc.value.trigger.url == cursor["url"] + assert exc.value.trigger.query_parameters == cursor["query_parameters"] + assert exc.value.trigger.pagination_link is True + assert exc.value.kwargs == {"query_parameters": cursor["query_parameters"]} + assert exc.value.method_name == "execute_complete" + + def test_last_page_returns_the_kept_pages_and_itself(self): + operator = users_operator(durable=True) + first_page = json.loads(users_event("users.json")["response"])["value"] + last_page = json.loads(users_event("next_users.json")["response"])["value"] + task_state_store = TaskStateStore( + **{ + f"{PAGE_KEY_PREFIX}0": first_page, + CURSOR_KEY: {"url": "users/delta()?$skiptoken=abc", "query_parameters": None, "pages": 1}, + } + ) + context = {"ti": mock.MagicMock(), "task_state_store": task_state_store} + + results = operator.execute_complete(context=context, event=users_event("next_users.json")) + + assert results == first_page + last_page + assert COMPLETED_KEY not in task_state_store.values + + def test_keeps_the_result_of_a_request_that_is_not_a_get(self): + operator = users_operator(durable=True, method="POST", data={"requestBody": "value"}) + task_state_store = TaskStateStore() + context = {"ti": mock.MagicMock(), "task_state_store": task_state_store} + + results = operator.execute_complete(context=context, event=users_event("next_users.json")) + + assert task_state_store.values == {COMPLETED_KEY: {"result": results}} + + def test_execute_returns_the_kept_result_without_sending_the_request_again(self): + operator = users_operator(durable=True, method="POST", data={"requestBody": "value"}) + context = { + "ti": mock.MagicMock(), + "task_state_store": TaskStateStore(**{COMPLETED_KEY: {"result": [{"id": "1"}]}}), + } + + with mock.patch.object(operator, "defer") as mock_defer: + assert operator.execute(context) == [{"id": "1"}] + + mock_defer.assert_not_called() + + def test_not_durable_leaves_the_task_state_store_alone(self): + operator = users_operator() + task_state_store = TaskStateStore() + context = mock_context(task=operator) + context["task_state_store"] = task_state_store + + with pytest.raises(TaskDeferred): + operator.execute_complete(context=context, event=users_event("users.json")) + + assert task_state_store.values == {} + + @mock.patch.object(MSGraphAsyncOperator, "log", new_callable=mock.PropertyMock) + def test_durable_without_task_state_store_keeps_pages_in_xcom(self, mock_log): + operator = users_operator(durable=True) + context = mock_context(task=operator) + + with pytest.raises(TaskDeferred): + operator.execute_complete(context=context, event=users_event("users.json")) + + assert context["ti"].xcom_pull(task_ids=operator.task_id, dag_id=operator.dag_id, key="return_value") + mock_log.return_value.warning.assert_not_called() + + @mock.patch.object(MSGraphAsyncOperator, "log", new_callable=mock.PropertyMock) + def test_durable_without_task_state_store_warns_in_execute(self, mock_log): + operator = users_operator(durable=True) + context = mock_context(task=operator) + + with pytest.raises(TaskDeferred): + operator.execute(context) + + mock_log.return_value.warning.assert_called_once_with( + "task_state_store not available in context, crash recovery is disabled for this run" + ) diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_powerbi.py b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_powerbi.py index 06f4dd5fb81..bda158c7329 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_powerbi.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_powerbi.py @@ -24,6 +24,7 @@ import pytest from airflow.providers.common.compat.sdk import AirflowException, BaseHook, TaskDeferred, timezone from airflow.providers.microsoft.azure.hooks.powerbi import ( + PowerBIDatasetRefreshException, PowerBIDatasetRefreshFields, PowerBIDatasetRefreshStatus, PowerBIHook, @@ -31,8 +32,8 @@ from airflow.providers.microsoft.azure.hooks.powerbi import ( from airflow.providers.microsoft.azure.operators.powerbi import PowerBIDatasetRefreshOperator from airflow.providers.microsoft.azure.triggers.powerbi import PowerBITrigger -from tests_common.test_utils.mock_context import mock_context -from tests_common.test_utils.operators.run_deferrable import execute_operator +from tests_common.test_utils.operators.run_deferrable import run_trigger +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS from unit.microsoft.azure.test_utils import get_airflow_connection DEFAULT_CONNECTION_CLIENT_SECRET = "powerbi_conn_id" @@ -56,6 +57,8 @@ CONFIG = { "request_body": REQUEST_BODY, } NEW_REFRESH_REQUEST_ID = "5e2d9921-e91b-491f-b7e1-e7d8db49194c" +PREVIOUS_REFRESH_REQUEST_ID = "0b4c1d2e-3f40-4a5b-8c6d-7e8f9a0b1c2d" +REFRESH_ID_KEY = "powerbi_dataset_refresh_id" SUCCESS_TRIGGER_EVENT = { "status": "success", @@ -92,48 +95,28 @@ IN_PROGRESS_REFRESH_DETAILS = { } -class TestPowerBIDatasetRefreshOperator: - @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection) - def test_execute_wait_for_completion_with_deferrable(self, connection): - operator = PowerBIDatasetRefreshOperator( - **CONFIG, - ) - context = mock_context(task=operator) - - with pytest.raises(TaskDeferred) as exc: - operator.execute(context) +def previous_refresh(status: str) -> dict[str, str]: + return { + PowerBIDatasetRefreshFields.REQUEST_ID.value: PREVIOUS_REFRESH_REQUEST_ID, + PowerBIDatasetRefreshFields.STATUS.value: status, + PowerBIDatasetRefreshFields.ERROR.value: "None", + } - assert isinstance(exc.value.trigger, PowerBITrigger) - assert exc.value.trigger.dataset_refresh_id is None - assert exc.value.trigger.wait_for_termination is False - assert exc.value.method_name == "handle_refresh" - @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection) - def test_powerbi_operator_async_get_refresh_status_success(self, connection): - """Test that execute first defers to trigger refresh when wait_for_completion=True""" - operator = PowerBIDatasetRefreshOperator( - **CONFIG, - wait_for_completion=True, # Explicitly set to True - ) - context = mock_context(task=operator) +class TaskStateStore: + """In-memory stand-in for ``context["task_state_store"]``.""" - with pytest.raises(TaskDeferred) as exc: - operator.execute(context) + def __init__(self, **values): + self.values = dict(values) - # Verify trigger is correct type - assert isinstance(exc.value.trigger, PowerBITrigger) + def get(self, key, default=None): + return self.values.get(key, default) - # Verify trigger has correct parameters - assert exc.value.trigger.dataset_id == DATASET_ID - assert exc.value.trigger.group_id == GROUP_ID - assert exc.value.trigger.wait_for_termination is False + def set(self, key, value, *, retention=None): + self.values[key] = value - # Verify callback method name - assert exc.value.method_name == "handle_refresh" - - # Verify dataset_refresh_id is None (trigger will create it) - assert exc.value.trigger.dataset_refresh_id is None +class TestPowerBIDatasetRefreshOperator: @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection) def test_handle_refresh_wait_for_completion(self, connection): """Test that handle_refresh defers to execute_complete when wait_for_completion=True.""" @@ -193,52 +176,6 @@ class TestPowerBIDatasetRefreshOperator: assert context["ti"].xcom_push.call_count == 0 - @mock.patch.object(PowerBIHook, "get_refresh_details_by_refresh_id") - @mock.patch.object(PowerBIHook, "trigger_dataset_refresh") - def test_execute_operator_wait_for_completion_full_lifecycle( - self, mock_trigger_dataset_refresh, mock_get_refresh_details_by_refresh_id - ): - """Assert the full deferrable lifecycle completes successfully.""" - mock_trigger_dataset_refresh.return_value = NEW_REFRESH_REQUEST_ID - mock_get_refresh_details_by_refresh_id.side_effect = [ - { - PowerBIDatasetRefreshFields.STATUS.value: PowerBIDatasetRefreshStatus.IN_PROGRESS, - PowerBIDatasetRefreshFields.ERROR.value: None, - }, - { - PowerBIDatasetRefreshFields.STATUS.value: PowerBIDatasetRefreshStatus.COMPLETED, - PowerBIDatasetRefreshFields.ERROR.value: None, - }, - ] - - operator = PowerBIDatasetRefreshOperator( - **CONFIG, - wait_for_completion=True, - ) - - result, events = execute_operator(operator) - - assert result is None - mock_trigger_dataset_refresh.assert_called_once_with( - dataset_id=DATASET_ID, - group_id=GROUP_ID, - request_body=REQUEST_BODY, - ) - assert mock_get_refresh_details_by_refresh_id.call_count == 2 - assert len(events) == 2 - assert events[0].payload == { - "status": "success", - "dataset_refresh_status": None, - "message": f"The dataset refresh {NEW_REFRESH_REQUEST_ID} has been triggered.", - "dataset_refresh_id": NEW_REFRESH_REQUEST_ID, - } - assert events[1].payload == { - "status": "success", - "dataset_refresh_status": PowerBIDatasetRefreshStatus.COMPLETED, - "message": f"The dataset refresh {NEW_REFRESH_REQUEST_ID} has {PowerBIDatasetRefreshStatus.COMPLETED}.", - "dataset_refresh_id": NEW_REFRESH_REQUEST_ID, - } - def test_powerbi_operator_async_execute_complete_success(self): """Assert that execute_complete processes success event correctly""" operator = PowerBIDatasetRefreshOperator(**CONFIG) @@ -335,44 +272,171 @@ class TestPowerBIDatasetRefreshOperator: assert url == EXPECTED_ITEM_RUN_OP_EXTRA_LINK - @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection) - def test_execute_fire_and_forget_mode(self, mock_connection): - """Test fire-and-forget mode defers to trigger refresh.""" - operator = PowerBIDatasetRefreshOperator( - **CONFIG, - wait_for_completion=False, - ) - context = mock_context(task=operator) + @mock.patch.object(PowerBIHook, "trigger_dataset_refresh", return_value=NEW_REFRESH_REQUEST_ID) + def test_execute_starts_refresh_on_worker_and_defers_to_poll(self, mock_trigger_dataset_refresh): + operator = PowerBIDatasetRefreshOperator(**CONFIG) + task_state_store = TaskStateStore() + context = {"ti": MagicMock(), "task_state_store": task_state_store} with pytest.raises(TaskDeferred) as exc: operator.execute(context) + mock_trigger_dataset_refresh.assert_called_once_with( + dataset_id=DATASET_ID, group_id=GROUP_ID, request_body=REQUEST_BODY + ) assert isinstance(exc.value.trigger, PowerBITrigger) - assert exc.value.trigger.wait_for_termination is False - assert exc.value.method_name == "handle_refresh" + assert exc.value.trigger.dataset_refresh_id == NEW_REFRESH_REQUEST_ID + assert exc.value.trigger.wait_for_termination is True + assert exc.value.method_name == "execute_complete" + context["ti"].xcom_push.assert_called_once_with( + key=f"{TASK_ID}.powerbi_dataset_refresh_id", value=NEW_REFRESH_REQUEST_ID + ) + if AIRFLOW_V_3_3_PLUS: + assert task_state_store.values == {REFRESH_ID_KEY: NEW_REFRESH_REQUEST_ID} + assert exc.value.trigger.cancel_on_error is False - @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection) - def test_execute_default_behavior_waits_for_completion(self, mock_connection): - """Test that default behavior (wait_for_completion=True) first defers to trigger refresh.""" - config_without_wait = { - "task_id": TASK_ID, - "conn_id": DEFAULT_CONNECTION_CLIENT_SECRET, - "group_id": GROUP_ID, - "dataset_id": DATASET_ID, - "request_body": REQUEST_BODY, - "check_interval": 1, - "timeout": 3, - # Deliberately exclude wait_for_completion - should default to True + @mock.patch.object(PowerBIHook, "trigger_dataset_refresh", return_value=NEW_REFRESH_REQUEST_ID) + def test_execute_fire_and_forget_mode(self, mock_trigger_dataset_refresh): + """Fire-and-forget starts the refresh and returns without deferring.""" + operator = PowerBIDatasetRefreshOperator(**CONFIG, wait_for_completion=False) + context = {"ti": MagicMock(), "task_state_store": TaskStateStore()} + + assert operator.execute(context) is None + + mock_trigger_dataset_refresh.assert_called_once() + context["ti"].xcom_push.assert_called_with( + key=f"{TASK_ID}.powerbi_dataset_refresh_id", value=NEW_REFRESH_REQUEST_ID + ) + + @mock.patch.object(PowerBIHook, "trigger_dataset_refresh", return_value=None) + def test_execute_raises_when_refresh_has_no_id(self, mock_trigger_dataset_refresh): + operator = PowerBIDatasetRefreshOperator(**CONFIG) + task_state_store = TaskStateStore() + context = {"ti": MagicMock(), "task_state_store": task_state_store} + + with pytest.raises(PowerBIDatasetRefreshException, match="Failed to trigger the dataset refresh."): + operator.execute(context) + + assert task_state_store.values == {} + + @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The task state store arrived in Airflow 3.3") + @mock.patch.object(PowerBIHook, "get_refresh_history") + @mock.patch.object(PowerBIHook, "trigger_dataset_refresh") + def test_execute_reconnects_to_refresh_in_progress( + self, mock_trigger_dataset_refresh, mock_get_refresh_history + ): + mock_get_refresh_history.return_value = [previous_refresh(PowerBIDatasetRefreshStatus.IN_PROGRESS)] + operator = PowerBIDatasetRefreshOperator(**CONFIG) + context = { + "ti": MagicMock(), + "task_state_store": TaskStateStore(**{REFRESH_ID_KEY: PREVIOUS_REFRESH_REQUEST_ID}), } - operator = PowerBIDatasetRefreshOperator(**config_without_wait) - context = mock_context(task=operator) + with pytest.raises(TaskDeferred) as exc: + operator.execute(context) + + mock_trigger_dataset_refresh.assert_not_called() + assert exc.value.trigger.dataset_refresh_id == PREVIOUS_REFRESH_REQUEST_ID + assert exc.value.method_name == "execute_complete" + + @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The task state store arrived in Airflow 3.3") + @mock.patch.object(PowerBIHook, "get_refresh_history") + @mock.patch.object(PowerBIHook, "trigger_dataset_refresh") + def test_execute_does_not_refresh_again_once_completed( + self, mock_trigger_dataset_refresh, mock_get_refresh_history + ): + mock_get_refresh_history.return_value = [previous_refresh(PowerBIDatasetRefreshStatus.COMPLETED)] + operator = PowerBIDatasetRefreshOperator(**CONFIG) + context = { + "ti": MagicMock(), + "task_state_store": TaskStateStore(**{REFRESH_ID_KEY: PREVIOUS_REFRESH_REQUEST_ID}), + } + + assert operator.execute(context) is None + + mock_trigger_dataset_refresh.assert_not_called() + xcom_data = { + call.kwargs["key"]: call.kwargs["value"] for call in context["ti"].xcom_push.call_args_list + } + assert xcom_data == { + f"{TASK_ID}.powerbi_dataset_refresh_id": PREVIOUS_REFRESH_REQUEST_ID, + f"{TASK_ID}.powerbi_dataset_refresh_status": PowerBIDatasetRefreshStatus.COMPLETED, + } + + @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The task state store arrived in Airflow 3.3") + @pytest.mark.parametrize( + "refresh_history", + [ + pytest.param([previous_refresh(PowerBIDatasetRefreshStatus.FAILED)], id="failed"), + pytest.param([previous_refresh(PowerBIDatasetRefreshStatus.DISABLED)], id="disabled"), + pytest.param([], id="no-longer-in-history"), + ], + ) + @mock.patch.object(PowerBIHook, "get_refresh_history") + @mock.patch.object(PowerBIHook, "trigger_dataset_refresh", return_value=NEW_REFRESH_REQUEST_ID) + def test_execute_starts_new_refresh_when_previous_cannot_be_resumed( + self, mock_trigger_dataset_refresh, mock_get_refresh_history, refresh_history + ): + mock_get_refresh_history.return_value = refresh_history + operator = PowerBIDatasetRefreshOperator(**CONFIG) + task_state_store = TaskStateStore(**{REFRESH_ID_KEY: PREVIOUS_REFRESH_REQUEST_ID}) + context = {"ti": MagicMock(), "task_state_store": task_state_store} - # Should defer for initial trigger step (because default is wait_for_completion=True) with pytest.raises(TaskDeferred) as exc: operator.execute(context) - # Verify it deferred with correct trigger - assert isinstance(exc.value.trigger, PowerBITrigger) - assert exc.value.trigger.wait_for_termination is False - assert exc.value.method_name == "handle_refresh" + mock_trigger_dataset_refresh.assert_called_once() + assert exc.value.trigger.dataset_refresh_id == NEW_REFRESH_REQUEST_ID + assert task_state_store.values == {REFRESH_ID_KEY: NEW_REFRESH_REQUEST_ID} + + @mock.patch.object(PowerBIHook, "trigger_dataset_refresh", return_value=NEW_REFRESH_REQUEST_ID) + def test_execute_not_durable_starts_new_refresh_and_cancels_on_error(self, mock_trigger_dataset_refresh): + operator = PowerBIDatasetRefreshOperator(**CONFIG, durable=False) + task_state_store = TaskStateStore(**{REFRESH_ID_KEY: PREVIOUS_REFRESH_REQUEST_ID}) + context = {"ti": MagicMock(), "task_state_store": task_state_store} + + with pytest.raises(TaskDeferred) as exc: + operator.execute(context) + + mock_trigger_dataset_refresh.assert_called_once() + assert exc.value.trigger.dataset_refresh_id == NEW_REFRESH_REQUEST_ID + assert exc.value.trigger.cancel_on_error is True + assert task_state_store.values == {REFRESH_ID_KEY: PREVIOUS_REFRESH_REQUEST_ID} + + @mock.patch.object(PowerBIHook, "get_refresh_details_by_refresh_id") + @mock.patch.object(PowerBIHook, "trigger_dataset_refresh", return_value=NEW_REFRESH_REQUEST_ID) + def test_execute_operator_wait_for_completion_full_lifecycle( + self, mock_trigger_dataset_refresh, mock_get_refresh_details_by_refresh_id + ): + """Assert the full lifecycle: start on the worker, poll in the trigger, complete on the worker.""" + mock_get_refresh_details_by_refresh_id.side_effect = [ + { + PowerBIDatasetRefreshFields.STATUS.value: PowerBIDatasetRefreshStatus.IN_PROGRESS, + PowerBIDatasetRefreshFields.ERROR.value: None, + }, + { + PowerBIDatasetRefreshFields.STATUS.value: PowerBIDatasetRefreshStatus.COMPLETED, + PowerBIDatasetRefreshFields.ERROR.value: None, + }, + ] + operator = PowerBIDatasetRefreshOperator(**CONFIG) + + context = {"ti": MagicMock(), "task_state_store": TaskStateStore()} + + with pytest.raises(TaskDeferred) as exc: + operator.execute(context) + events = run_trigger(exc.value.trigger) + + assert getattr(operator, exc.value.method_name)(context=context, event=events[0].payload) is None + mock_trigger_dataset_refresh.assert_called_once_with( + dataset_id=DATASET_ID, group_id=GROUP_ID, request_body=REQUEST_BODY + ) + assert mock_get_refresh_details_by_refresh_id.call_count == 2 + assert [event.payload for event in events] == [ + { + "status": "success", + "dataset_refresh_status": PowerBIDatasetRefreshStatus.COMPLETED, + "message": f"The dataset refresh {NEW_REFRESH_REQUEST_ID} has {PowerBIDatasetRefreshStatus.COMPLETED}.", + "dataset_refresh_id": NEW_REFRESH_REQUEST_ID, + } + ] diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_powerbi.py b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_powerbi.py index a648f1c2f9b..d462ec9e6d6 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_powerbi.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_powerbi.py @@ -126,6 +126,7 @@ class TestPowerBITrigger: "check_interval": CHECK_INTERVAL, "wait_for_termination": True, "request_body": REQUEST_BODY, + "cancel_on_error": True, } @pytest.mark.asyncio @@ -247,6 +248,38 @@ class TestPowerBITrigger: assert response in task mock_cancel_dataset_refresh.assert_called_once() + @pytest.mark.asyncio + @mock.patch.object(PowerBIHook, "cancel_dataset_refresh") + @mock.patch.object(PowerBIHook, "get_refresh_details_by_refresh_id") + async def test_powerbi_trigger_run_exception_does_not_cancel_when_cancel_on_error_is_false( + self, mock_get_refresh_details_by_refresh_id, mock_cancel_dataset_refresh + ): + """A durable operator keeps the refresh running, so that its retry can reconnect to it.""" + mock_get_refresh_details_by_refresh_id.side_effect = Exception("Test exception") + powerbi_trigger = PowerBITrigger( + conn_id=POWERBI_CONN_ID, + dataset_id=DATASET_ID, + dataset_refresh_id=DATASET_REFRESH_ID, + group_id=GROUP_ID, + check_interval=CHECK_INTERVAL, + timeout=TIMEOUT, + cancel_on_error=False, + ) + + events = [event async for event in powerbi_trigger.run()] + + assert events == [ + TriggerEvent( + { + "status": "error", + "dataset_refresh_status": None, + "message": "An error occurred: Test exception", + "dataset_refresh_id": DATASET_REFRESH_ID, + } + ) + ] + mock_cancel_dataset_refresh.assert_not_called() + @pytest.mark.asyncio @mock.patch.object(PowerBIHook, "cancel_dataset_refresh") @mock.patch.object(PowerBIHook, "get_refresh_details_by_refresh_id")
