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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
