diff --git a/providers/vertica/pyproject.toml b/providers/vertica/pyproject.toml index 79ede554c5912..57cb90be53bba 100644 --- a/providers/vertica/pyproject.toml +++ b/providers/vertica/pyproject.toml @@ -63,6 +63,13 @@ dependencies = [ "vertica-python>=1.3.0", ] +# 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", +] + [dependency-groups] dev = [ "apache-airflow", @@ -71,6 +78,7 @@ dev = [ "apache-airflow-providers-common-sql", # Additional devel dependencies (do not remove this line and add extra development dependencies) "apache-airflow-providers-common-sql[pandas,polars]", + "apache-airflow-providers-vertica[sqlalchemy]" ] # To build docs: diff --git a/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py b/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py index 6e1ea11c67c70..b0e65bfa26319 100644 --- a/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py +++ b/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py @@ -17,14 +17,17 @@ from __future__ import annotations from collections.abc import Callable, Iterable, Mapping -from typing import Any, overload +from typing import TYPE_CHECKING, Any, overload -from sqlalchemy.engine import URL from vertica_python import connect +from airflow.exceptions import AirflowOptionalProviderFeatureException from airflow.providers.common.sql.hooks.handlers import fetch_all_handler from airflow.providers.common.sql.hooks.sql import DbApiHook +if TYPE_CHECKING: + from sqlalchemy.engine import URL + def vertica_fetch_all_handler(cursor) -> list[tuple] | None: """ @@ -137,6 +140,13 @@ def get_conn(self) -> connect: @property def sqlalchemy_url(self) -> URL: """Return a SQLAlchemy URL object with properly formatted query parameters.""" + try: + from sqlalchemy.engine import URL + except (ImportError, ModuleNotFoundError) as err: + raise AirflowOptionalProviderFeatureException( + "The 'sqlalchemy' library is required to use this feature. " + "Please install it with: pip install 'apache-airflow-providers-vertica[sqlalchemy]'" + ) from err conn = self.get_connection(self.get_conn_id()) extra = conn.extra_dejson or {}