This is an automated email from the ASF dual-hosted git repository.
taragolis 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 830b71b3be fix: try002 for provider amazon (#38789)
830b71b3be is described below
commit 830b71b3be66b970340e85cb8f4a297625c71ea1
Author: Sebastian Daum <[email protected]>
AuthorDate: Mon Apr 8 09:36:48 2024 +0200
fix: try002 for provider amazon (#38789)
---
airflow/providers/amazon/aws/operators/athena.py | 4 ++--
airflow/providers/amazon/aws/operators/emr.py | 2 +-
pyproject.toml | 3 ---
tests/providers/amazon/aws/operators/test_athena.py | 8 ++++----
4 files changed, 7 insertions(+), 10 deletions(-)
diff --git a/airflow/providers/amazon/aws/operators/athena.py
b/airflow/providers/amazon/aws/operators/athena.py
index 18fb165fc2..a9e09fb325 100644
--- a/airflow/providers/amazon/aws/operators/athena.py
+++ b/airflow/providers/amazon/aws/operators/athena.py
@@ -165,12 +165,12 @@ class AthenaOperator(AwsBaseOperator[AthenaHook]):
if query_status in AthenaHook.FAILURE_STATES:
error_message =
self.hook.get_state_change_reason(self.query_execution_id)
- raise Exception(
+ raise AirflowException(
f"Final state of Athena job is {query_status},
query_execution_id is "
f"{self.query_execution_id}. Error: {error_message}"
)
elif not query_status or query_status in
AthenaHook.INTERMEDIATE_STATES:
- raise Exception(
+ raise AirflowException(
f"Final state of Athena job is {query_status}. Max tries of
poll status exceeded, "
f"query_execution_id is {self.query_execution_id}."
)
diff --git a/airflow/providers/amazon/aws/operators/emr.py
b/airflow/providers/amazon/aws/operators/emr.py
index 01e1567eab..ec1be30f91 100644
--- a/airflow/providers/amazon/aws/operators/emr.py
+++ b/airflow/providers/amazon/aws/operators/emr.py
@@ -576,7 +576,7 @@ class EmrContainerOperator(BaseOperator):
stacklevel=2,
)
if max_polling_attempts and max_polling_attempts != max_tries:
- raise Exception("max_polling_attempts must be the same value
as max_tries")
+ raise ValueError("max_polling_attempts must be the same value
as max_tries")
else:
self.max_polling_attempts = max_tries
diff --git a/pyproject.toml b/pyproject.toml
index c3711b8548..e7284bc3ba 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -379,9 +379,6 @@ combine-as-imports = true
"tests/providers/snowflake/operators/test_snowflake_sql.py" = ["E402"]
# All the providers modules which do not follow TRY002 yet
-# amazon
-"airflow/providers/amazon/aws/operators/athena.py" = ["TRY002"]
-"airflow/providers/amazon/aws/operators/emr.py" = ["TRY002"]
# cncf.kubernetes
"airflow/providers/cncf/kubernetes/operators/pod.py" = ["TRY002"]
# common.sql
diff --git a/tests/providers/amazon/aws/operators/test_athena.py
b/tests/providers/amazon/aws/operators/test_athena.py
index 4c7e564ccc..d4ccf521a7 100644
--- a/tests/providers/amazon/aws/operators/test_athena.py
+++ b/tests/providers/amazon/aws/operators/test_athena.py
@@ -29,7 +29,7 @@ from openlineage.client.facet import (
)
from openlineage.client.run import Dataset
-from airflow.exceptions import TaskDeferred
+from airflow.exceptions import AirflowException, TaskDeferred
from airflow.models import DAG, DagRun, TaskInstance
from airflow.providers.amazon.aws.hooks.athena import AthenaHook
from airflow.providers.amazon.aws.operators.athena import AthenaOperator
@@ -180,7 +180,7 @@ class TestAthenaOperator:
mock_check_query_status,
mock_get_state_change_reason,
):
- with pytest.raises(Exception):
+ with pytest.raises(AirflowException):
self.athena.execute({})
mock_run_query.assert_called_once_with(
MOCK_DATA["query"],
@@ -195,7 +195,7 @@ class TestAthenaOperator:
@mock.patch.object(AthenaHook, "run_query", return_value=ATHENA_QUERY_ID)
@mock.patch.object(AthenaHook, "get_conn")
def test_hook_run_cancelled_query(self, mock_conn, mock_run_query,
mock_check_query_status):
- with pytest.raises(Exception):
+ with pytest.raises(AirflowException):
self.athena.execute({})
mock_run_query.assert_called_once_with(
MOCK_DATA["query"],
@@ -209,7 +209,7 @@ class TestAthenaOperator:
@mock.patch.object(AthenaHook, "run_query", return_value=ATHENA_QUERY_ID)
@mock.patch.object(AthenaHook, "get_conn")
def test_hook_run_failed_query_with_max_tries(self, mock_conn,
mock_run_query, mock_check_query_status):
- with pytest.raises(Exception):
+ with pytest.raises(AirflowException):
self.athena.execute({})
mock_run_query.assert_called_once_with(
MOCK_DATA["query"],
