You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "dabla (via GitHub)" <gi...@apache.org> on 2026/10/13 23:21:40 UTC
[PR] Make the PowerBI refresh and MSGraph operators resume after a retry [airflow]
dabla opened a new pull request, #74675:
URL: https://github.com/apache/airflow/pull/74675
A retry of `PowerBIDatasetRefreshOperator`, for instance after a worker crash, starts a new dataset refresh while the first one may still be running. A retry of `MSGraphAsyncOperator` sends its request again and fetches every page from the first one. There is a second gap: the refresh is started by the trigger, and its id lives only in the trigger's memory until the first event. So a triggerer restart in that window starts a second refresh as well.
This builds on `ResumableJobMixin` from `common.compat` (#74673).
**PowerBI refresh**
- The refresh is now started from the worker.
- Its id is kept in the task state store before the task defers to poll.
- A retry reconnects to a refresh still in progress and does not start a new one once it has completed.
- A refresh that failed, or that dropped out of the dataset's refresh history, is started again.
- Since the polling trigger gets the id from the start, a triggerer restart no longer starts a second refresh, on any Airflow version.
- A failing poll no longer cancels the refresh when the operator is durable, so the retry can pick it up.
- As far as I can see, this is the first deferrable operator built on the mixin: `poll_until_complete` defers, and the mixin has already stored the id by then.
**MSGraph**
- New opt-in `durable` flag. Every processed page and the link to the next one go to the task state store, so a retry continues at the page it stopped at.
- A request that is not a `GET` keeps its result, so a retry returns it instead of sending the request again.
- It is opt-in because pages are stored as they are, and Graph pages can be larger than the store's default value limit.
Resuming needs Airflow 3.3. On older versions, both operators behave as before.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (please specify the tool below)
Generated-by: Claude Code (Opus 5.5) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make the PowerBI refresh and MSGraph operators resume after a retry [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #74675:
URL: https://github.com/apache/airflow/pull/74675#discussion_r4268268281
##########
providers/microsoft/azure/docs/operators/msgraph.rst:
##########
@@ -80,6 +80,21 @@ 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]
+Resuming after a retry
+^^^^^^^^^^^^^^^^^^^^^^
+
+By default a retry, for instance after a worker crash, sends the request again and fetches every page from
Review Comment:
Good point, thanks. Both operator pages now keep only what is specific to the operator, under a "Durable execution" sub-section like the Databricks operators, and link to the resumable tasks and task state store docs for the concept itself.
---
Drafted-by: Claude Code (Opus 5.5); reviewed by @dabla before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make the PowerBI refresh and MSGraph operators durable [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #74675:
URL: https://github.com/apache/airflow/pull/74675#issuecomment-6248283466
Thanks, both right. The durable path now reads the stored pages only on the last page, so an intermediate page reads just the cursor, and a test counts the reads. The warning about the missing task state store is logged once from `execute` instead of on every page.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make the PowerBI refresh and MSGraph operators resume after a retry [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #74675:
URL: https://github.com/apache/airflow/pull/74675#discussion_r4268268281
##########
providers/microsoft/azure/docs/operators/msgraph.rst:
##########
@@ -80,6 +80,21 @@ 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]
+Resuming after a retry
+^^^^^^^^^^^^^^^^^^^^^^
+
+By default a retry, for instance after a worker crash, sends the request again and fetches every page from
Review Comment:
Good point, thanks. Both operator pages now keep only what is specific to the operator, under a "Durable execution" sub-section like the Databricks operators, and link to the resumable tasks and task state store docs for the concept itself.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make the PowerBI refresh and MSGraph operators resume after a retry [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #74675:
URL: https://github.com/apache/airflow/pull/74675#discussion_r4266026002
##########
providers/microsoft/azure/docs/operators/msgraph.rst:
##########
@@ -80,6 +80,21 @@ 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]
+Resuming after a retry
+^^^^^^^^^^^^^^^^^^^^^^
+
+By default a retry, for instance after a worker crash, sends the request again and fetches every page from
Review Comment:
The text here is a general explnation of durable. Maybe we should link to the durable docs? Otherwise it means explnation will be in every operator page
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make the PowerBI refresh and MSGraph operators durable [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla merged PR #74675:
URL: https://github.com/apache/airflow/pull/74675
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make the PowerBI refresh and MSGraph operators durable [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #74675:
URL: https://github.com/apache/airflow/pull/74675#discussion_r4334635863
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py:
##########
@@ -161,8 +185,26 @@ def __init__(
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"]
Review Comment:
Dead code. Nothing reads `self.query_parameters` between that line and the `TaskDeferred` raised by `defer_to_next_link` on the next line, which takes the value as an argument. `execute_complete` sets it again from its own `query_parameters` kwarg at line 246. Drop the assignment.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py:
##########
@@ -161,8 +185,26 @@ def __init__(
self.result_processor = result_processor
self.event_handler = event_handler or default_event_handler
self.serializer: ResponseSerializer = serializer()
+ self.durable = durable
Review Comment:
PowerBI uses the mixin, so durable below 3.3 raises a `UserWarning` at parse time. MSGraph only logs at execute time. Match MSGraph to the mixin.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py:
##########
@@ -161,8 +185,26 @@ def __init__(
self.result_processor = result_processor
self.event_handler = event_handler or default_event_handler
self.serializer: ResponseSerializer = serializer()
+ self.durable = durable
Review Comment:
Or use mixin itself?
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py:
##########
Review Comment:
The task state store keys are never cleaned up. A 500 page run leaves 501 keys per task instance until retention expires them. We have a delete and clear. Calling `clear()` once the result is returned would keep the table small. Say in the docs that it is deliberate if it is.
##########
providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_powerbi.py:
##########
@@ -335,44 +272,171 @@ def test_powerbi_link(self, dag_maker, create_task_instance_of_operator):
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(
Review Comment:
`assert_called_with only` checks the last call, so it hides a duplicate xcom write. Fire and forget pushes the refresh id twice: once in `poll_until_complete`, again in `get_job_result`. Use `assert_called_once_with`, and skip the second push when `wait_for_completion` is false.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py:
##########
@@ -39,6 +41,20 @@
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"
Review Comment:
These dont have to be public interfaces...
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py:
##########
@@ -236,21 +281,67 @@ def execute_complete(
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}"))
Review Comment:
`pull_pages` appends whatever get returns for each page key. A key that expired under `[state_store] default_retention_days`, or that a custom backend dropped, reads back as `None`. The operator then returns `[page0, None, page2]` and the task succeeds with a hole in its data. Wrapping the page as `{"result": ...}` lets this be told apart from a stored None and raised.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py:
##########
@@ -236,21 +281,67 @@ def execute_complete(
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)
Review Comment:
`push_page` writes the page result straight into the task state store store. The task state store refuses `None`. So a `durable` paginated run dies the moment `result_processor` returns None for any page that is not the last.
Take this scenario - `MSGraphAsyncOperator(url="users/delta()", durable=True, result_processor=lambda result, **context: result.get("value"))`, which is the exact processor the PR's own test helper uses. Page 1 comes back without a "value" key, so result is None. `execute_complete` catches `TaskDeferred`, calls `push_page`, and `task_state_store.set("msgraph_page_0", None)` raises `ValueError: Cannot set value as None`. The task fails instead of deferring. The same operator with `durable=False` works, because `append_result` appends `None` to the list without complaint ([msgraph.py:352-356](vscode-webview://1dkhhega5psqueji9am5949j63htqtc3np0ikcrp60r3krob8h43/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py#L352-L356)).
The fix is perhaps just one method away. `COMPLETED_KEY` is stored as `{"result": results}` at line 304 precisely so a falsy result survives. Do the same for a page: `set(f"{PAGE_KEY_PREFIX}{page}", {"result": result})`, and unwrap in `pull_pages`.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/powerbi.py:
##########
@@ -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"
Review Comment:
Two defaults change for users who set nothing, with no changelog entry too! The Power BI call moves from the triggerer to the worker, so a deployment that only gave the triggerer egress to `api.powerbi.com` breaks. And durable defaults to True, so `cancel_on_error` is False and a polling error no longer cancels the refresh. Worth calling better in docs?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org