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 a3f9792d45664..ccbf025effeb9 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 @@ -66,12 +66,13 @@ 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_completion`` is True. + :param timeout: Time in seconds to wait for a dataset to reach a terminal status for asynchronous waits. :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``. + :param deferrable: This parameter is deprecated and no longer has any effect. The operator now always + uses deferrable execution. """ template_fields: Sequence[str] = ( @@ -99,9 +100,7 @@ def __init__( ) -> 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.log.warning("The PowerBIDatasetRefreshOperator now always uses deferrable execution.") 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 @@ -121,40 +120,56 @@ def api_version(self) -> str | None: def execute(self, context: Context): """Refresh the Power BI Dataset.""" - 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 - ) - - dataset_refresh_id = hook.trigger_dataset_refresh( - dataset_id=self.dataset_id, + 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__, + ) + + def handle_refresh(self, context: Context, event: dict[str, str] | None) -> None: + """ + Handle refresh-trigger event and optionally defer again to wait for refresh completion. + + :param context: Airflow context dictionary + :param event: Event dict from trigger with status and dataset_refresh_id + """ + if not event: + return + + dataset_refresh_id = event.get("dataset_refresh_id") + if dataset_refresh_id: + context["ti"].xcom_push( + key=f"{self.task_id}.powerbi_dataset_refresh_id", + value=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") + if event["status"] == "error": + raise AirflowException(event["message"]) + + if not self.wait_for_completion: + self.log.info("Triggered dataset refresh %s (fire-and-forget)", dataset_refresh_id) 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, + 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_completion, - request_body=self.request_body, + wait_for_termination=True, ), method_name=self.execute_complete.__name__, ) 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 b5be85f3c7f2e..ad8408455f9ea 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 @@ -26,17 +26,19 @@ from airflow.providers.microsoft.azure.hooks.powerbi import ( PowerBIDatasetRefreshFields, PowerBIDatasetRefreshStatus, + PowerBIHook, ) from airflow.providers.microsoft.azure.operators.powerbi import PowerBIDatasetRefreshOperator from airflow.providers.microsoft.azure.triggers.powerbi import PowerBITrigger from tests_common.test_utils.mock_context import mock_context +from tests_common.test_utils.operators.run_deferrable import execute_operator from unit.microsoft.azure.test_utils import get_airflow_connection try: from airflow.sdk import timezone except ImportError: - from airflow.utils import timezone # type: ignore[no-redef] + from airflow.utils import timezone # type: ignore[attr-defined, no-redef] DEFAULT_CONNECTION_CLIENT_SECRET = "powerbi_conn_id" @@ -109,10 +111,12 @@ def test_execute_wait_for_completion_with_deferrable(self, connection): assert isinstance(exc.value.trigger, PowerBITrigger) assert exc.value.trigger.dataset_refresh_id is None + assert exc.value.trigger.wait_for_termination is False + assert exc.value.method_name == "handle_refresh" @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection) def test_powerbi_operator_async_get_refresh_status_success(self, connection): - """Test that execute defers once when wait_for_completion=True""" + """Test that execute first defers to trigger refresh when wait_for_completion=True""" operator = PowerBIDatasetRefreshOperator( **CONFIG, wait_for_completion=True, # Explicitly set to True @@ -128,14 +132,119 @@ def test_powerbi_operator_async_get_refresh_status_success(self, connection): # 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 + assert exc.value.trigger.wait_for_termination is False # Verify callback method name - assert exc.value.method_name == "execute_complete" + assert exc.value.method_name == "handle_refresh" # Verify dataset_refresh_id is None (trigger will create it) assert exc.value.trigger.dataset_refresh_id is None + @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection) + def test_handle_refresh_wait_for_completion(self, connection): + """Test that handle_refresh defers to execute_complete when wait_for_completion=True.""" + operator = PowerBIDatasetRefreshOperator( + **CONFIG, + wait_for_completion=True, + ) + context = {"ti": MagicMock()} + + with pytest.raises(TaskDeferred) as exc: + operator.handle_refresh( + context=context, + event=SUCCESS_TRIGGER_EVENT, + ) + + assert context["ti"].xcom_push.call_count == 1 + assert isinstance(exc.value.trigger, PowerBITrigger) + 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" + + @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection) + def test_handle_refresh_fire_and_forget(self, connection): + """Test that handle_refresh finishes immediately for fire-and-forget mode.""" + operator = PowerBIDatasetRefreshOperator( + **CONFIG, + wait_for_completion=False, + ) + context = {"ti": MagicMock()} + + operator.handle_refresh( + context=context, + event=SUCCESS_TRIGGER_EVENT, + ) + + assert context["ti"].xcom_push.call_count == 1 + + @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection) + def test_handle_refresh_failure(self, connection): + """Test that handle_refresh raises exception on trigger error.""" + operator = PowerBIDatasetRefreshOperator( + **CONFIG, + wait_for_completion=False, + ) + context = {"ti": MagicMock()} + + with pytest.raises(AirflowException, match="Failed to trigger the dataset refresh."): + operator.handle_refresh( + context=context, + event={ + "status": "error", + "dataset_refresh_status": None, + "message": "Failed to trigger the dataset refresh.", + "dataset_refresh_id": None, + }, + ) + + assert context["ti"].xcom_push.call_count == 0 + + @mock.patch.object(PowerBIHook, "get_refresh_details_by_refresh_id") + @mock.patch.object(PowerBIHook, "trigger_dataset_refresh") + def test_execute_operator_wait_for_completion_full_lifecycle( + self, mock_trigger_dataset_refresh, mock_get_refresh_details_by_refresh_id + ): + """Assert the full deferrable lifecycle completes successfully.""" + mock_trigger_dataset_refresh.return_value = NEW_REFRESH_REQUEST_ID + mock_get_refresh_details_by_refresh_id.side_effect = [ + { + PowerBIDatasetRefreshFields.STATUS.value: PowerBIDatasetRefreshStatus.IN_PROGRESS, + PowerBIDatasetRefreshFields.ERROR.value: None, + }, + { + PowerBIDatasetRefreshFields.STATUS.value: PowerBIDatasetRefreshStatus.COMPLETED, + PowerBIDatasetRefreshFields.ERROR.value: None, + }, + ] + + operator = PowerBIDatasetRefreshOperator( + **CONFIG, + wait_for_completion=True, + ) + + result, events = execute_operator(operator) + + assert result is None + mock_trigger_dataset_refresh.assert_called_once_with( + dataset_id=DATASET_ID, + group_id=GROUP_ID, + request_body=REQUEST_BODY, + ) + assert mock_get_refresh_details_by_refresh_id.call_count == 2 + assert len(events) == 2 + assert events[0].payload == { + "status": "success", + "dataset_refresh_status": None, + "message": f"The dataset refresh {NEW_REFRESH_REQUEST_ID} has been triggered.", + "dataset_refresh_id": NEW_REFRESH_REQUEST_ID, + } + assert events[1].payload == { + "status": "success", + "dataset_refresh_status": PowerBIDatasetRefreshStatus.COMPLETED, + "message": f"The dataset refresh {NEW_REFRESH_REQUEST_ID} has {PowerBIDatasetRefreshStatus.COMPLETED}.", + "dataset_refresh_id": NEW_REFRESH_REQUEST_ID, + } + def test_powerbi_operator_async_execute_complete_success(self): """Assert that execute_complete processes success event correctly""" operator = PowerBIDatasetRefreshOperator(**CONFIG) @@ -232,63 +341,25 @@ 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 - + 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 = {"ti": MagicMock()} - context["ti"].task_id = TASK_ID + context = mock_context(task=operator) - # Should raise AirflowException - with pytest.raises(AirflowException, match="Failed to trigger dataset refresh"): + with pytest.raises(TaskDeferred) as exc: operator.execute(context) - # Should not push to XCom on failure - assert context["ti"].xcom_push.call_count == 0 + assert isinstance(exc.value.trigger, PowerBITrigger) + assert exc.value.trigger.wait_for_termination is False + assert exc.value.method_name == "handle_refresh" @mock.patch.object(BaseHook, "get_connection", side_effect=get_airflow_connection) def test_execute_default_behavior_waits_for_completion(self, mock_connection): - """Test that default behavior (wait_for_completion=True) defers and waits""" + """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, @@ -303,11 +374,11 @@ def test_execute_default_behavior_waits_for_completion(self, mock_connection): operator = PowerBIDatasetRefreshOperator(**config_without_wait) context = mock_context(task=operator) - # Should defer (because default is wait_for_completion=True) + # Should defer for initial trigger step (because default is wait_for_completion=True) with pytest.raises(TaskDeferred) as exc: operator.execute(context) # Verify it deferred with correct trigger assert isinstance(exc.value.trigger, PowerBITrigger) - assert exc.value.trigger.wait_for_termination is True - assert exc.value.method_name == "execute_complete" + assert exc.value.trigger.wait_for_termination is False + assert exc.value.method_name == "handle_refresh"