This is an automated email from the ASF dual-hosted git repository.
potiuk 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 0440172082e Add deferrable mode support to
AzureSynapseRunPipelineOperator (#63614)
0440172082e is described below
commit 0440172082e16fb3588667e2dc0fd50860b290d6
Author: SameerMesiah97 <[email protected]>
AuthorDate: Tue Apr 28 00:41:55 2026 +0100
Add deferrable mode support to AzureSynapseRunPipelineOperator (#63614)
* Add deferrable mode support to AzureSynapseRunPipelineOperator.
Introduce AzureSynapsePipelineTrigger and update the operator to defer
execution when deferrable=True by delegating pipeline run monitoring to
the trigger. Implement execute_complete to handle trigger events and
raise AzureSynapsePipelineRunException for failure conditions.
Add unit tests for the deferrable execution path covering pre-deferral
terminal states, deferral for intermediate states, and execute_complete
handling of success and error trigger events.
* Changed object name reference strings to fully qualified import paths.
* Align Synapse failure states with operator expectations
Map FAILED/CANCELLED to error (consistent with ADF) so execute_complete
handles them correctly. Add parametrized round-trip test and update
newsfragment, operator docs and example DAG.
* Delete providers/microsoft/azure/newsfragments/63614.feature.rst
* Fix race condition in trigger to avoid cancelling already-completed
pipeline runs. Wrap cancellation in try/except to ensure timeout event is
always emitted. Add additional unit tests.
---------
Co-authored-by: Sameer Mesiah <[email protected]>
Co-authored-by: Kaxil Naik <[email protected]>
---
.../azure/docs/operators/azure_synapse.rst | 9 +
providers/microsoft/azure/provider.yaml | 3 +
.../providers/microsoft/azure/get_provider_info.py | 4 +
.../providers/microsoft/azure/operators/synapse.py | 91 ++++++--
.../providers/microsoft/azure/triggers/synapse.py | 186 +++++++++++++++++
.../azure/example_synapse_run_pipeline.py | 12 +-
.../unit/microsoft/azure/operators/test_synapse.py | 232 ++++++++++++++++++++-
.../unit/microsoft/azure/triggers/test_synapse.py | 230 ++++++++++++++++++++
scripts/ci/prek/known_airflow_exceptions.txt | 2 +-
9 files changed, 745 insertions(+), 24 deletions(-)
diff --git a/providers/microsoft/azure/docs/operators/azure_synapse.rst
b/providers/microsoft/azure/docs/operators/azure_synapse.rst
index 53b8e4df245..223f7258aae 100644
--- a/providers/microsoft/azure/docs/operators/azure_synapse.rst
+++ b/providers/microsoft/azure/docs/operators/azure_synapse.rst
@@ -52,6 +52,15 @@ The operator will Execute a Synapse Pipeline.
:start-after: [START howto_operator_azure_synapse_run_pipeline]
:end-before: [END howto_operator_azure_synapse_run_pipeline]
+Below is an example of using this operator to execute an Synapse pipeline with
a deferrable flag
+so that polling for the status of the pipeline run occurs on the Airflow
Triggerer.
+
+ .. exampleinclude::
/../tests/system/microsoft/azure/example_synapse_run_pipeline.py
+ :language: python
+ :dedent: 4
+ :start-after: [START
howto_operator_azure_synapse_run_pipeline_with_deferrable_flag]
+ :end-before: [END
howto_operator_azure_synapse_run_pipeline_with_deferrable_flag]
+
Reference
---------
diff --git a/providers/microsoft/azure/provider.yaml
b/providers/microsoft/azure/provider.yaml
index cebcdbebd19..a2a21cc270c 100644
--- a/providers/microsoft/azure/provider.yaml
+++ b/providers/microsoft/azure/provider.yaml
@@ -326,6 +326,9 @@ triggers:
- integration-name: Microsoft Azure Service Bus
python-modules:
- airflow.providers.microsoft.azure.triggers.message_bus
+ - integration-name: Microsoft Azure Synapse
+ python-modules:
+ - airflow.providers.microsoft.azure.triggers.synapse
queues:
-
airflow.providers.microsoft.azure.queues.asb.AzureServiceBusMessageQueueProvider
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
index 4ca9bc5c620..944b296d708 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
@@ -308,6 +308,10 @@ def get_provider_info():
"integration-name": "Microsoft Azure Service Bus",
"python-modules":
["airflow.providers.microsoft.azure.triggers.message_bus"],
},
+ {
+ "integration-name": "Microsoft Azure Synapse",
+ "python-modules":
["airflow.providers.microsoft.azure.triggers.synapse"],
+ },
],
"queues":
["airflow.providers.microsoft.azure.queues.asb.AzureServiceBusMessageQueueProvider"],
"transfers": [
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/synapse.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/synapse.py
index 5bffb6dc0d8..b0c171a13bd 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/synapse.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/synapse.py
@@ -16,6 +16,8 @@
# under the License.
from __future__ import annotations
+import time
+import warnings
from collections.abc import Sequence
from functools import cached_property
from typing import TYPE_CHECKING, Any
@@ -27,6 +29,7 @@ from airflow.providers.common.compat.sdk import (
BaseOperator,
BaseOperatorLink,
XCom,
+ conf,
)
from airflow.providers.microsoft.azure.hooks.synapse import (
AzureSynapseHook,
@@ -35,6 +38,9 @@ from airflow.providers.microsoft.azure.hooks.synapse import (
AzureSynapsePipelineRunStatus,
AzureSynapseSparkBatchRunStatus,
)
+from airflow.providers.microsoft.azure.triggers.synapse import (
+ AzureSynapsePipelineTrigger,
+)
if TYPE_CHECKING:
from azure.synapse.spark.models import SparkBatchJobOptions
@@ -194,11 +200,12 @@ class AzureSynapseRunPipelineOperator(BaseOperator):
:param parameters: Parameters of the pipeline run. These parameters are
referenced in a pipeline via
``@pipeline().parameters.parameterName`` and will be used only if the
``reference_pipeline_run_id`` is
not specified.
- :param timeout: Time in seconds to wait for a pipeline to reach a terminal
status for non-asynchronous
- waits. Used only if ``wait_for_termination`` is True.
+ :param timeout: Maximum time in seconds to wait for the pipeline run to
reach a terminal state. In synchronous
+ mode this controls how long the operator polls for completion. In
deferrable mode it determines the maximum
+ time the trigger will wait before emitting a timeout event. Used only
if ``wait_for_termination`` is True.
:param check_interval: Time in seconds to check on a pipeline run's status
for non-asynchronous waits.
Used only if ``wait_for_termination`` is True.
-
+ :param deferrable: Run operator in deferrable mode.
"""
template_fields: Sequence[str] = ("azure_synapse_conn_id",)
@@ -217,6 +224,7 @@ class AzureSynapseRunPipelineOperator(BaseOperator):
parameters: dict[str, Any] | None = None,
timeout: int = 60 * 60 * 24 * 7,
check_interval: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
**kwargs,
) -> None:
super().__init__(**kwargs)
@@ -230,6 +238,7 @@ class AzureSynapseRunPipelineOperator(BaseOperator):
self.parameters = parameters
self.timeout = timeout
self.check_interval = check_interval
+ self.deferrable = deferrable
@cached_property
def hook(self):
@@ -255,30 +264,72 @@ class AzureSynapseRunPipelineOperator(BaseOperator):
context["ti"].xcom_push(key="run_id", value=self.run_id)
if self.wait_for_termination:
- self.log.info("Waiting for pipeline run %s to terminate.",
self.run_id)
-
- if self.hook.wait_for_pipeline_run_status(
- run_id=self.run_id,
- expected_statuses=AzureSynapsePipelineRunStatus.SUCCEEDED,
- check_interval=self.check_interval,
- timeout=self.timeout,
- ):
- self.log.info("Pipeline run %s has completed successfully.",
self.run_id)
+ if not self.deferrable:
+ self.log.info("Waiting for pipeline run %s to terminate.",
self.run_id)
+ if self.hook.wait_for_pipeline_run_status(
+ run_id=self.run_id,
+ expected_statuses=AzureSynapsePipelineRunStatus.SUCCEEDED,
+ check_interval=self.check_interval,
+ timeout=self.timeout,
+ ):
+ self.log.info("Pipeline run %s has completed
successfully.", self.run_id)
+ else:
+ raise AzureSynapsePipelineRunException(
+ f"Pipeline run {self.run_id} has failed or has been
cancelled."
+ )
else:
- raise AzureSynapsePipelineRunException(
- f"Pipeline run {self.run_id} has failed or has been
cancelled."
+ end_time = time.time() + self.timeout
+
+ pipeline_run_status =
self.hook.get_pipeline_run_status(self.run_id)
+
+ if pipeline_run_status not in
AzureSynapsePipelineRunStatus.TERMINAL_STATUSES:
+ self.defer(
+ timeout=self.execution_timeout,
+ trigger=AzureSynapsePipelineTrigger(
+ azure_synapse_conn_id=self.azure_synapse_conn_id,
+
azure_synapse_workspace_dev_endpoint=self.azure_synapse_workspace_dev_endpoint,
+ run_id=self.run_id,
+ check_interval=self.check_interval,
+ end_time=end_time,
+ ),
+ method_name=self.execute_complete.__name__,
+ )
+
+ elif pipeline_run_status ==
AzureSynapsePipelineRunStatus.SUCCEEDED:
+ self.log.info("Pipeline run %s has completed
successfully.", self.run_id)
+
+ elif pipeline_run_status in
AzureSynapsePipelineRunStatus.FAILURE_STATES:
+ raise AzureSynapsePipelineRunException(
+ f"Pipeline run {self.run_id} has failed or has been
cancelled."
+ )
+ else:
+ if self.deferrable is True:
+ warnings.warn(
+ "Argument `wait_for_termination` is False and `deferrable`
is True, "
+ "hence `deferrable` parameter doesn't have any effect",
+ UserWarning,
+ stacklevel=2,
)
- def execute_complete(self, event: dict[str, str]) -> None:
+ def execute_complete(self, context: Context, event: dict[str, str] | None)
-> None:
"""
Return immediately - callback for when the trigger fires.
- Relies on trigger to throw an exception, otherwise it assumes
execution was successful.
+ The trigger communicates the terminal pipeline state through the event
payload.
"""
- if event:
- if event["status"] == "error":
- raise AirflowException(event["message"])
- self.log.info(event["message"])
+ if event is None:
+ raise AzureSynapsePipelineRunException("Trigger returned no
event.")
+
+ status = event.get("status")
+ message = event.get("message", "No message returned from trigger.")
+
+ if status == "success":
+ self.log.info(message)
+ return
+ if status == "error":
+ raise AzureSynapsePipelineRunException(message)
+
+ raise AzureSynapsePipelineRunException(f"Unexpected trigger event
received: {event}")
def on_kill(self) -> None:
if self.run_id:
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/synapse.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/synapse.py
new file mode 100644
index 00000000000..57c8d03613a
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/synapse.py
@@ -0,0 +1,186 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import AsyncIterator
+from typing import Any
+
+from azure.core.exceptions import ServiceRequestError
+
+from airflow.providers.microsoft.azure.hooks.synapse import (
+ AzureSynapsePipelineAsyncHook,
+ AzureSynapsePipelineRunStatus,
+)
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class AzureSynapsePipelineTrigger(BaseTrigger):
+ """
+ Trigger when an Azure Synapse pipeline run reaches a terminal state.
+
+ If an unexpected exception occurs while polling, the trigger attempts to
cancel
+ the pipeline run to avoid leaving orphaned executions.
+
+ :param run_id: Pipeline run identifier.
+ :param azure_synapse_conn_id: Azure Synapse connection id.
+ :param azure_synapse_workspace_dev_endpoint: Synapse workspace dev
endpoint.
+ :param end_time: Epoch timestamp when the trigger should timeout.
+ :param check_interval: Poll interval in seconds.
+ """
+
+ def __init__(
+ self,
+ run_id: str,
+ azure_synapse_conn_id: str,
+ azure_synapse_workspace_dev_endpoint: str,
+ end_time: float,
+ check_interval: int = 60,
+ ):
+ super().__init__()
+
+ self.run_id = run_id
+ self.azure_synapse_conn_id = azure_synapse_conn_id
+ self.azure_synapse_workspace_dev_endpoint =
azure_synapse_workspace_dev_endpoint
+ self.end_time = end_time
+ self.check_interval = check_interval
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize trigger arguments and classpath."""
+ return (
+ f"{self.__class__.__module__}.{self.__class__.__name__}",
+ {
+ "run_id": self.run_id,
+ "azure_synapse_conn_id": self.azure_synapse_conn_id,
+ "azure_synapse_workspace_dev_endpoint":
self.azure_synapse_workspace_dev_endpoint,
+ "end_time": self.end_time,
+ "check_interval": self.check_interval,
+ },
+ )
+
+ def _build_trigger_event(self, pipeline_status: str) -> TriggerEvent |
None:
+ """
+ Convert Synapse pipeline status to TriggerEvent.
+
+ Returns None if the pipeline is still running.
+ """
+ if pipeline_status == AzureSynapsePipelineRunStatus.SUCCEEDED:
+ return TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Pipeline run {self.run_id} succeeded.",
+ "run_id": self.run_id,
+ }
+ )
+
+ if pipeline_status in AzureSynapsePipelineRunStatus.FAILURE_STATES:
+ return TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Pipeline run {self.run_id} finished with
state {pipeline_status}.",
+ "run_id": self.run_id,
+ }
+ )
+
+ return None
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ """Poll Synapse pipeline status until completion or timeout."""
+ async with AzureSynapsePipelineAsyncHook(
+ azure_synapse_conn_id=self.azure_synapse_conn_id,
+
azure_synapse_workspace_dev_endpoint=self.azure_synapse_workspace_dev_endpoint,
+ ) as hook:
+ try:
+ executed_after_token_refresh = True
+
+ while time.time() < self.end_time:
+ try:
+ pipeline_status = await
hook.get_pipeline_run_status(self.run_id)
+
+ executed_after_token_refresh = True
+
+ event = self._build_trigger_event(pipeline_status)
+
+ if event:
+ yield event
+ return
+
+ self.log.info(
+ "Pipeline %s not finished yet (state=%s). Sleeping
%s seconds.",
+ self.run_id,
+ pipeline_status,
+ self.check_interval,
+ )
+
+ await asyncio.sleep(self.check_interval)
+
+ except ServiceRequestError:
+ # Azure token may expire during long pipeline runs.
+ if not executed_after_token_refresh:
+ raise
+
+ await hook.refresh_conn()
+ executed_after_token_refresh = False
+
+ try:
+ pipeline_status = await
hook.get_pipeline_run_status(self.run_id)
+
+ except ServiceRequestError:
+ # Azure token may expire after timeout period has elapsed.
+ await hook.refresh_conn()
+ pipeline_status = await
hook.get_pipeline_run_status(self.run_id)
+
+ # Yielding terminal event in case pipeline reaches terminal
state
+ # i.e. FAILED/SUCCEEDED during last sleep.
+ event = self._build_trigger_event(pipeline_status)
+
+ if event:
+ yield event
+ return
+
+ try:
+ await hook.cancel_pipeline_run(self.run_id)
+ except Exception:
+ self.log.exception("Pipeline %s cancellation failed.",
self.run_id)
+ finally:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Timeout waiting for pipeline run
{self.run_id}.",
+ "run_id": self.run_id,
+ }
+ )
+
+ except Exception as e:
+ self.log.exception(e)
+
+ if self.run_id:
+ try:
+ self.log.info("Cancelling pipeline run %s",
self.run_id)
+ await hook.cancel_pipeline_run(self.run_id)
+ except Exception:
+ self.log.exception("Failed to cancel pipeline run %s",
self.run_id)
+
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": str(e),
+ "run_id": self.run_id,
+ }
+ )
diff --git
a/providers/microsoft/azure/tests/system/microsoft/azure/example_synapse_run_pipeline.py
b/providers/microsoft/azure/tests/system/microsoft/azure/example_synapse_run_pipeline.py
index 421ad86a13f..e456df50332 100644
---
a/providers/microsoft/azure/tests/system/microsoft/azure/example_synapse_run_pipeline.py
+++
b/providers/microsoft/azure/tests/system/microsoft/azure/example_synapse_run_pipeline.py
@@ -41,7 +41,17 @@ with DAG(
azure_synapse_workspace_dev_endpoint="azure_synapse_workspace_dev_endpoint",
)
# [END howto_operator_azure_synapse_run_pipeline]
- begin >> run_pipeline1
+
+ # [START howto_operator_azure_synapse_run_pipeline_with_deferrable_flag]
+ run_pipeline2 = AzureSynapseRunPipelineOperator(
+ task_id="run_pipeline2",
+ azure_synapse_conn_id="azure_synapse_connection",
+ pipeline_name="Pipeline 2",
+
azure_synapse_workspace_dev_endpoint="azure_synapse_workspace_dev_endpoint",
+ deferrable=True,
+ )
+ # [END howto_operator_azure_synapse_run_pipeline_with_deferrable_flag]
+ begin >> run_pipeline1 >> run_pipeline2
from tests_common.test_utils.watcher import watcher
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_synapse.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_synapse.py
index 418981c9364..f7c256f04c3 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_synapse.py
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_synapse.py
@@ -16,13 +16,18 @@
# under the License.
from __future__ import annotations
+import functools
+import time
from unittest import mock
from unittest.mock import MagicMock, patch
+import pendulum
import pytest
-from airflow.models import Connection
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.models import DAG, Connection
+from airflow.models.dagrun import DagRun
+from airflow.models.taskinstance import TaskInstance
+from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred
from airflow.providers.microsoft.azure.hooks.synapse import (
AzureSynapsePipelineHook,
AzureSynapsePipelineRunException,
@@ -33,8 +38,11 @@ from airflow.providers.microsoft.azure.operators.synapse
import (
AzureSynapseRunPipelineOperator,
AzureSynapseRunSparkBatchOperator,
)
+from airflow.providers.microsoft.azure.triggers.synapse import
AzureSynapsePipelineTrigger
from airflow.utils import timezone
+from airflow.utils.types import DagRunType
+from tests_common.test_utils.taskinstance import create_task_instance as
_create_task_instance
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
if AIRFLOW_V_3_0_PLUS:
@@ -337,3 +345,223 @@ class TestAzureSynapseRunPipelineOperator:
workspace_url =
"https://web.azuresynapse.net?workspace=%2Fsubscriptions%2Fspam-egg";
with pytest.raises(ValueError, match="Workspace expected at least 5
segments"):
AzureSynapsePipelineRunLink().get_fields_from_url(workspace_url=workspace_url)
+
+
[email protected]
+def create_task_instance(create_task_instance_of_operator, session):
+ def _create_task_instance(operator_class, **kwargs):
+ return functools.partial(
+ create_task_instance_of_operator,
+ session=session,
+ operator_class=operator_class,
+ dag_id="adhoc_airflow",
+ )(**kwargs)
+
+ return _create_task_instance
+
+
+class TestAzureSynapseRunPipelineOperatorWithDeferrable:
+ @pytest.fixture(autouse=True)
+ def setup_operator(self, dag_maker, create_task_instance):
+ """Fixture to set up the operator using create_task_instance."""
+ self.ti = create_task_instance(
+ operator_class=AzureSynapseRunPipelineOperator,
+ task_id="run_pipeline",
+ pipeline_name=PIPELINE_NAME,
+
azure_synapse_workspace_dev_endpoint=AZURE_SYNAPSE_WORKSPACE_DEV_ENDPOINT,
+ azure_synapse_conn_id=AZURE_SYNAPSE_CONN_ID,
+ deferrable=True,
+ )
+ self.task = dag_maker.dag.get_task(self.ti.task_id)
+
+ def create_context(self, task, dag=None):
+ if dag is None:
+ dag = DAG(dag_id="dag", schedule=None)
+
+ tzinfo = pendulum.timezone("UTC")
+ logical_date = timezone.datetime(2022, 1, 1, 1, 0, 0, tzinfo=tzinfo)
+
+ if AIRFLOW_V_3_0_PLUS:
+ dag_run = DagRun(
+ dag_id=dag.dag_id,
+ logical_date=logical_date,
+ run_id=DagRun.generate_run_id(
+ run_type=DagRunType.MANUAL,
+ logical_date=logical_date,
+ run_after=logical_date,
+ ),
+ )
+ else:
+ dag_run = DagRun(
+ dag_id=dag.dag_id,
+ execution_date=logical_date,
+ run_id=DagRun.generate_run_id(DagRunType.MANUAL, logical_date),
+ )
+
+ if AIRFLOW_V_3_0_PLUS:
+ task_instance = _create_task_instance(task=task,
dag_version_id=mock.MagicMock())
+ else:
+ task_instance = TaskInstance(task=task)
+
+ task_instance.dag_run = dag_run
+ task_instance.xcom_push = mock.Mock()
+
+ date_key = "logical_date" if AIRFLOW_V_3_0_PLUS else "execution_date"
+
+ return {
+ "dag": dag,
+ "ts": logical_date.isoformat(),
+ "task": task,
+ "ti": task_instance,
+ "task_instance": task_instance,
+ "run_id": dag_run.run_id,
+ "dag_run": dag_run,
+ "data_interval_end": logical_date,
+ date_key: logical_date,
+ }
+
+ @pytest.mark.db_test
+
@mock.patch("airflow.providers.microsoft.azure.operators.synapse.AzureSynapseRunPipelineOperator.defer")
+ @mock.patch(
+
"airflow.providers.microsoft.azure.hooks.synapse.AzureSynapsePipelineHook.get_pipeline_run_status",
+ return_value=AzureSynapsePipelineRunStatus.SUCCEEDED,
+ )
+
@mock.patch("airflow.providers.microsoft.azure.hooks.synapse.AzureSynapsePipelineHook.run_pipeline")
+ def test_synapse_run_pipeline_operator_async_succeeded_before_deferred(
+ self, mock_run_pipeline, mock_get_status, mock_defer
+ ):
+ class CreateRunResponse:
+ pass
+
+ CreateRunResponse.run_id = PIPELINE_RUN_RESPONSE["run_id"]
+ mock_run_pipeline.return_value = CreateRunResponse
+
+ self.task.execute(context=self.create_context(self.task))
+ assert not mock_defer.called
+
+ @pytest.mark.db_test
+ @pytest.mark.parametrize(
+ "status",
+ sorted(AzureSynapsePipelineRunStatus.FAILURE_STATES),
+ )
+
@mock.patch("airflow.providers.microsoft.azure.operators.synapse.AzureSynapseRunPipelineOperator.defer")
+ @mock.patch(
+
"airflow.providers.microsoft.azure.hooks.synapse.AzureSynapsePipelineHook.get_pipeline_run_status"
+ )
+
@mock.patch("airflow.providers.microsoft.azure.hooks.synapse.AzureSynapsePipelineHook.run_pipeline")
+ def test_synapse_run_pipeline_operator_async_error_before_deferred(
+ self, mock_run_pipeline, mock_get_status, mock_defer, status
+ ):
+ mock_get_status.return_value = status
+
+ class CreateRunResponse:
+ pass
+
+ CreateRunResponse.run_id = PIPELINE_RUN_RESPONSE["run_id"]
+ mock_run_pipeline.return_value = CreateRunResponse
+
+ with pytest.raises(AzureSynapsePipelineRunException):
+ self.task.execute(context=self.create_context(self.task))
+
+ assert not mock_defer.called
+
+ @pytest.mark.db_test
+ @pytest.mark.parametrize(
+ "status",
+ sorted(AzureSynapsePipelineRunStatus.INTERMEDIATE_STATES),
+ )
+ @mock.patch(
+
"airflow.providers.microsoft.azure.hooks.synapse.AzureSynapsePipelineHook.get_pipeline_run_status"
+ )
+
@mock.patch("airflow.providers.microsoft.azure.hooks.synapse.AzureSynapsePipelineHook.run_pipeline")
+ def test_synapse_run_pipeline_operator_async(self, mock_run_pipeline,
mock_get_status, status):
+ """Assert that AzureSynapseRunPipelineOperator(..., deferrable=True)
deferred."""
+
+ mock_get_status.return_value = status
+
+ class CreateRunResponse:
+ pass
+
+ CreateRunResponse.run_id = PIPELINE_RUN_RESPONSE["run_id"]
+ mock_run_pipeline.return_value = CreateRunResponse
+
+ with pytest.raises(TaskDeferred) as exc:
+ self.task.execute(context=self.create_context(self.task))
+
+ trigger = exc.value.trigger
+
+ assert isinstance(trigger, AzureSynapsePipelineTrigger)
+
+ assert trigger.run_id == PIPELINE_RUN_RESPONSE["run_id"]
+ assert trigger.check_interval == self.task.check_interval
+
+ @pytest.mark.db_test
+ def
test_synapse_run_pipeline_operator_async_execute_complete_success(self):
+ """Assert that execute_complete logs success message."""
+ with mock.patch.object(self.task.log, "info") as mock_log_info:
+ self.task.execute_complete(
+ context={},
+ event={
+ "status": "success",
+ "message": "success",
+ "run_id": PIPELINE_RUN_RESPONSE["run_id"],
+ },
+ )
+
+ mock_log_info.assert_called_with("success")
+
+ @pytest.mark.db_test
+ def test_synapse_run_pipeline_operator_async_execute_complete_fail(self):
+ """Assert that execute_complete raises exception on error."""
+ with pytest.raises(AzureSynapsePipelineRunException):
+ self.task.execute_complete(
+ context={},
+ event={
+ "status": "error",
+ "message": "error",
+ "run_id": PIPELINE_RUN_RESPONSE["run_id"],
+ },
+ )
+
+ @pytest.mark.db_test
+ def test_execute_complete_no_event(self):
+ with pytest.raises(AzureSynapsePipelineRunException, match="no event"):
+ self.task.execute_complete(context={}, event=None)
+
+ @pytest.mark.db_test
+ def test_execute_complete_unexpected_event(self):
+ with pytest.raises(AzureSynapsePipelineRunException,
match="Unexpected"):
+ self.task.execute_complete(
+ context={},
+ event={"status": "unknown", "message": "??"},
+ )
+
+ @pytest.mark.db_test
+ @pytest.mark.parametrize(
+ "pipeline_state",
+ [
+ AzureSynapsePipelineRunStatus.FAILED,
+ AzureSynapsePipelineRunStatus.CANCELLED,
+ ],
+ )
+ def test_failure_states_roundtrip(self, pipeline_state):
+ trigger = AzureSynapsePipelineTrigger(
+ run_id=PIPELINE_RUN_RESPONSE["run_id"],
+ azure_synapse_conn_id=AZURE_SYNAPSE_CONN_ID,
+
azure_synapse_workspace_dev_endpoint=AZURE_SYNAPSE_WORKSPACE_DEV_ENDPOINT,
+ end_time=time.time() + 100,
+ check_interval=1,
+ )
+
+ event = trigger._build_trigger_event(pipeline_state)
+
+ assert event is not None
+
+ payload = event.payload
+
+ assert payload["status"] == "error"
+
+ with pytest.raises(AzureSynapsePipelineRunException) as exc:
+ self.task.execute_complete(context={}, event=payload)
+
+ assert f"state {pipeline_state}" in str(exc.value)
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_synapse.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_synapse.py
new file mode 100644
index 00000000000..3c517426eaf
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_synapse.py
@@ -0,0 +1,230 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import time
+from unittest import mock
+
+import pytest
+from azure.core.exceptions import ServiceRequestError
+
+from airflow.providers.microsoft.azure.triggers.synapse import (
+ AzureSynapsePipelineRunStatus,
+ AzureSynapsePipelineTrigger,
+)
+from airflow.triggers.base import TriggerEvent
+
+AZURE_SYNAPSE_CONN_ID = "azure_synapse_default"
+AZURE_SYNAPSE_WORKSPACE_DEV_ENDPOINT = "azure_synapse_workspace_dev_endpoint"
+RUN_ID = "7f8c6c72-c093-11ec-a83d-0242ac120007"
+POKE_INTERVAL = 5
+AZ_PIPELINE_END_TIME = time.time() + 60 * 60 * 24 * 7
+MODULE = "airflow.providers.microsoft.azure"
+
+
+class TestAzureSynapsePipelineTrigger:
+ TRIGGER = AzureSynapsePipelineTrigger(
+ run_id=RUN_ID,
+ azure_synapse_conn_id=AZURE_SYNAPSE_CONN_ID,
+
azure_synapse_workspace_dev_endpoint=AZURE_SYNAPSE_WORKSPACE_DEV_ENDPOINT,
+ check_interval=POKE_INTERVAL,
+ end_time=AZ_PIPELINE_END_TIME,
+ )
+
+ def test_synapse_trigger_serialization(self):
+ classpath, kwargs = self.TRIGGER.serialize()
+
+ assert classpath ==
f"{MODULE}.triggers.synapse.AzureSynapsePipelineTrigger"
+
+ assert kwargs == {
+ "run_id": RUN_ID,
+ "azure_synapse_conn_id": AZURE_SYNAPSE_CONN_ID,
+ "azure_synapse_workspace_dev_endpoint":
AZURE_SYNAPSE_WORKSPACE_DEV_ENDPOINT,
+ "check_interval": POKE_INTERVAL,
+ "end_time": AZ_PIPELINE_END_TIME,
+ }
+
+ @pytest.mark.parametrize(
+ ("status", "expected"),
+ [
+ ("Succeeded", "success"),
+ ("Failed", "error"),
+ ("Cancelled", "error"),
+ ],
+ )
+ def test_build_trigger_event_terminal_states(self, status, expected):
+ event = self.TRIGGER._build_trigger_event(status)
+
+ assert event is not None
+ assert event.payload["status"] == expected
+ assert event.payload["run_id"] == RUN_ID
+
+ @pytest.mark.parametrize(
+ "status",
+ ["Queued", "InProgress", "Canceling"],
+ )
+ def test_build_trigger_event_non_terminal_states(self, status):
+ event = self.TRIGGER._build_trigger_event(status)
+ assert event is None
+
+ @pytest.mark.asyncio
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.get_pipeline_run_status")
+ async def test_trigger_run_success(self, mock_status):
+ mock_status.return_value = AzureSynapsePipelineRunStatus.SUCCEEDED
+
+ generator = self.TRIGGER.run()
+ actual = await generator.asend(None)
+
+ assert actual == TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Pipeline run {RUN_ID} succeeded.",
+ "run_id": RUN_ID,
+ }
+ )
+
+ @pytest.mark.asyncio
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.get_pipeline_run_status")
+ async def test_trigger_run_failed(self, mock_status):
+ mock_status.return_value = AzureSynapsePipelineRunStatus.FAILED
+
+ generator = self.TRIGGER.run()
+ actual = await generator.asend(None)
+
+ assert actual == TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Pipeline run {RUN_ID} finished with state
Failed.",
+ "run_id": RUN_ID,
+ }
+ )
+
+ @pytest.mark.asyncio
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.cancel_pipeline_run")
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.get_pipeline_run_status")
+ async def test_trigger_exception(self, mock_status, mock_cancel):
+ mock_status.side_effect = Exception("API failure")
+
+ events = [event async for event in self.TRIGGER.run()]
+
+ mock_cancel.assert_awaited_once_with(RUN_ID)
+
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "error",
+ "message": "API failure",
+ "run_id": RUN_ID,
+ }
+ )
+ ]
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.triggers.synapse.asyncio.sleep",
new_callable=mock.AsyncMock)
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.refresh_conn")
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.get_pipeline_run_status")
+ async def test_trigger_refresh_on_service_request_error(self, mock_status,
mock_refresh, _mock_sleep):
+ mock_status.side_effect = [
+ ServiceRequestError("token expired"),
+ AzureSynapsePipelineRunStatus.SUCCEEDED,
+ ]
+
+ events = [event async for event in self.TRIGGER.run()]
+
+ mock_refresh.assert_awaited_once()
+
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Pipeline run {RUN_ID} succeeded.",
+ "run_id": RUN_ID,
+ }
+ )
+ ]
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.triggers.synapse.time")
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.cancel_pipeline_run")
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.get_pipeline_run_status")
+ async def test_trigger_timeout_pipeline_already_succeeded(self,
mock_status, mock_cancel, mock_time):
+ mock_status.return_value = AzureSynapsePipelineRunStatus.SUCCEEDED
+
+ mock_time.time.return_value = AZ_PIPELINE_END_TIME + 60
+
+ events = [event async for event in self.TRIGGER.run()]
+
+ mock_cancel.assert_not_awaited()
+
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Pipeline run {RUN_ID} succeeded.",
+ "run_id": RUN_ID,
+ }
+ )
+ ]
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.triggers.synapse.time")
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.cancel_pipeline_run")
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.get_pipeline_run_status")
+ async def test_trigger_timeout_cancellation(self, mock_status,
mock_cancel, mock_time):
+
+ mock_status.return_value = AzureSynapsePipelineRunStatus.IN_PROGRESS
+
+ mock_time.time.return_value = AZ_PIPELINE_END_TIME + 60
+
+ events = [event async for event in self.TRIGGER.run()]
+
+ mock_cancel.assert_awaited_once_with(RUN_ID)
+
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Timeout waiting for pipeline run {RUN_ID}.",
+ "run_id": RUN_ID,
+ }
+ )
+ ]
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.triggers.synapse.time")
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.cancel_pipeline_run")
+
@mock.patch(f"{MODULE}.hooks.synapse.AzureSynapsePipelineAsyncHook.get_pipeline_run_status")
+ async def test_trigger_timeout_cancel_failure_still_yields_timeout(
+ self, mock_status, mock_cancel, mock_time
+ ):
+ mock_status.return_value = AzureSynapsePipelineRunStatus.IN_PROGRESS
+ mock_cancel.side_effect = Exception("cancel failed")
+
+ mock_time.time.return_value = AZ_PIPELINE_END_TIME + 60
+
+ events = [event async for event in self.TRIGGER.run()]
+
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Timeout waiting for pipeline run {RUN_ID}.",
+ "run_id": RUN_ID,
+ }
+ )
+ ]
diff --git a/scripts/ci/prek/known_airflow_exceptions.txt
b/scripts/ci/prek/known_airflow_exceptions.txt
index f093214991f..6175e6ad00c 100644
--- a/scripts/ci/prek/known_airflow_exceptions.txt
+++ b/scripts/ci/prek/known_airflow_exceptions.txt
@@ -358,7 +358,7 @@
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/contai
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/data_factory.py::1
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py::1
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/powerbi.py::4
-providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/synapse.py::2
+providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/synapse.py::1
providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/data_factory.py::1
providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/msgraph.py::1
providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/wasb.py::4
