diff --git a/providers/postgres/pyproject.toml b/providers/postgres/pyproject.toml index d0febbd0ba195..6929ce837f5b0 100644 --- a/providers/postgres/pyproject.toml +++ b/providers/postgres/pyproject.toml @@ -69,6 +69,9 @@ dependencies = [ # The optional dependencies should be modified in place in the generated file # Any change in the dependencies is preserved when the file is regenerated [project.optional-dependencies] +sqlalchemy = [ + "sqlalchemy>=1.4.49" +] "amazon" = [ "apache-airflow-providers-amazon>=2.6.0", ] diff --git a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py index 1bc554cc9985d..a87f2d03b7061 100644 --- a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py +++ b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py @@ -27,7 +27,6 @@ import psycopg2.extras from more_itertools import chunked from psycopg2.extras import DictCursor, NamedTupleCursor, RealDictCursor, execute_batch -from sqlalchemy.engine import URL from airflow.exceptions import AirflowOptionalProviderFeatureException from airflow.providers.common.compat.sdk import AirflowException, Connection, conf @@ -57,6 +56,8 @@ from airflow.providers.common.sql.dialects.dialect import Dialect from airflow.providers.openlineage.sqlparser import DatabaseInfo + from sqlalchemy.engine import URL + if USE_PSYCOPG3: from psycopg.errors import Diagnostic @@ -167,6 +168,13 @@ def __cast_nullable(value, dst_type: type) -> Any: @property def sqlalchemy_url(self) -> URL: + try: + from sqlalchemy.engine import URL + except (ImportError, ModuleNotFoundError) as err: + raise AirflowOptionalProviderFeatureException( + "SQLAlchemy is not installed. Please install it with " + "`pip install apache-airflow-providers-postgres[sqlalchemy]`." + ) from err conn = self.connection query = conn.extra_dejson.get("sqlalchemy_query", {}) if not isinstance(query, dict):