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..518eabeb5a60a 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 @@ -54,6 +53,7 @@ if TYPE_CHECKING: from pandas import DataFrame as PandasDataFrame from polars import DataFrame as PolarsDataFrame + from sqlalchemy.engine import URL from airflow.providers.common.sql.dialects.dialect import Dialect from airflow.providers.openlineage.sqlparser import DatabaseInfo @@ -167,6 +167,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):