This is an automated email from the ASF dual-hosted git repository.
hussein-awala pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1040a319e0e Fix the durable MSGraph and PowerBI operators from
post-merge review (#75272)
1040a319e0e is described below
commit 1040a319e0eb6a87a9ce1d1133c13d09e1dca8cc
Author: David Blain <[email protected]>
AuthorDate: Thu Oct 29 11:46:10 2026 +0100
Fix the durable MSGraph and PowerBI operators from post-merge review
(#75272)
* Fix the durable MSGraph and PowerBI operators from post-merge review
Follow-up to #74675, addressing the post-merge review
(https://github.com/apache/airflow/pull/74675#pullrequestreview-5596928776).
MSGraphAsyncOperator (durable):
- A page is stored as {"result": ...}, so a page whose result_processor
returned None no longer fails
the task, and a page gone from the store (expired or dropped) fails the
try with
MSGraphPageNotInTaskStateStoreError and clears the store, so the next try
starts from the first page.
- The pages and the cursor are removed once the result is returned. A
request that is not a GET keeps
only its result, written before the pages are removed.
- durable=True below Airflow 3.3 warns at parse time and is turned off,
like ResumableJobMixin.
- The store keys and the cursor type are private, and a dead assignment is
gone.
PowerBIDatasetRefreshOperator:
- Fire-and-forget pushes the refresh id once. A retry that finds the
refresh in progress counts it as
done, since the task only had to start it.
- The changelog and the operator docs say that the refresh now starts from
the worker and that a
durable refresh is not cancelled when polling fails.
* Remove the MSGraph operator's own keys from the task state store instead
of clearing it
`task_state_store.clear()` deletes every key of the task instance, so a key
a
`result_processor` or a subclass wrote was lost when a durable GET finished
or a
page was missing. The pages and the cursor are now deleted by name in both
places, and the tests keep a key of their own to show it survives. A page
kept
unwrapped by 15.3.0rc1 is treated like a missing page, so a retry after the
upgrade starts again from the first page instead of failing on every try.
The changelog warning is dropped: the 15.3.0 note already says it. The
comment
on `is_job_succeeded` says that counting an in-progress fire-and-forget
refresh
as done is logged and metered by the mixin as `already_succeeded`, and why.
* Let the trigger start the Power BI refresh when the operator is not
durable
With `durable=False` the operator defers to `PowerBITrigger` without a
refresh
id, as it did before durable execution, so only the triggerer talks to
Power BI
and every try starts a new refresh. `durable=False` thereby restores the
whole
earlier behaviour, and deployments below Airflow 3.3, where `durable` is
forced
off, keep it without any change. The worker starts the refresh only when the
operator is durable, which the docs and the docstrings now say.
`test_execute_not_durable_starts_new_refresh_and_cancels_on_error` is
replaced by
a test of the deferral and a full lifecycle test of the non-durable path;
the
tests that start the refresh on the worker are now limited to Airflow 3.3+.
* Update
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/powerbi.py
Co-authored-by: Amogh Desai <[email protected]>
* Update
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/powerbi.py
Co-authored-by: Amogh Desai <[email protected]>
* Update
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
Co-authored-by: Amogh Desai <[email protected]>
* Update
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
Co-authored-by: Amogh Desai <[email protected]>
* Update
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
Co-authored-by: Amogh Desai <[email protected]>
* Update
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
Co-authored-by: Amogh Desai <[email protected]>
* Update
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
Co-authored-by: Amogh Desai <[email protected]>
* Update
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
Co-authored-by: Amogh Desai <[email protected]>
* fix: Version the pages a durable MSGraphAsyncOperator keeps in the task
state store
An unwrapped page kept by 15.3.0rc1 whose result is a dict with a "result"
key
passed for a wrapped one and was unwrapped on resume, so an upgrade could
drop
the result's structure instead of starting the pagination again. The cursor
now
carries a version; pages under a cursor of another version are removed and
the
next try starts again from the first page, on both resume points
(execute_complete through push_page or pull_pages). execute removes them and
starts again in the same try.
---------
Co-authored-by: Amogh Desai <[email protected]>
---
.../microsoft/azure/docs/operators/msgraph.rst | 5 +-
.../microsoft/azure/docs/operators/powerbi.rst | 14 +-
.../providers/microsoft/azure/exceptions.py | 9 ++
.../providers/microsoft/azure/operators/msgraph.py | 124 ++++++++++++++----
.../providers/microsoft/azure/operators/powerbi.py | 53 ++++++--
.../providers/microsoft/azure/version_compat.py | 2 +
.../unit/microsoft/azure/operators/test_msgraph.py | 145 ++++++++++++++++++---
.../unit/microsoft/azure/operators/test_powerbi.py | 87 ++++++++++++-
8 files changed, 376 insertions(+), 63 deletions(-)
diff --git a/providers/microsoft/azure/docs/operators/msgraph.rst
b/providers/microsoft/azure/docs/operators/msgraph.rst
index 25e2db76297..875d3df4772 100644
--- a/providers/microsoft/azure/docs/operators/msgraph.rst
+++ b/providers/microsoft/azure/docs/operators/msgraph.rst
@@ -87,7 +87,10 @@ Set ``durable=True`` to make the operator :doc:`durable
<apache-airflow:core-con
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.
+of sending the request again. The pages and the link are removed from the
store once the result is
+returned. A page that is no longer in the store, for instance because it
expired under
+``[state_store] default_retention_days``, fails the try and removes the pages
and the link, so the next try starts again
+from the first page.
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.
diff --git a/providers/microsoft/azure/docs/operators/powerbi.rst
b/providers/microsoft/azure/docs/operators/powerbi.rst
index eb03eff0c50..453a6e94191 100644
--- a/providers/microsoft/azure/docs/operators/powerbi.rst
+++ b/providers/microsoft/azure/docs/operators/powerbi.rst
@@ -47,8 +47,18 @@ The operator is :doc:`durable
<apache-airflow:core-concepts/resumable-tasks>` by
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.
+failed, or that is no longer in the dataset's refresh history, is started
again. Durable execution
+requires Airflow 3.3 or newer.
+
+Two things differ from earlier versions of the provider when the operator is
durable:
+
+* The refresh is started by the worker, so the workers need network access to
``api.powerbi.com`` too.
+ The triggerer still polls the refresh.
+* The refresh is not cancelled when polling it fails, so the retry can
reconnect to it.
+
+Set ``durable=False`` to keep the earlier behaviour: the trigger starts the
refresh, so only the triggerer
+talks to Power BI, every try starts a new refresh, and it is cancelled when
polling fails. Below Airflow
+3.3 the operator behaves as with ``durable=False``.
.. _howto/operator:PowerBIDatasetListOperator:
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
index 461463e100c..a6d46897ba5 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
@@ -62,3 +62,12 @@ class
AzureBatchRequiredParametersMissingError(AirflowException):
class AzureAISearchIndexingError(AirflowException):
"""Raised when Azure AI Search rejects some documents of a write."""
+
+
+class MSGraphPageNotInTaskStateStoreError(AirflowException):
+ """
+ Raised when the pages kept by a durable MSGraphAsyncOperator cannot be
read back from the task state store.
+
+ A page is gone, or an earlier version of the provider kept the pages; the
next try starts again from the
+ first page.
+ """
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 6b6d383e567..59b6da91dab 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
@@ -29,11 +29,13 @@ from typing import (
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.compat.sdk import XCOM_RETURN_KEY,
AirflowException, BaseOperator, TaskDeferred
+from airflow.providers.microsoft.azure.exceptions import
MSGraphPageNotInTaskStateStoreError
from airflow.providers.microsoft.azure.hooks.msgraph import
KiotaRequestAdapterHook
from airflow.providers.microsoft.azure.triggers.msgraph import (
MSGraphTrigger,
ResponseSerializer,
)
+from airflow.providers.microsoft.azure.version_compat import AIRFLOW_V_3_3_PLUS
if TYPE_CHECKING:
from io import BytesIO
@@ -44,14 +46,18 @@ if TYPE_CHECKING:
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"
+_CURSOR_KEY = "msgraph_cursor"
+_PAGE_KEY_PREFIX = "msgraph_page_"
+_COMPLETED_KEY = "msgraph_completed"
+# The version of the kept pages and their cursor. 15.3.0rc1 kept the pages
unwrapped, under a cursor without
+# a version; a cursor of another version is not read back, its pages go and
the pagination starts again.
+_CURSOR_VERSION = 1
-class PaginationCursor(TypedDict):
+class _PaginationCursor(TypedDict):
"""Where a durable MSGraphAsyncOperator continues: the next link and the
number of pages kept."""
+ version: int
url: str
query_parameters: dict[str, Any] | None
pages: int
@@ -126,9 +132,10 @@ class MSGraphAsyncOperator(BaseOperator):
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.
+ crash between a response and its write to the store still repeats that
one request. The pages
+ are removed from the store once the result is returned. 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] = (
@@ -167,6 +174,9 @@ class MSGraphAsyncOperator(BaseOperator):
**kwargs: Any,
):
super().__init__(**kwargs)
+ if durable and not AIRFLOW_V_3_3_PLUS:
+ warnings.warn("`durable` has no effect on Airflow versions below
3.3.", UserWarning, stacklevel=2)
+ durable = False
self.url = url
self.response_type = response_type
self.path_parameters = path_parameters
@@ -189,17 +199,23 @@ class MSGraphAsyncOperator(BaseOperator):
def execute(self, context: Context) -> Any:
if task_state_store := self.get_task_state_store(context):
- if completed := task_state_store.get(COMPLETED_KEY):
+ 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__,
- )
+ if cursor.get("version") != _CURSOR_VERSION:
+ self.log.info(
+ "Removing the pages kept by an earlier version of the
provider, "
+ "starting again from the first page."
+ )
+ self._remove_pages(task_state_store, cursor)
+ else:
+ self.log.info("Resuming at page %s: %s", cursor["pages"] +
1, cursor["url"])
+ 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"
@@ -300,8 +316,8 @@ class MSGraphAsyncOperator(BaseOperator):
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})
+ if task_state_store:
+ self._clean_up(task_state_store, results=results)
return results
return None
@@ -310,15 +326,48 @@ class MSGraphAsyncOperator(BaseOperator):
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 get_cursor(task_state_store: TaskStateStoreAccessor) ->
_PaginationCursor | None:
+ return cast("_PaginationCursor | None",
task_state_store.get(_CURSOR_KEY))
+
+ def _get_cursor_of_this_version(
+ self, task_state_store: TaskStateStoreAccessor
+ ) -> _PaginationCursor | None:
+ """
+ Return the cursor the kept pages can be read back with, or none when
no page is kept.
+
+ The pages kept by an earlier version of the provider are not read
back: an unwrapped page of
+ 15.3.0rc1 that happens to be a dict with a ``result`` key would pass
for a wrapped one. They go,
+ with their cursor, so the next try starts again from the first page.
+ """
+ cursor = self.get_cursor(task_state_store)
+ if cursor is None or cursor.get("version") == _CURSOR_VERSION:
+ return cursor
+ self._remove_pages(task_state_store, cursor)
+ raise MSGraphPageNotInTaskStateStoreError(
+ "The pages were kept by an earlier version of the provider, so the
next try starts again from "
+ "the first page."
+ )
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."""
+ """
+ Return the results of the pages processed so far, kept in the task
state store.
+
+ A page is kept wrapped in a dict, so a page whose result is ``None``
is told apart from a page
+ that is gone, for instance because it expired under ``[state_store]
default_retention_days``.
+ A missing page removes the pages and the cursor, so the next try
starts again from the first page
+ instead of returning results with a hole in them.
+ """
results: list[Any] = []
- cursor = self.get_cursor(task_state_store)
+ cursor = self._get_cursor_of_this_version(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}"))
+ kept = task_state_store.get(f"{_PAGE_KEY_PREFIX}{page}")
+ if not isinstance(kept, dict) or "result" not in kept:
+ self._remove_pages(task_state_store, cursor)
+ raise MSGraphPageNotInTaskStateStoreError(
+ f"Page {page + 1} is no longer in the task state store, so
the next try starts again from "
+ "the first page."
+ )
+ self.append_result(results=results, result=cast("dict[str, Any]",
kept)["result"])
return results
def push_page(
@@ -330,18 +379,43 @@ class MSGraphAsyncOperator(BaseOperator):
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)
+ cursor = self._get_cursor_of_this_version(task_state_store)
page = cursor["pages"] if cursor else 0
- task_state_store.set(f"{PAGE_KEY_PREFIX}{page}", result)
+ task_state_store.set(f"{_PAGE_KEY_PREFIX}{page}", {"result": result})
task_state_store.set(
- CURSOR_KEY,
+ _CURSOR_KEY,
{
+ "version": _CURSOR_VERSION,
"url": deferred.trigger.url,
"query_parameters": (deferred.kwargs or
{}).get("query_parameters"),
"pages": page + 1,
},
)
+ def _clean_up(self, task_state_store: TaskStateStoreAccessor, results:
Any) -> None:
+ """
+ Remove the pages and the cursor from the task state store once the
result is returned.
+
+ A non GET result is written before the pages go, so a crash in between
does not resend the request.
+ """
+ if self.method.upper() != "GET":
+ task_state_store.set(_COMPLETED_KEY, {"result": results})
+ self._remove_pages(task_state_store, self.get_cursor(task_state_store))
+
+ @staticmethod
+ def _remove_pages(task_state_store: TaskStateStoreAccessor, cursor:
_PaginationCursor | None) -> None:
+ """
+ Delete the pages and the cursor by their keys.
+
+ ``task_state_store.clear()`` would delete every key of the task
instance, including keys a
+ ``result_processor`` or a subclass wrote.
+ """
+ if cursor is None:
+ return
+ for page in range(cursor["pages"]):
+ task_state_store.delete(f"{_PAGE_KEY_PREFIX}{page}")
+ task_state_store.delete(_CURSOR_KEY)
+
@classmethod
def append_result(
cls,
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 8ae2b979389..aa45d84e653 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
@@ -69,10 +69,11 @@ 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.
+ When durable, 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. Otherwise the trigger starts the refresh, as in
earlier versions of this
+ provider: only the triggerer talks to Power BI, and every try starts a new
refresh.
.. seealso::
For more information on how to use this operator, take a look at the
guide:
@@ -88,8 +89,11 @@ class PowerBIDatasetRefreshOperator(ResumableJobMixin,
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.
+ :param durable: Start the refresh from the worker and reconnect to it on a
retry instead of starting
+ a new one (default is True). A durable refresh is not cancelled when
polling it fails, so that
+ the retry can reconnect to it. Set it to False to keep the behaviour
of earlier versions: the
+ trigger starts the refresh, every try starts a new one, and it is
cancelled when polling fails.
+ Needs Airflow 3.3 or later; below it the operator behaves as with
``durable=False``.
"""
external_id_key = "powerbi_dataset_refresh_id"
@@ -139,6 +143,23 @@ class PowerBIDatasetRefreshOperator(ResumableJobMixin,
BaseOperator):
def execute(self, context: Context):
"""Refresh the Power BI Dataset."""
+ if not self.durable:
+ # As in earlier versions of this provider: the trigger starts the
refresh, so only the
+ # triggerer talks to Power BI, and every try starts a new refresh.
+ self.defer(
+ trigger=PowerBITrigger(
+ conn_id=self.conn_id,
+ group_id=self.group_id,
+ dataset_id=self.dataset_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,
+ ),
+ method_name=self.handle_refresh.__name__,
+ )
return self.execute_resumable(context)
def submit_job(self, context: Context) -> str:
@@ -172,19 +193,24 @@ class PowerBIDatasetRefreshOperator(ResumableJobMixin,
BaseOperator):
return refresh_details[PowerBIDatasetRefreshFields.STATUS.value]
def is_job_active(self, status: str) -> bool:
- return status == PowerBIDatasetRefreshStatus.IN_PROGRESS
+ # Without wait_for_completion a retry does not wait for a refresh in
progress either:
+ # is_job_succeeded counts it as done.
+ return self.wait_for_completion and status ==
PowerBIDatasetRefreshStatus.IN_PROGRESS
def is_job_succeeded(self, status: str) -> bool:
- return status == PowerBIDatasetRefreshStatus.COMPLETED
+ if status == PowerBIDatasetRefreshStatus.COMPLETED:
+ return True
+ # Without wait_for_completion, is_job_succeeded counts an in-progress
refresh as done.
+ return not self.wait_for_completion and status ==
PowerBIDatasetRefreshStatus.IN_PROGRESS
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:
+ # get_job_result pushes the refresh id.
self.log.info("Triggered dataset refresh %s (fire-and-forget)",
dataset_refresh_id)
return
+ self._push_refresh_xcom(context, dataset_refresh_id)
self.defer(
trigger=PowerBITrigger(
conn_id=self.conn_id,
@@ -202,7 +228,8 @@ class PowerBIDatasetRefreshOperator(ResumableJobMixin,
BaseOperator):
)
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.
+ # With wait_for_completion the mixin only gets here for a refresh an
earlier try saw complete,
+ # otherwise poll_until_complete defers and execute_complete pushes the
id and the status.
self._push_refresh_xcom(
context,
str(external_id),
@@ -222,8 +249,8 @@ class PowerBIDatasetRefreshOperator(ResumableJobMixin,
BaseOperator):
"""
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.
+ The trigger started the refresh: the task is not durable, or was
deferred by an earlier version
+ of this provider.
: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/version_compat.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py
index 209e8b63f35..8b5cdcf5cc4 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py
@@ -34,8 +34,10 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0)
+AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0)
__all__ = [
"AIRFLOW_V_3_0_PLUS",
"AIRFLOW_V_3_1_PLUS",
+ "AIRFLOW_V_3_3_PLUS",
]
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 772be3739df..3d6ace80136 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
@@ -28,10 +28,11 @@ import pytest
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.compat.sdk import AirflowException, Context,
TaskDeferred, timezone
+from airflow.providers.microsoft.azure.exceptions import
MSGraphPageNotInTaskStateStoreError
from airflow.providers.microsoft.azure.operators.msgraph import (
- COMPLETED_KEY,
- CURSOR_KEY,
- PAGE_KEY_PREFIX,
+ _COMPLETED_KEY,
+ _CURSOR_KEY,
+ _PAGE_KEY_PREFIX,
MSGraphAsyncOperator,
execute_callable,
)
@@ -40,6 +41,7 @@ from airflow.triggers.base import TriggerEvent
from tests_common.test_utils.file_loading import load_file_from_resources,
load_json_from_resources
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.version_compat import AIRFLOW_V_3_3_PLUS
from unit.microsoft.azure.test_utils import (
mock_json_response,
mock_response,
@@ -59,8 +61,13 @@ class TaskStateStore:
return self.values.get(key, default)
def set(self, key, value, *, retention=None):
+ if value is None:
+ raise ValueError("Cannot set value as None")
self.values[key] = value
+ def delete(self, key):
+ self.values.pop(key, None)
+
def users_event(name: str) -> dict[str, Any]:
users = load_json_from_resources(dirname(__file__), "..", "resources",
name)
@@ -77,6 +84,11 @@ def users_operator(**kwargs) -> MSGraphAsyncOperator:
)
+def kept_cursor(pages: int = 1) -> dict[str, Any]:
+ """A cursor as this version of the provider keeps it, pointing at the next
page."""
+ return {"version": 1, "url": "users/delta()?$skiptoken=abc",
"query_parameters": None, "pages": pages}
+
+
class TestMSGraphAsyncOperator:
def test_execute_with_old_result_processor_signature(self):
users = load_json_from_resources(dirname(__file__), "..", "resources",
"users.json")
@@ -548,6 +560,15 @@ class TestMSGraphAsyncOperator:
assert len(recorded_warnings) == 0
[email protected]("airflow.providers.microsoft.azure.operators.msgraph.AIRFLOW_V_3_3_PLUS",
False)
+def test_durable_below_airflow_3_3_warns_and_is_disabled():
+ with pytest.warns(UserWarning, match="`durable` has no effect on Airflow
versions below 3.3"):
+ operator = users_operator(durable=True)
+
+ assert operator.durable is False
+
+
[email protected](not AIRFLOW_V_3_3_PLUS, reason="The task state store
arrived in Airflow 3.3")
class TestDurableMSGraphAsyncOperator:
def test_keeps_the_page_and_the_link_to_the_next_one(self):
operator = users_operator(durable=True)
@@ -559,34 +580,54 @@ class TestDurableMSGraphAsyncOperator:
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},
+ f"{_PAGE_KEY_PREFIX}0": {"result": users["value"]},
+ _CURSOR_KEY: {
+ "version": 1,
+ "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_keeps_a_page_whose_result_is_none(self):
+ operator = users_operator(durable=True)
+ task_state_store = TaskStateStore()
+ context = {"ti": mock.MagicMock(), "task_state_store":
task_state_store}
+ page_without_value = {
+ "@odata.nextLink":
"https://graph.microsoft.com/v1.0/users/delta()?$skiptoken=abc"
+ }
+ event = {"status": "success", "type": "builtins.dict", "response":
json.dumps(page_without_value)}
+
+ with pytest.raises(TaskDeferred):
+ operator.execute_complete(context=context, event=event)
+
+ assert task_state_store.values[f"{_PAGE_KEY_PREFIX}0"] == {"result":
None}
+
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}},
+ **{f"{_PAGE_KEY_PREFIX}{page}": {"result": [{"id": str(page)}]}
for page in range(3)},
+ **{_CURSOR_KEY: kept_cursor(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
+ 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 = {
+ "version": 1,
"url": "users/delta()?$skiptoken=abc",
"query_parameters": {"$top": 12, "$skip": 12},
"pages": 1,
}
- context = {"ti": mock.MagicMock(), "task_state_store":
TaskStateStore(**{CURSOR_KEY: cursor})}
+ context = {"ti": mock.MagicMock(), "task_state_store":
TaskStateStore(**{_CURSOR_KEY: cursor})}
with pytest.raises(TaskDeferred) as exc:
operator.execute(context)
@@ -603,8 +644,9 @@ class TestDurableMSGraphAsyncOperator:
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},
+ f"{_PAGE_KEY_PREFIX}0": {"result": first_page},
+ _CURSOR_KEY: kept_cursor(),
+ "written_by_the_dag": "kept",
}
)
context = {"ti": mock.MagicMock(), "task_state_store":
task_state_store}
@@ -612,7 +654,65 @@ class TestDurableMSGraphAsyncOperator:
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
+ assert task_state_store.values == {"written_by_the_dag": "kept"}
+
+ def test_page_no_longer_in_the_store_fails_and_removes_the_pages(self):
+ operator = users_operator(durable=True)
+ task_state_store = TaskStateStore(
+ **{
+ f"{_PAGE_KEY_PREFIX}0": {"result": [{"id": "0"}]},
+ _CURSOR_KEY: kept_cursor(pages=2),
+ "written_by_the_dag": "kept",
+ }
+ )
+ context = {"ti": mock.MagicMock(), "task_state_store":
task_state_store}
+
+ with pytest.raises(MSGraphPageNotInTaskStateStoreError, match="Page 2
is no longer"):
+ operator.execute_complete(context=context,
event=users_event("next_users.json"))
+
+ assert task_state_store.values == {"written_by_the_dag": "kept"}
+
+ @pytest.mark.parametrize(
+ "page",
+ [[{"id": "0"}], {"result": "value", "metadata": "keep"}],
+ ids=["list", "dict_with_a_result_key"],
+ )
+ @pytest.mark.parametrize("event", ["users.json", "next_users.json"],
ids=["next_page", "last_page"])
+ def test_pages_kept_by_an_earlier_version_fail_and_are_removed(self, page,
event):
+ """15.3.0rc1 kept the pages unwrapped, under a cursor without a
version."""
+ operator = users_operator(durable=True)
+ task_state_store = TaskStateStore(
+ **{
+ f"{_PAGE_KEY_PREFIX}0": page,
+ _CURSOR_KEY: {"url": "users/delta()?$skiptoken=abc",
"query_parameters": None, "pages": 1},
+ "written_by_the_dag": "kept",
+ }
+ )
+ context = {"ti": mock.MagicMock(), "task_state_store":
task_state_store}
+
+ with pytest.raises(MSGraphPageNotInTaskStateStoreError, match="earlier
version of the provider"):
+ operator.execute_complete(context=context,
event=users_event(event))
+
+ assert task_state_store.values == {"written_by_the_dag": "kept"}
+ context["ti"].xcom_push.assert_not_called()
+
+ def
test_execute_removes_the_pages_kept_by_an_earlier_version_and_starts_again(self):
+ operator = users_operator(durable=True)
+ task_state_store = TaskStateStore(
+ **{
+ f"{_PAGE_KEY_PREFIX}0": {"result": "value", "metadata":
"keep"},
+ _CURSOR_KEY: {"url": "users/delta()?$skiptoken=abc",
"query_parameters": None, "pages": 1},
+ "written_by_the_dag": "kept",
+ }
+ )
+ context = {"ti": mock.MagicMock(), "task_state_store":
task_state_store}
+
+ with pytest.raises(TaskDeferred) as exc:
+ operator.execute(context)
+
+ assert exc.value.trigger.url == operator.url
+ assert exc.value.trigger.pagination_link is False
+ assert task_state_store.values == {"written_by_the_dag": "kept"}
def test_keeps_the_result_of_a_request_that_is_not_a_get(self):
operator = users_operator(durable=True, method="POST",
data={"requestBody": "value"})
@@ -621,13 +721,28 @@ class TestDurableMSGraphAsyncOperator:
results = operator.execute_complete(context=context,
event=users_event("next_users.json"))
- assert task_state_store.values == {COMPLETED_KEY: {"result": results}}
+ assert task_state_store.values == {_COMPLETED_KEY: {"result": results}}
+
+ def
test_keeps_only_the_result_of_a_paginated_request_that_is_not_a_get(self):
+ operator = users_operator(durable=True, method="POST",
data={"requestBody": "value"})
+ task_state_store = TaskStateStore(
+ **{
+ f"{_PAGE_KEY_PREFIX}0": {"result": [{"id": "0"}]},
+ _CURSOR_KEY: kept_cursor(),
+ }
+ )
+ context = {"ti": mock.MagicMock(), "task_state_store":
task_state_store}
+
+ results = operator.execute_complete(context=context,
event=users_event("next_users.json"))
+
+ assert results[0] == {"id": "0"}
+ 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"}]}}),
+ "task_state_store": TaskStateStore(**{_COMPLETED_KEY: {"result":
[{"id": "1"}]}}),
}
with mock.patch.object(operator, "defer") as mock_defer:
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 5a0dcb781e0..4a9ae25cee3 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
@@ -273,6 +273,7 @@ class TestPowerBIDatasetRefreshOperator:
assert url == EXPECTED_ITEM_RUN_OP_EXTRA_LINK
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The worker starts the
refresh only from Airflow 3.3")
@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)
@@ -292,10 +293,10 @@ class TestPowerBIDatasetRefreshOperator:
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
+ assert task_state_store.values == {REFRESH_ID_KEY:
NEW_REFRESH_REQUEST_ID}
+ assert exc.value.trigger.cancel_on_error is False
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The worker starts the
refresh only from Airflow 3.3")
@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."""
@@ -305,10 +306,34 @@ class TestPowerBIDatasetRefreshOperator:
assert operator.execute(context) is None
mock_trigger_dataset_refresh.assert_called_once()
- context["ti"].xcom_push.assert_called_with(
+ context["ti"].xcom_push.assert_called_once_with(
key=f"{TASK_ID}.powerbi_dataset_refresh_id",
value=NEW_REFRESH_REQUEST_ID
)
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The task state store
arrived in Airflow 3.3")
+ @pytest.mark.parametrize(
+ "status", [PowerBIDatasetRefreshStatus.IN_PROGRESS,
PowerBIDatasetRefreshStatus.COMPLETED]
+ )
+ @mock.patch.object(PowerBIHook, "get_refresh_history")
+ @mock.patch.object(PowerBIHook, "trigger_dataset_refresh")
+ def test_execute_fire_and_forget_does_not_refresh_again_once_started(
+ self, mock_trigger_dataset_refresh, mock_get_refresh_history, status
+ ):
+ mock_get_refresh_history.return_value = [previous_refresh(status)]
+ operator = PowerBIDatasetRefreshOperator(**CONFIG,
wait_for_completion=False)
+ 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()
+ context["ti"].xcom_push.assert_called_once_with(
+ key=f"{TASK_ID}.powerbi_dataset_refresh_id",
value=PREVIOUS_REFRESH_REQUEST_ID
+ )
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The worker starts the
refresh only from Airflow 3.3")
@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)
@@ -485,7 +510,10 @@ class TestPowerBIDatasetRefreshOperator:
assert task_state_store.values == {REFRESH_ID_KEY:
PREVIOUS_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):
+ def test_execute_not_durable_defers_to_the_trigger_that_starts_the_refresh(
+ self, mock_trigger_dataset_refresh
+ ):
+ """Without durable execution the worker never talks to Power BI, as
before the operator was durable."""
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}
@@ -493,11 +521,56 @@ class TestPowerBIDatasetRefreshOperator:
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
+ mock_trigger_dataset_refresh.assert_not_called()
+ context["ti"].xcom_push.assert_not_called()
+ assert isinstance(exc.value.trigger, PowerBITrigger)
+ assert exc.value.trigger.dataset_refresh_id is None
+ assert exc.value.trigger.request_body == REQUEST_BODY
+ assert exc.value.trigger.wait_for_termination is False
assert exc.value.trigger.cancel_on_error is True
+ assert exc.value.method_name == "handle_refresh"
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_not_durable_full_lifecycle(
+ self, mock_trigger_dataset_refresh,
mock_get_refresh_details_by_refresh_id
+ ):
+ """Assert the full lifecycle: start in the trigger, defer again to
poll, complete on the worker."""
+ mock_get_refresh_details_by_refresh_id.return_value = {
+ PowerBIDatasetRefreshFields.STATUS.value:
PowerBIDatasetRefreshStatus.COMPLETED,
+ PowerBIDatasetRefreshFields.ERROR.value: None,
+ }
+ operator = PowerBIDatasetRefreshOperator(**CONFIG, durable=False)
+ context = {"ti": MagicMock(), "task_state_store": TaskStateStore()}
+
+ with pytest.raises(TaskDeferred) as started:
+ operator.execute(context)
+ [started_event] = run_trigger(started.value.trigger)
+ with pytest.raises(TaskDeferred) as polled:
+ getattr(operator, started.value.method_name)(context=context,
event=started_event.payload)
+ [completed_event] = run_trigger(polled.value.trigger)
+ assert (
+ getattr(operator, polled.value.method_name)(context=context,
event=completed_event.payload)
+ is None
+ )
+
+ mock_trigger_dataset_refresh.assert_called_once_with(
+ dataset_id=DATASET_ID, group_id=GROUP_ID, request_body=REQUEST_BODY
+ )
+ assert started_event.payload["dataset_refresh_id"] ==
NEW_REFRESH_REQUEST_ID
+ assert polled.value.trigger.dataset_refresh_id ==
NEW_REFRESH_REQUEST_ID
+ assert completed_event.payload["dataset_refresh_status"] ==
PowerBIDatasetRefreshStatus.COMPLETED
+ assert context["ti"].xcom_push.call_args_list == [
+ mock.call(key=f"{TASK_ID}.powerbi_dataset_refresh_id",
value=NEW_REFRESH_REQUEST_ID),
+ mock.call(key=f"{TASK_ID}.powerbi_dataset_refresh_id",
value=NEW_REFRESH_REQUEST_ID),
+ mock.call(
+ key=f"{TASK_ID}.powerbi_dataset_refresh_status",
value=PowerBIDatasetRefreshStatus.COMPLETED
+ ),
+ ]
+ assert context["task_state_store"].values == {}
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The worker starts the
refresh only from Airflow 3.3")
@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(
