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 a4320d95df4 Replace generic AirflowException in AzureBatchOperator
input validation (#71981)
a4320d95df4 is described below
commit a4320d95df4e5646027dba43cf6bb731577522b7
Author: Baha Bouali <[email protected]>
AuthorDate: Wed Sep 9 00:58:18 2026 +0100
Replace generic AirflowException in AzureBatchOperator input validation
(#71981)
Signed-off-by: baha-bouali <[email protected]>
---
generated/known_airflow_exceptions.txt | 1 -
.../providers/microsoft/azure/exceptions.py | 60 +++++++++++
.../providers/microsoft/azure/operators/batch.py | 30 ++++--
.../unit/microsoft/azure/operators/test_batch.py | 110 ++++++++++++++++++++-
.../tests/unit/microsoft/azure/test_exceptions.py | 61 ++++++++++++
5 files changed, 247 insertions(+), 15 deletions(-)
diff --git a/generated/known_airflow_exceptions.txt
b/generated/known_airflow_exceptions.txt
index fa58bca8214..fd54d247222 100644
--- a/generated/known_airflow_exceptions.txt
+++ b/generated/known_airflow_exceptions.txt
@@ -346,7 +346,6 @@
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_facto
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_lake.py::1
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py::3
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/wasb.py::2
-providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/batch.py::8
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py::10
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
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
new file mode 100644
index 00000000000..7235826189a
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
@@ -0,0 +1,60 @@
+# 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.
+
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# ... (full ASF header, copy verbatim from any file in the repo)
+# under the License.
+# Note: Any AirflowException raised is expected to cause the TaskInstance
+# to be marked in an ERROR state
+"""Exceptions used by Microsoft Azure Provider."""
+
+from __future__ import annotations
+
+from airflow.providers.common.compat.sdk import AirflowException
+
+
+class AzureBatchVmPublisherMissingError(AirflowException):
+ """Raised when vm_publisher is not provided."""
+
+
+class AzureBatchLatestImageSpecIncompleteError(AirflowException):
+ """Raised when use_latest_verified_vm_image_and_sku is requested without a
complete image spec."""
+
+
+class AzureBatchVmImageSpecIncompleteError(AirflowException):
+ """Raised when vm_publisher is provided without vm_sku, vm_offer and
vm_node_agent_sku_id."""
+
+
+class AzureBatchPoolSizingMissingError(AirflowException):
+ """Raised when neither target_dedicated_nodes nor enable_auto_scale is
provided."""
+
+
+class AzureBatchPoolSizingConflictError(AirflowException):
+ """Raised when enable_auto_scale is combined with explicit node counts."""
+
+
+class AzureBatchAutoScaleFormulaMissingError(AirflowException):
+ """Raised when enable_auto_scale is set without auto_scale_formula."""
+
+
+class AzureBatchJobPreparationTaskMissingError(AirflowException):
+ """Raised when batch_job_release_task is provided without
batch_job_preparation_task."""
+
+
+class AzureBatchRequiredParametersMissingError(AirflowException):
+ """Raised when one or more required pool, job or task parameters are
missing."""
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/batch.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/batch.py
index 8210a18012a..815dd1247a6 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/batch.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/batch.py
@@ -24,7 +24,17 @@ from typing import TYPE_CHECKING, Any
from azure.batch import models as batch_models
-from airflow.providers.common.compat.sdk import AirflowException,
BaseOperator, conf
+from airflow.providers.common.compat.sdk import BaseOperator, conf
+from airflow.providers.microsoft.azure.exceptions import (
+ AzureBatchAutoScaleFormulaMissingError,
+ AzureBatchJobPreparationTaskMissingError,
+ AzureBatchLatestImageSpecIncompleteError,
+ AzureBatchPoolSizingConflictError,
+ AzureBatchPoolSizingMissingError,
+ AzureBatchRequiredParametersMissingError,
+ AzureBatchVmImageSpecIncompleteError,
+ AzureBatchVmPublisherMissingError,
+)
from airflow.providers.microsoft.azure.hooks.batch import AzureBatchHook
from airflow.providers.microsoft.azure.triggers.batch import AzureBatchTrigger
@@ -188,18 +198,18 @@ class AzureBatchOperator(BaseOperator):
def _check_inputs(self) -> Any:
if not self.vm_publisher:
- raise AirflowException("You must specify vm_publisher")
+ raise AzureBatchVmPublisherMissingError("You must specify
vm_publisher")
if self.use_latest_image:
if not self.vm_publisher or not self.vm_offer:
- raise AirflowException(
+ raise AzureBatchLatestImageSpecIncompleteError(
f"If use_latest_image_and_sku is set to True then the
parameters vm_publisher, "
f"vm_offer, must all be set. "
f"Found vm_publisher={self.vm_publisher},
vm_offer={self.vm_offer}"
)
if self.vm_publisher:
if not all([self.vm_sku, self.vm_offer,
self.vm_node_agent_sku_id]):
- raise AirflowException(
+ raise AzureBatchVmImageSpecIncompleteError(
"If vm_publisher is set, then the parameters vm_sku,
vm_offer,"
"vm_node_agent_sku_id must be set. Found "
f"vm_publisher={self.vm_publisher},
vm_offer={self.vm_offer} "
@@ -208,21 +218,23 @@ class AzureBatchOperator(BaseOperator):
)
if not self.target_dedicated_nodes and not self.enable_auto_scale:
- raise AirflowException(
+ raise AzureBatchPoolSizingMissingError(
"Either target_dedicated_nodes or enable_auto_scale must be
set. None was set"
)
if self.enable_auto_scale:
if self.target_dedicated_nodes or self.target_low_priority_nodes:
- raise AirflowException(
+ raise AzureBatchPoolSizingConflictError(
f"If enable_auto_scale is set, then the parameters
target_dedicated_nodes and "
f"target_low_priority_nodes must not be set. Found "
f"target_dedicated_nodes={self.target_dedicated_nodes}, "
f"target_low_priority_nodes={self.target_low_priority_nodes}"
)
if not self.auto_scale_formula:
- raise AirflowException("The auto_scale_formula is required
when enable_auto_scale is set")
+ raise AzureBatchAutoScaleFormulaMissingError(
+ "The auto_scale_formula is required when enable_auto_scale
is set"
+ )
if self.batch_job_release_task and not self.batch_job_preparation_task:
- raise AirflowException(
+ raise AzureBatchJobPreparationTaskMissingError(
"A batch_job_release_task cannot be specified without also "
" specifying a batch_job_preparation_task for the Job."
)
@@ -235,7 +247,7 @@ class AzureBatchOperator(BaseOperator):
self.batch_task_command_line,
]
):
- raise AirflowException(
+ raise AzureBatchRequiredParametersMissingError(
"Some required parameters are missing.Please you must set all
the required parameters. "
)
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_batch.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_batch.py
index 47a84521ead..39e97b3736b 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_batch.py
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_batch.py
@@ -23,7 +23,17 @@ from unittest import mock
import pytest
from airflow.models import Connection
-from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred
+from airflow.providers.common.compat.sdk import TaskDeferred
+from airflow.providers.microsoft.azure.exceptions import (
+ AzureBatchAutoScaleFormulaMissingError,
+ AzureBatchJobPreparationTaskMissingError,
+ AzureBatchLatestImageSpecIncompleteError,
+ AzureBatchPoolSizingConflictError,
+ AzureBatchPoolSizingMissingError,
+ AzureBatchRequiredParametersMissingError,
+ AzureBatchVmImageSpecIncompleteError,
+ AzureBatchVmPublisherMissingError,
+)
from airflow.providers.microsoft.azure.hooks.batch import AzureBatchHook
from airflow.providers.microsoft.azure.operators.batch import
AzureBatchOperator
from airflow.providers.microsoft.azure.triggers.batch import AzureBatchTrigger
@@ -145,6 +155,68 @@ class TestAzureBatchOperator:
target_dedicated_nodes=1,
timeout=2,
)
+ self.operator_latest_image_incomplete = AzureBatchOperator(
+ task_id=TASK_ID,
+ batch_pool_id=BATCH_POOL_ID,
+ batch_pool_vm_size=BATCH_VM_SIZE,
+ batch_job_id=BATCH_JOB_ID,
+ batch_task_id=BATCH_TASK_ID,
+ vm_publisher=self.test_vm_publisher,
+ vm_sku=self.test_vm_sku,
+ vm_node_agent_sku_id=self.test_node_agent_sku,
+ batch_task_command_line="echo hello",
+ azure_batch_conn_id=self.test_vm_conn_id,
+ use_latest_verified_vm_image_and_sku=True,
+ target_dedicated_nodes=1,
+ timeout=2,
+ )
+ self.operator_vm_image_incomplete = AzureBatchOperator(
+ task_id=TASK_ID,
+ batch_pool_id=BATCH_POOL_ID,
+ batch_pool_vm_size=BATCH_VM_SIZE,
+ batch_job_id=BATCH_JOB_ID,
+ batch_task_id=BATCH_TASK_ID,
+ vm_publisher=self.test_vm_publisher,
+ vm_offer=self.test_vm_offer,
+ vm_node_agent_sku_id=self.test_node_agent_sku,
+ batch_task_command_line="echo hello",
+ azure_batch_conn_id=self.test_vm_conn_id,
+ target_dedicated_nodes=1,
+ timeout=2,
+ )
+ self.operator_sizing_conflict = AzureBatchOperator(
+ task_id=TASK_ID,
+ batch_pool_id=BATCH_POOL_ID,
+ batch_pool_vm_size=BATCH_VM_SIZE,
+ batch_job_id=BATCH_JOB_ID,
+ batch_task_id=BATCH_TASK_ID,
+ vm_publisher=self.test_vm_publisher,
+ vm_offer=self.test_vm_offer,
+ vm_sku=self.test_vm_sku,
+ vm_node_agent_sku_id=self.test_node_agent_sku,
+ batch_task_command_line="echo hello",
+ azure_batch_conn_id=self.test_vm_conn_id,
+ enable_auto_scale=True,
+ auto_scale_formula=FORMULA,
+ target_dedicated_nodes=1,
+ timeout=2,
+ )
+ self.operator_release_without_preparation = AzureBatchOperator(
+ task_id=TASK_ID,
+ batch_pool_id=BATCH_POOL_ID,
+ batch_pool_vm_size=BATCH_VM_SIZE,
+ batch_job_id=BATCH_JOB_ID,
+ batch_task_id=BATCH_TASK_ID,
+ vm_publisher=self.test_vm_publisher,
+ vm_offer=self.test_vm_offer,
+ vm_sku=self.test_vm_sku,
+ vm_node_agent_sku_id=self.test_node_agent_sku,
+ batch_task_command_line="echo hello",
+ azure_batch_conn_id=self.test_vm_conn_id,
+ target_dedicated_nodes=1,
+ batch_job_release_task=mock.MagicMock(),
+ timeout=2,
+ )
@mock.patch.object(AzureBatchHook, "wait_for_all_node_state")
def test_execute_without_failures(self, wait_mock):
@@ -166,7 +238,7 @@ class TestAzureBatchOperator:
def test_execute_with_failures(self, wait_mock):
wait_mock.return_value = True
self.operator.batch_pool_id = None
- with pytest.raises(AirflowException):
+ with pytest.raises(AzureBatchRequiredParametersMissingError):
self.operator.execute(None)
@mock.patch.object(AzureBatchHook, "wait_for_all_node_state")
@@ -181,7 +253,7 @@ class TestAzureBatchOperator:
@mock.patch.object(AzureBatchHook, "wait_for_all_node_state")
def test_operator_fails_no_dedicated_nodes_or_autoscale(self, wait_mock):
wait_mock.return_value = True
- with pytest.raises(AirflowException) as ctx:
+ with pytest.raises(AzureBatchPoolSizingMissingError) as ctx:
self.operator_fail.execute(None)
assert (
str(ctx.value) == "Either target_dedicated_nodes or
enable_auto_scale must be set. None was set"
@@ -190,17 +262,45 @@ class TestAzureBatchOperator:
@mock.patch.object(AzureBatchHook, "wait_for_all_node_state")
def test_operator_fails_no_formula(self, wait_mock):
wait_mock.return_value = True
- with pytest.raises(AirflowException) as ctx:
+ with pytest.raises(AzureBatchAutoScaleFormulaMissingError) as ctx:
self.operator_no_formula.execute(None)
assert str(ctx.value) == "The auto_scale_formula is required when
enable_auto_scale is set"
@mock.patch.object(AzureBatchHook, "wait_for_all_node_state")
def test_operator_fails_invalid_args(self, wait_mock):
wait_mock.return_value = True
- with pytest.raises(AirflowException) as ctx:
+ with pytest.raises(AzureBatchVmPublisherMissingError) as ctx:
self.operator_invalid.execute(None)
assert str(ctx.value) == "You must specify vm_publisher"
+ @mock.patch.object(AzureBatchHook, "wait_for_all_node_state")
+ def test_operator_fails_latest_image_without_vm_offer(self, wait_mock):
+ wait_mock.return_value = True
+ with pytest.raises(AzureBatchLatestImageSpecIncompleteError) as ctx:
+ self.operator_latest_image_incomplete.execute(None)
+ assert "use_latest_image_and_sku" in str(ctx.value)
+
+ @mock.patch.object(AzureBatchHook, "wait_for_all_node_state")
+ def test_operator_fails_vm_publisher_without_vm_sku(self, wait_mock):
+ wait_mock.return_value = True
+ with pytest.raises(AzureBatchVmImageSpecIncompleteError) as ctx:
+ self.operator_vm_image_incomplete.execute(None)
+ assert "vm_sku" in str(ctx.value)
+
+ @mock.patch.object(AzureBatchHook, "wait_for_all_node_state")
+ def test_operator_fails_auto_scale_with_dedicated_nodes(self, wait_mock):
+ wait_mock.return_value = True
+ with pytest.raises(AzureBatchPoolSizingConflictError) as ctx:
+ self.operator_sizing_conflict.execute(None)
+ assert "must not be set" in str(ctx.value)
+
+ @mock.patch.object(AzureBatchHook, "wait_for_all_node_state")
+ def test_operator_fails_release_task_without_preparation_task(self,
wait_mock):
+ wait_mock.return_value = True
+ with pytest.raises(AzureBatchJobPreparationTaskMissingError) as ctx:
+ self.operator_release_without_preparation.execute(None)
+ assert "batch_job_preparation_task" in str(ctx.value)
+
def test_cleaning_works(self):
self.operator.clean_up(job_id="myjob")
self.batch_client.begin_delete_job.assert_called_once_with("myjob")
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/test_exceptions.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/test_exceptions.py
new file mode 100644
index 00000000000..545b5f6664a
--- /dev/null
+++ b/providers/microsoft/azure/tests/unit/microsoft/azure/test_exceptions.py
@@ -0,0 +1,61 @@
+# 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 pytest
+
+from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.microsoft.azure.exceptions import (
+ AzureBatchAutoScaleFormulaMissingError,
+ AzureBatchJobPreparationTaskMissingError,
+ AzureBatchLatestImageSpecIncompleteError,
+ AzureBatchPoolSizingConflictError,
+ AzureBatchPoolSizingMissingError,
+ AzureBatchRequiredParametersMissingError,
+ AzureBatchVmImageSpecIncompleteError,
+ AzureBatchVmPublisherMissingError,
+)
+
+ALL_EXCEPTIONS = [
+ AzureBatchAutoScaleFormulaMissingError,
+ AzureBatchJobPreparationTaskMissingError,
+ AzureBatchLatestImageSpecIncompleteError,
+ AzureBatchPoolSizingConflictError,
+ AzureBatchPoolSizingMissingError,
+ AzureBatchRequiredParametersMissingError,
+ AzureBatchVmImageSpecIncompleteError,
+ AzureBatchVmPublisherMissingError,
+]
+
+
[email protected]("exception_class", ALL_EXCEPTIONS)
+def test_inherits_from_airflow_exception(exception_class):
+ assert issubclass(exception_class, AirflowException)
+
+
[email protected]("exception_class", ALL_EXCEPTIONS)
+def test_can_be_raised_with_message(exception_class):
+ with pytest.raises(exception_class, match="boom"):
+ raise exception_class("boom")
+
+
[email protected]("exception_class", ALL_EXCEPTIONS)
+def test_caught_as_airflow_exception(exception_class):
+ with pytest.raises(AirflowException):
+ raise exception_class("boom")
