From c80590e52a6f5356585ac1285ed6e207b6ba8ae8 Mon Sep 17 00:00:00 2001 From: Dr Alex Mitre <30060514+mitre88@users.noreply.github.com> Date: Mon, 27 Jul 2026 09:43:16 -0600 Subject: [PATCH] Use built-in exceptions for Postgres hook input validation A wrong type for sqlalchemy_query is a TypeError, and requesting the psycopg3-unsupported realdictcursor is a ValueError like the invalid cursor branch right next to it; both were raised as the broad AirflowException the community is actively reducing. Follows the clean-up pattern of #66279. --- generated/known_airflow_exceptions.txt | 1 - .../src/airflow/providers/postgres/hooks/postgres.py | 7 ++----- .../postgres/tests/unit/postgres/hooks/test_postgres.py | 4 ++-- 3 files changed, 4 insertions(+), 8 deletions(-) diff --git a/generated/known_airflow_exceptions.txt b/generated/known_airflow_exceptions.txt index 1fb461854905d..6c0067c564fcf 100644 --- a/generated/known_airflow_exceptions.txt +++ b/generated/known_airflow_exceptions.txt @@ -376,7 +376,6 @@ providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py::1 providers/opensearch/src/airflow/providers/opensearch/operators/opensearch.py::9 providers/pagerduty/src/airflow/providers/pagerduty/hooks/pagerduty.py::1 providers/pagerduty/src/airflow/providers/pagerduty/hooks/pagerduty_events.py::2 -providers/postgres/src/airflow/providers/postgres/hooks/postgres.py::2 providers/presto/src/airflow/providers/presto/hooks/presto.py::1 providers/samba/src/airflow/providers/samba/transfers/gcs_to_samba.py::1 providers/segment/src/airflow/providers/segment/hooks/segment.py::2 diff --git a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py index 3821779e261fd..afbb2094639da 100644 --- a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py +++ b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py @@ -26,7 +26,6 @@ from more_itertools import chunked from airflow.providers.common.compat.sdk import ( - AirflowException, AirflowOptionalProviderFeatureException, Connection, conf, @@ -189,7 +188,7 @@ def sqlalchemy_url(self) -> URL: conn = self.connection query = conn.extra_dejson.get("sqlalchemy_query", {}) if not isinstance(query, dict): - raise AirflowException("The parameter 'sqlalchemy_query' must be of type dict!") + raise TypeError("The parameter 'sqlalchemy_query' must be of type dict!") if conn.extra_dejson.get("iam", False): conn.login, conn.password, conn.port = self.get_iam_token(conn) return URL.create( @@ -222,9 +221,7 @@ def _get_cursor(self, raw_cursor: str) -> CursorType: if _cursor == "namedtuplecursor": return namedtuple_row if _cursor == "realdictcursor": - raise AirflowException( - "realdictcursor is not supported with psycopg3. Use dictcursor instead." - ) + raise ValueError("realdictcursor is not supported with psycopg3. Use dictcursor instead.") valid_cursors = "dictcursor, namedtuplecursor" raise ValueError(f"Invalid cursor passed {_cursor}. Valid options are: {valid_cursors}") diff --git a/providers/postgres/tests/unit/postgres/hooks/test_postgres.py b/providers/postgres/tests/unit/postgres/hooks/test_postgres.py index a9eff1bcb6f6a..12fe869a212d9 100644 --- a/providers/postgres/tests/unit/postgres/hooks/test_postgres.py +++ b/providers/postgres/tests/unit/postgres/hooks/test_postgres.py @@ -27,7 +27,7 @@ import sqlalchemy from airflow.models import Connection -from airflow.providers.common.compat.sdk import AirflowException, AirflowOptionalProviderFeatureException +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.postgres.dialects.postgres import PostgresDialect from airflow.providers.postgres.hooks.postgres import PostgresHook @@ -110,7 +110,7 @@ def test_sqlalchemy_url_with_wrong_sqlalchemy_query_value(self): ) hook = PostgresHook(connection=conn) - with pytest.raises(AirflowException): + with pytest.raises(TypeError, match="'sqlalchemy_query' must be of type dict"): hook.sqlalchemy_url @pytest.mark.parametrize("aws_conn_id", [NOTSET, None, "mock_aws_conn"])