Add wait_for_termination parameter and fix double-deferral in PowerBIDatasetRefreshOperator (#60369)

Author: shreyas-dev

providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/powerbi.py

@@ -20,6 +20,7 @@
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
+from airflow.configuration import conf
 from airflow.providers.common.compat.sdk import AirflowException, BaseOperator, BaseOperatorLink
 from airflow.providers.microsoft.azure.hooks.powerbi import PowerBIHook
 from airflow.providers.microsoft.azure.triggers.powerbi import (
@@ -61,10 +62,12 @@ class PowerBIDatasetRefreshOperator(BaseOperator):
     :param dataset_id: The dataset id.
     :param group_id: The workspace id.
     :param conn_id: Airflow Connection ID that contains the connection information for the Power BI account used for authentication.
-    :param timeout: Time in seconds to wait for a dataset to reach a terminal status for asynchronous waits. Used only if ``wait_for_termination`` is True.
+    :param timeout: Time in seconds to wait for a dataset to reach a terminal status for asynchronous waits. Used only if ``wait_for_completion`` is True.
     :param check_interval: Number of seconds to wait before rechecking the
         refresh status.
     :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 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 when ``wait_for_completion=True``.
     """
 
     template_fields: Sequence[str] = (
@@ -86,13 +89,19 @@ def __init__(
         api_version: APIVersion | str | None = None,
         check_interval: int = 60,
         request_body: dict[str, Any] | None = None,
+        wait_for_completion: bool = True,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
+        if "deferrable" in kwargs or deferrable is not True:
+            self.log.warning(
+                "The PowerBIDatasetRefreshOperator now always uses deferrable execution when wait_for_completion=True."
+            )
         self.hook = PowerBIHook(conn_id=conn_id, proxies=proxies, api_version=api_version, timeout=timeout)
         self.dataset_id = dataset_id
         self.group_id = group_id
-        self.wait_for_termination = True
+        self.wait_for_completion = wait_for_completion
         self.conn_id = conn_id
         self.timeout = timeout
         self.check_interval = check_interval
@@ -108,63 +117,76 @@ def api_version(self) -> str | None:
 
     def execute(self, context: Context):
         """Refresh the Power BI Dataset."""
-        if self.wait_for_termination:
-            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=self.wait_for_termination,
-                    request_body=self.request_body,
-                ),
-                method_name=self.get_refresh_status.__name__,
+        if not self.wait_for_completion:
+            # Fire and forget - synchronous execution, no deferral
+            hook = PowerBIHook(
+                conn_id=self.conn_id, proxies=self.proxies, api_version=self.api_version, timeout=self.timeout
             )
 
-    def get_refresh_status(self, context: Context, event: dict[str, str] | None = None):
-        """Push the refresh Id to XCom then runs the Trigger to wait for refresh completion."""
-        if event:
-            if event["status"] == "error":
-                raise AirflowException(event["message"])
+            dataset_refresh_id = hook.trigger_dataset_refresh(
+                dataset_id=self.dataset_id,
+                group_id=self.group_id,
+                request_body=self.request_body,
+            )
 
-            dataset_refresh_id = event["dataset_refresh_id"]
+            if dataset_refresh_id:
+                self.log.info("Triggered dataset refresh %s (fire-and-forget)", dataset_refresh_id)
+                context["ti"].xcom_push(
+                    key=f"{self.task_id}.powerbi_dataset_refresh_id",
+                    value=dataset_refresh_id,
+                )
+            else:
+                raise AirflowException("Failed to trigger dataset refresh")
+            return
+
+        # Wait for termination - use deferrable trigger
+        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=self.wait_for_completion,
+                request_body=self.request_body,
+            ),
+            method_name=self.execute_complete.__name__,
+        )
+
+    def execute_complete(self, context: Context, event: dict[str, str]) -> None:
+        """
+        Handle trigger completion and push results to XCom or raise an exception.
+
+        :param context: Airflow context dictionary
+        :param event: Event dict from trigger with status and dataset_refresh_id
+        """
+        if not event:
+            return
+
+        # Success - push both ID and status to XCom
+        dataset_refresh_id = event.get("dataset_refresh_id")
+        dataset_refresh_status = event.get("dataset_refresh_status")
 
         if dataset_refresh_id:
             context["ti"].xcom_push(
-                key=f"{self.task_id}.powerbi_dataset_refresh_Id",
+                key=f"{self.task_id}.powerbi_dataset_refresh_id",
                 value=dataset_refresh_id,
             )
-            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=self.wait_for_termination,
-                ),
-                method_name=self.execute_complete.__name__,
-            )
-
-    def execute_complete(self, context: Context, event: dict[str, str]) -> Any:
-        """
-        Return immediately - callback for when the trigger fires.
 
-        Relies on trigger to throw an exception, otherwise it assumes execution was successful.
-        """
-        if event:
+        if dataset_refresh_status:
             context["ti"].xcom_push(
                 key=f"{self.task_id}.powerbi_dataset_refresh_status",
-                value=event["dataset_refresh_status"],
+                value=dataset_refresh_status,
             )
-            if event["status"] == "error":
-                raise AirflowException(event["message"])
+
+        if event["status"] == "error":
+            raise AirflowException(event["message"])
+
+        self.log.info(
+            "Dataset refresh %s completed with status: %s", dataset_refresh_id, dataset_refresh_status
+        )
 
 
 class PowerBIWorkspaceListOperator(BaseOperator):

providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_powerbi.py

@@ -98,7 +98,7 @@
 
 class TestPowerBIDatasetRefreshOperator:
     @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection)
-    def test_execute_wait_for_termination_with_deferrable(self, connection):
+    def test_execute_wait_for_completion_with_deferrable(self, connection):
         operator = PowerBIDatasetRefreshOperator(
             **CONFIG,
         )
@@ -112,34 +112,52 @@ def test_execute_wait_for_termination_with_deferrable(self, connection):
 
     @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection)
     def test_powerbi_operator_async_get_refresh_status_success(self, connection):
-        """Assert that get_refresh_status log success message"""
+        """Test that execute defers once when wait_for_completion=True"""
         operator = PowerBIDatasetRefreshOperator(
             **CONFIG,
+            wait_for_completion=True,  # Explicitly set to True
         )
-        context = {"ti": MagicMock()}
-        context["ti"].task_id = TASK_ID
+        context = mock_context(task=operator)
 
         with pytest.raises(TaskDeferred) as exc:
-            operator.get_refresh_status(
-                context=context,
-                event=SUCCESS_TRIGGER_EVENT,
-            )
+            operator.execute(context)
 
+        # Verify trigger is correct type
         assert isinstance(exc.value.trigger, PowerBITrigger)
-        assert exc.value.trigger.dataset_refresh_id is NEW_REFRESH_REQUEST_ID
-        assert context["ti"].xcom_push.call_count == 1
+
+        # 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 True
+
+        # Verify callback method name
+        assert exc.value.method_name == "execute_complete"
+
+        # Verify dataset_refresh_id is None (trigger will create it)
+        assert exc.value.trigger.dataset_refresh_id is None
 
     def test_powerbi_operator_async_execute_complete_success(self):
-        """Assert that execute_complete log success message"""
-        operator = PowerBIDatasetRefreshOperator(
-            **CONFIG,
-        )
+        """Assert that execute_complete processes success event correctly"""
+        operator = PowerBIDatasetRefreshOperator(**CONFIG)
         context = {"ti": MagicMock()}
+
         operator.execute_complete(
             context=context,
             event=SUCCESS_REFRESH_EVENT,
         )
-        assert context["ti"].xcom_push.call_count == 1
+
+        # Should push both refresh_id and status
+        assert context["ti"].xcom_push.call_count == 2
+
+        # Verify the XCom keys and values
+        calls = context["ti"].xcom_push.call_args_list
+        xcom_data = {call[1]["key"]: call[1]["value"] for call in calls}
+
+        assert f"{TASK_ID}.powerbi_dataset_refresh_id" in xcom_data
+        assert xcom_data[f"{TASK_ID}.powerbi_dataset_refresh_id"] == NEW_REFRESH_REQUEST_ID
+
+        assert f"{TASK_ID}.powerbi_dataset_refresh_status" in xcom_data
+        assert xcom_data[f"{TASK_ID}.powerbi_dataset_refresh_status"] == PowerBIDatasetRefreshStatus.COMPLETED
 
     def test_powerbi_operator_async_execute_complete_fail(self):
         """Assert that execute_complete raise exception on error"""
@@ -176,7 +194,7 @@ def test_powerbi_operator_refresh_fail(self):
                     "dataset_refresh_id": "1234",
                 },
             )
-        assert context["ti"].xcom_push.call_count == 1
+        assert context["ti"].xcom_push.call_count == 2
         assert str(exc.value) == "error message"
 
     def test_execute_complete_no_event(self):
@@ -213,3 +231,83 @@ def test_powerbi_link(self, dag_maker, create_task_instance_of_operator):
         )
 
         assert url == EXPECTED_ITEM_RUN_OP_EXTRA_LINK
+
+    @mock.patch("airflow.providers.microsoft.azure.operators.powerbi.PowerBIHook")
+    @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection)
+    def test_execute_fire_and_forget_mode(self, mock_connection, mock_hook_class):
+        """Test fire-and-forget mode (wait_for_completion=False)"""
+        mock_hook_instance = mock_hook_class.return_value
+        mock_hook_instance.trigger_dataset_refresh.return_value = NEW_REFRESH_REQUEST_ID
+
+        operator = PowerBIDatasetRefreshOperator(
+            **CONFIG,
+            wait_for_completion=False,
+        )
+        context = {"ti": MagicMock()}
+        context["ti"].task_id = TASK_ID
+
+        # Should not raise TaskDeferred
+        result = operator.execute(context)
+
+        # Verify hook was called correctly
+        mock_hook_instance.trigger_dataset_refresh.assert_called_once_with(
+            dataset_id=DATASET_ID,
+            group_id=GROUP_ID,
+            request_body=REQUEST_BODY,
+        )
+
+        # Verify XCom push
+        assert context["ti"].xcom_push.call_count == 1
+        call_args = context["ti"].xcom_push.call_args
+        assert call_args[1]["key"] == f"{TASK_ID}.powerbi_dataset_refresh_id"
+        assert call_args[1]["value"] == NEW_REFRESH_REQUEST_ID
+
+        # Should return None (completes synchronously)
+        assert result is None
+
+    @mock.patch("airflow.providers.microsoft.azure.operators.powerbi.PowerBIHook")
+    @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection)
+    def test_execute_fire_and_forget_mode_failure(self, mock_connection, mock_hook_class):
+        """Test fire-and-forget mode raises exception when trigger fails"""
+        mock_hook_instance = mock_hook_class.return_value
+        mock_hook_instance.trigger_dataset_refresh.return_value = None
+
+        operator = PowerBIDatasetRefreshOperator(
+            **CONFIG,
+            wait_for_completion=False,
+        )
+        context = {"ti": MagicMock()}
+        context["ti"].task_id = TASK_ID
+
+        # Should raise AirflowException
+        with pytest.raises(AirflowException, match="Failed to trigger dataset refresh"):
+            operator.execute(context)
+
+        # Should not push to XCom on failure
+        assert context["ti"].xcom_push.call_count == 0
+
+    @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) defers and waits"""
+        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
+        }
+
+        operator = PowerBIDatasetRefreshOperator(**config_without_wait)
+        context = mock_context(task=operator)
+
+        # Should defer (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 True
+        assert exc.value.method_name == "execute_complete"