diff --git a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py index 1f9d1b0c5b906..b4a8a91e31d19 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py +++ b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py @@ -37,7 +37,6 @@ from sqlalchemy import create_engine from airflow.providers.common.compat.sdk import ( - AirflowException, AirflowOptionalProviderFeatureException, Connection, conf, @@ -586,7 +585,7 @@ def get_private_key(self) -> PrivateKeyTypes | None: p_key = None if private_key_content and private_key_file: - raise AirflowException( + raise ValueError( "The private_key_file and private_key_content extra fields are mutually exclusive. " "Please remove one." ) diff --git a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py index 5479e5a8727d6..bd0ea6b6485d9 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py +++ b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py @@ -37,7 +37,6 @@ ) from airflow.exceptions import AirflowProviderDeprecationWarning -from airflow.providers.common.compat.sdk import AirflowException from airflow.providers.common.sql.hooks.lineage import send_sql_hook_lineage from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook from airflow.providers.snowflake.utils.sql_api_generate_jwt import JWTGenerator @@ -161,7 +160,7 @@ def execute_query( headers = self.get_headers() sql_is_multi_stmt = ";" in sql.strip() if not isinstance(bindings, dict) and bindings is not None: - raise AirflowException("Bindings should be a dictionary or None.") + raise TypeError("Bindings should be a dictionary or None.") if bindings and sql_is_multi_stmt: self.log.warning( "Bindings are not supported for multi-statement queries. Bindings will be ignored." @@ -191,7 +190,7 @@ def execute_query( elif "statementHandle" in json_response: self.query_ids.append(json_response["statementHandle"]) else: - raise AirflowException("No statementHandle/statementHandles present in response") + raise RuntimeError("No statementHandle/statementHandles present in response") # Send Hook Level Lineage len_query_ids = len(self.query_ids) @@ -238,7 +237,7 @@ def get_headers(self) -> dict[str, Any]: if conn_config.get("authenticator") == "programmatic_access_token": pat = conn_config.get("password") if not pat: - raise AirflowException( + raise ValueError( "Programmatic Access Token (PAT) authentication requires the connection password " "field to contain the PAT token value." ) diff --git a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py index 0cd5929a1ba27..2b02d259e7152 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py @@ -23,7 +23,7 @@ from functools import cached_property from typing import TYPE_CHECKING, Any, SupportsAbs, cast -from airflow.providers.common.compat.sdk import AirflowException, conf +from airflow.providers.common.compat.sdk import conf from airflow.providers.common.sql.operators.sql import ( SQLCheckOperator, SQLExecuteQueryOperator, @@ -443,7 +443,7 @@ def execute(self, context: Context) -> None: if statement_status.get("status") == "success": succeeded_query_ids.append(query_id) else: - raise AirflowException(f"{statement_status.get('status')}: {statement_status.get('message')}") + raise RuntimeError(f"{statement_status.get('status')}: {statement_status.get('message')}") if len(self.query_ids) == len(succeeded_query_ids): self.log.info("%s completed successfully.", self.task_id) @@ -465,7 +465,7 @@ def execute(self, context: Context) -> None: while True: statement_status = self.poll_on_queries() if statement_status["error"]: - raise AirflowException(statement_status["error"]) + raise RuntimeError(str(statement_status["error"])) if not statement_status["running"]: break @@ -509,7 +509,7 @@ def execute_complete(self, context: Context, event: dict[str, str | list[str]] | if event: if "status" in event and event["status"] == "error": msg = f"{event['status']}: {event['message']}" - raise AirflowException(msg) + raise RuntimeError(msg) if "status" in event and event["status"] == "success": self.query_ids = cast("list[str]", event["statement_query_ids"]) self._hook.check_query_output(self.query_ids) diff --git a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py index 062bee3a67f32..7b2d397e8ab7e 100644 --- a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py +++ b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py @@ -35,7 +35,6 @@ from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.models import Connection -from airflow.providers.common.compat.sdk import AirflowException from airflow.providers.snowflake.hooks.snowflake_sql_api import SnowflakeSqlApiHook if TYPE_CHECKING: @@ -391,14 +390,14 @@ def test_execute_query_exception_without_statement_handle( mock_requests, ): """ - Test execute_query method by mocking the exception response and raise airflow exception + Test execute_query method by mocking the exception response and raise RuntimeError without statementHandle in the response """ # status_code, json payload without statementHandle mock_make_api_call.return_value = (None, {"foo": "bar"}) hook = SnowflakeSqlApiHook("mock_conn_id") - with pytest.raises(AirflowException): + with pytest.raises(RuntimeError): hook.execute_query(sql, statement_count) @pytest.mark.parametrize( @@ -517,11 +516,11 @@ def test_get_headers_should_support_pat(self, mock_conn_param): @mock.patch(f"{HOOK_PATH}._get_conn_params") def test_get_headers_pat_raises_when_password_missing(self, mock_conn_param): - """Test get_headers raises AirflowException when PAT authenticator is set but password is empty.""" + """Test get_headers raises ValueError when PAT authenticator is set but password is empty.""" conn_params_pat_no_password = {**CONN_PARAMS_PAT, "password": ""} mock_conn_param.return_value = conn_params_pat_no_password hook = SnowflakeSqlApiHook(snowflake_conn_id="mock_conn_id") - with pytest.raises(AirflowException, match="Programmatic Access Token"): + with pytest.raises(ValueError, match="Programmatic Access Token"): hook.get_headers() @mock.patch("airflow.providers.snowflake.hooks.snowflake.HTTPBasicAuth") @@ -629,7 +628,7 @@ def test_get_private_key_raise_exception( "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri() ), pytest.raises( - AirflowException, + ValueError, match="The private_key_file and private_key_content extra fields are mutually " "exclusive. Please remove one.", ), diff --git a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py index 400b813d64be1..cca046f041548 100644 --- a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py +++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py @@ -27,7 +27,7 @@ from airflow.models.dag import DAG from airflow.models.dagrun import DagRun from airflow.models.taskinstance import TaskInstance -from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred +from airflow.providers.common.compat.sdk import TaskDeferred from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator from airflow.providers.snowflake.operators.snowflake import ( SnowflakeCheckOperator, @@ -339,7 +339,7 @@ def test_snowflake_sql_api_to_fails_when_one_query_fails( ) mock_execute_query.return_value = ["uuid1", "uuid2"] mock_get_sql_api_query_status.side_effect = [{"status": "error"}, {"status": "success"}] - with pytest.raises(AirflowException): + with pytest.raises(RuntimeError): operator.execute(context=None) @pytest.mark.parametrize( @@ -373,7 +373,7 @@ def test_snowflake_sql_api_execute_operator_async( ) def test_snowflake_sql_api_execute_complete_failure(self): - """Test SnowflakeSqlApiOperator raise AirflowException of error event""" + """Test SnowflakeSqlApiOperator raise RuntimeError of error event""" operator = SnowflakeSqlApiOperator( task_id=TASK_ID, @@ -382,7 +382,7 @@ def test_snowflake_sql_api_execute_complete_failure(self): statement_count=4, deferrable=True, ) - with pytest.raises(AirflowException): + with pytest.raises(RuntimeError): operator.execute_complete( context=None, event={"status": "error", "message": "Test failure message", "type": "FAILED_WITH_ERROR"}, @@ -469,7 +469,7 @@ def test_snowflake_sql_api_execute_operator_failed_before_defer( ) mock_execute_query.return_value = ["uuid1"] mock_get_sql_api_query_status.side_effect = [{"status": "error"}] - with pytest.raises(AirflowException): + with pytest.raises(RuntimeError): operator.execute(create_context(operator)) assert not mock_defer.called @@ -551,7 +551,7 @@ def test_snowflake_sql_api_execute_operator_polling_failed( self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output ): """ - Tests that the execute method raises AirflowException if any query fails during polling + Tests that the execute method raises RuntimeError if any query fails during polling when ``deferrable=False`` """ operator = SnowflakeSqlApiOperator( @@ -572,7 +572,7 @@ def test_snowflake_sql_api_execute_operator_polling_failed( {"status": "error"}, ] - with pytest.raises(AirflowException): + with pytest.raises(RuntimeError): operator.execute(context=None) mock_check_query_output.assert_not_called() diff --git a/scripts/ci/prek/known_airflow_exceptions.txt b/scripts/ci/prek/known_airflow_exceptions.txt index 6175e6ad00c00..b9dd36f80c295 100644 --- a/scripts/ci/prek/known_airflow_exceptions.txt +++ b/scripts/ci/prek/known_airflow_exceptions.txt @@ -392,10 +392,6 @@ providers/slack/src/airflow/providers/slack/transfers/sql_to_slack.py::1 providers/slack/src/airflow/providers/slack/transfers/sql_to_slack_webhook.py::3 providers/smtp/src/airflow/providers/smtp/hooks/smtp.py::16 providers/smtp/src/airflow/providers/smtp/operators/smtp.py::3 -providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py::1 -providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py::3 -providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py::3 -providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py::1 providers/ssh/src/airflow/providers/ssh/hooks/ssh.py::5 providers/ssh/src/airflow/providers/ssh/operators/ssh.py::3 providers/ssh/src/airflow/providers/ssh/operators/ssh_remote_job.py::7