From 56bb1a7b37af2b8d1a15ad0f78c88f57181f2905 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Tue, 1 Sep 2026 21:38:50 +0900 Subject: [PATCH] Make SQLAlchemy an optional dependency for JDBC and YDB providers Signed-off-by: PoAn Yang --- providers/jdbc/docs/index.rst | 1 + providers/jdbc/pyproject.toml | 4 +++ .../src/airflow/providers/jdbc/hooks/jdbc.py | 12 +++++++- .../jdbc/tests/unit/jdbc/hooks/test_jdbc.py | 12 ++++++++ providers/ydb/docs/index.rst | 17 +++++++++++ providers/ydb/pyproject.toml | 8 +++++ .../src/airflow/providers/ydb/hooks/ydb.py | 15 ++++++++-- .../ydb/tests/unit/ydb/hooks/test_ydb.py | 29 +++++++++++++++++++ uv.lock | 17 ++++++++++- 9 files changed, 111 insertions(+), 4 deletions(-) diff --git a/providers/jdbc/docs/index.rst b/providers/jdbc/docs/index.rst index 657e30e8db249..458739c19b255 100644 --- a/providers/jdbc/docs/index.rst +++ b/providers/jdbc/docs/index.rst @@ -150,6 +150,7 @@ Install them when installing from PyPI. For example: Extra Dependencies =============== ======================================== ``openlineage`` ``apache-airflow-providers-openlineage`` +``sqlalchemy`` ``sqlalchemy>=1.4.54`` =============== ======================================== Downloading official packages diff --git a/providers/jdbc/pyproject.toml b/providers/jdbc/pyproject.toml index a1cdb15ab04b1..1e10b5e24bbfe 100644 --- a/providers/jdbc/pyproject.toml +++ b/providers/jdbc/pyproject.toml @@ -84,6 +84,9 @@ dependencies = [ "openlineage" = [ "apache-airflow-providers-openlineage" ] +"sqlalchemy" = [ + "sqlalchemy>=1.4.54", +] [dependency-groups] dev = [ @@ -94,6 +97,7 @@ dev = [ "apache-airflow-providers-common-sql", "apache-airflow-providers-openlineage", # Additional devel dependencies (do not remove this line and add extra development dependencies) + "apache-airflow-providers-jdbc[sqlalchemy]", ] # To build docs: diff --git a/providers/jdbc/src/airflow/providers/jdbc/hooks/jdbc.py b/providers/jdbc/src/airflow/providers/jdbc/hooks/jdbc.py index c4b6fbf193c60..778ae69a393fa 100644 --- a/providers/jdbc/src/airflow/providers/jdbc/hooks/jdbc.py +++ b/providers/jdbc/src/airflow/providers/jdbc/hooks/jdbc.py @@ -25,11 +25,13 @@ from urllib.parse import quote_plus, urlencode import jaydebeapi -from sqlalchemy.engine import URL +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.common.sql.hooks.sql import DbApiHook if TYPE_CHECKING: + from sqlalchemy.engine import URL + from airflow.providers.common.compat.sdk import Connection @@ -150,6 +152,14 @@ def driver_class(self) -> str | None: @property def sqlalchemy_url(self) -> URL: + 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-jdbc[sqlalchemy]'" + ) from err + conn = self.connection sqlalchemy_query = conn.extra_dejson.get("sqlalchemy_query", {}) if not isinstance(sqlalchemy_query, dict): diff --git a/providers/jdbc/tests/unit/jdbc/hooks/test_jdbc.py b/providers/jdbc/tests/unit/jdbc/hooks/test_jdbc.py index afea430885085..74242d9de5b3d 100644 --- a/providers/jdbc/tests/unit/jdbc/hooks/test_jdbc.py +++ b/providers/jdbc/tests/unit/jdbc/hooks/test_jdbc.py @@ -31,6 +31,7 @@ import pytest from airflow.models import Connection +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.jdbc.hooks.jdbc import JdbcHook, suppress_and_warn jdbc_conn_mock = Mock(name="jdbc_conn") @@ -245,6 +246,17 @@ def test_sqlalchemy_url_with_sqlalchemy_scheme_and_wrong_query_value(self): with pytest.raises(TypeError, match="'sqlalchemy_query' must be of type dict"): hook.sqlalchemy_url + def test_sqlalchemy_url_raises_when_sqlalchemy_is_not_installed(self): + conn_params = dict(extra=json.dumps(dict(sqlalchemy_scheme="mssql"))) + hook = get_hook(conn_params=conn_params) + + with patch.dict("sys.modules", {"sqlalchemy.engine": None}): + with pytest.raises( + AirflowOptionalProviderFeatureException, + match=r"apache-airflow-providers-jdbc\[sqlalchemy\]", + ): + hook.sqlalchemy_url + def test_get_sqlalchemy_engine_verify_creator_is_being_used(self): jdbc_hook = get_hook( conn_params=dict(extra={"sqlalchemy_scheme": "sqlite"}), diff --git a/providers/ydb/docs/index.rst b/providers/ydb/docs/index.rst index 47d548ba7d960..1da6751d85053 100644 --- a/providers/ydb/docs/index.rst +++ b/providers/ydb/docs/index.rst @@ -107,6 +107,23 @@ PIP package Version required ``ydb-dbapi`` ``>=0.1.0`` ========================================== ================== +Optional dependencies +--------------------- + +These extras install optional third-party libraries that enable additional features of the provider. +Install them when installing from PyPI. For example: + +.. code-block:: bash + + pip install apache-airflow-providers-ydb[sqlalchemy] + + +============== ====================== +Extra Dependencies +============== ====================== +``sqlalchemy`` ``sqlalchemy>=1.4.54`` +============== ====================== + Downloading official packages ----------------------------- diff --git a/providers/ydb/pyproject.toml b/providers/ydb/pyproject.toml index d1e46ecac6723..2621e75be22ca 100644 --- a/providers/ydb/pyproject.toml +++ b/providers/ydb/pyproject.toml @@ -66,6 +66,13 @@ dependencies = [ "ydb-dbapi>=0.1.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.54", +] + [dependency-groups] dev = [ "apache-airflow", @@ -74,6 +81,7 @@ dev = [ "apache-airflow-providers-common-compat", "apache-airflow-providers-common-sql", # Additional devel dependencies (do not remove this line and add extra development dependencies) + "apache-airflow-providers-ydb[sqlalchemy]", ] # To build docs: diff --git a/providers/ydb/src/airflow/providers/ydb/hooks/ydb.py b/providers/ydb/src/airflow/providers/ydb/hooks/ydb.py index dd47bcb9f223b..0596b4d01c2f0 100644 --- a/providers/ydb/src/airflow/providers/ydb/hooks/ydb.py +++ b/providers/ydb/src/airflow/providers/ydb/hooks/ydb.py @@ -20,10 +20,12 @@ from typing import TYPE_CHECKING, Any import ydb -from sqlalchemy.engine import URL from ydb_dbapi import Connection as DbApiConnection -from airflow.providers.common.compat.sdk import AirflowException +from airflow.providers.common.compat.sdk import ( + AirflowException, + AirflowOptionalProviderFeatureException, +) from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.providers.ydb.utils.credentials import get_credentials_from_connection from airflow.providers.ydb.utils.defaults import CONN_NAME_ATTR, CONN_TYPE, DEFAULT_CONN_NAME @@ -31,6 +33,7 @@ DEFAULT_YDB_GRPCS_PORT: int = 2135 if TYPE_CHECKING: + from sqlalchemy.engine import URL from ydb_dbapi import Cursor as DbApiCursor try: @@ -231,6 +234,14 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]: @property def sqlalchemy_url(self) -> URL: + 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-ydb[sqlalchemy]'" + ) from err + conn: Connection = self.get_connection(self.get_conn_id()) return URL.create( drivername="ydb", diff --git a/providers/ydb/tests/unit/ydb/hooks/test_ydb.py b/providers/ydb/tests/unit/ydb/hooks/test_ydb.py index 587ec68e6033a..0e29cb740d23a 100644 --- a/providers/ydb/tests/unit/ydb/hooks/test_ydb.py +++ b/providers/ydb/tests/unit/ydb/hooks/test_ydb.py @@ -23,6 +23,7 @@ from unittest.mock import PropertyMock, patch from airflow.models import Connection +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.ydb.hooks.ydb import YDBHook try: @@ -206,3 +207,31 @@ def test_get_df_by_chunks_hook_lineage( assert call_kw["context"] is hook assert call_kw["sql"] == sql assert call_kw["sql_parameters"] == parameters + + +@patch(f"{BASEHOOK_PATCH_PATH}.get_connection") +@patch("ydb.Driver") +@patch("ydb.QuerySessionPool") +def test_sqlalchemy_url_raises_when_sqlalchemy_is_not_installed( + mock_session_pool, mock_driver, mock_get_connection +): + mock_get_connection.return_value = Connection( + conn_type="ydb", + host="grpc://localhost", + port=2135, + login="my_user", + password="my_pwd", + extra={"database": "/my_db1"}, + ) + driver_instance = FakeDriver() + mock_driver.return_value = driver_instance + mock_session_pool.return_value = FakeSessionPool(driver_instance) + + hook = YDBHook() + + with patch.dict("sys.modules", {"sqlalchemy.engine": None}): + with pytest.raises( + AirflowOptionalProviderFeatureException, + match=r"apache-airflow-providers-ydb\[sqlalchemy\]", + ): + hook.sqlalchemy_url diff --git a/uv.lock b/uv.lock index afe82113f142f..f8a57aa5ece51 100644 --- a/uv.lock +++ b/uv.lock @@ -6467,6 +6467,9 @@ dependencies = [ openlineage = [ { name = "apache-airflow-providers-openlineage" }, ] +sqlalchemy = [ + { name = "sqlalchemy" }, +] [package.dev-dependencies] dev = [ @@ -6474,6 +6477,7 @@ dev = [ { name = "apache-airflow-devel-common" }, { name = "apache-airflow-providers-common-compat" }, { name = "apache-airflow-providers-common-sql" }, + { name = "apache-airflow-providers-jdbc", extra = ["sqlalchemy"] }, { name = "apache-airflow-providers-openlineage" }, { name = "apache-airflow-task-sdk" }, ] @@ -6497,8 +6501,9 @@ requires-dist = [ { name = "jpype1", marker = "(python_full_version == '3.11.*' and platform_machine != 'arm64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')", specifier = ">=1.4.1" }, { name = "jpype1", marker = "(python_full_version == '3.12.*' and platform_machine != 'arm64') or (python_full_version == '3.12.*' and sys_platform != 'darwin')", specifier = ">=1.5.0" }, { name = "jpype1", marker = "(python_full_version == '3.13.*' and platform_machine != 'arm64') or (python_full_version == '3.13.*' and sys_platform != 'darwin')", specifier = ">=1.5.1" }, + { name = "sqlalchemy", marker = "extra == 'sqlalchemy'", specifier = ">=1.4.54" }, ] -provides-extras = ["openlineage"] +provides-extras = ["openlineage", "sqlalchemy"] [package.metadata.requires-dev] dev = [ @@ -6506,6 +6511,7 @@ dev = [ { name = "apache-airflow-devel-common", editable = "devel-common" }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, { name = "apache-airflow-providers-common-sql", editable = "providers/common/sql" }, + { name = "apache-airflow-providers-jdbc", extras = ["sqlalchemy"], editable = "providers/jdbc" }, { name = "apache-airflow-providers-openlineage", editable = "providers/openlineage" }, { name = "apache-airflow-task-sdk", editable = "task-sdk" }, ] @@ -8708,12 +8714,18 @@ dependencies = [ { name = "ydb-dbapi" }, ] +[package.optional-dependencies] +sqlalchemy = [ + { name = "sqlalchemy" }, +] + [package.dev-dependencies] dev = [ { name = "apache-airflow" }, { name = "apache-airflow-devel-common" }, { name = "apache-airflow-providers-common-compat" }, { name = "apache-airflow-providers-common-sql" }, + { name = "apache-airflow-providers-ydb", extra = ["sqlalchemy"] }, { name = "apache-airflow-task-sdk" }, ] docs = [ @@ -8725,9 +8737,11 @@ requires-dist = [ { name = "apache-airflow", editable = "." }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, { name = "apache-airflow-providers-common-sql", editable = "providers/common/sql" }, + { name = "sqlalchemy", marker = "extra == 'sqlalchemy'", specifier = ">=1.4.54" }, { name = "ydb", specifier = ">=3.18.8" }, { name = "ydb-dbapi", specifier = ">=0.1.0" }, ] +provides-extras = ["sqlalchemy"] [package.metadata.requires-dev] dev = [ @@ -8735,6 +8749,7 @@ dev = [ { name = "apache-airflow-devel-common", editable = "devel-common" }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, { name = "apache-airflow-providers-common-sql", editable = "providers/common/sql" }, + { name = "apache-airflow-providers-ydb", extras = ["sqlalchemy"], editable = "providers/ydb" }, { name = "apache-airflow-task-sdk", editable = "task-sdk" }, ] docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "devel-common" }]