Make SQLAlchemy an optional dependency for JDBC and YDB providers (#72373)

Signed-off-by: PoAn Yang <payang@apache.org>

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

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:

providers/jdbc/src/airflow/providers/jdbc/hooks/jdbc.py

@@ -25,11 +25,13 @@ from typing import TYPE_CHECKING, Any, cast
 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 @@ class JdbcHook(DbApiHook):
 
     @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):

providers/jdbc/tests/unit/jdbc/hooks/test_jdbc.py

@@ -31,6 +31,7 @@ import jaydebeapi
 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 @@ class TestJdbcHook:
         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"}),

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
 -----------------------------
 

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:

providers/ydb/src/airflow/providers/ydb/hooks/ydb.py

@@ -20,10 +20,12 @@ from collections.abc import Mapping, Sequence
 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 @@ from airflow.providers.ydb.utils.defaults import CONN_NAME_ATTR, CONN_TYPE, DEFA
 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 @@ class YDBHook(DbApiHook):
 
     @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",

providers/ydb/tests/unit/ydb/hooks/test_ydb.py

@@ -23,6 +23,7 @@ ydb = pytest.importorskip("ydb")
 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

uv.lock

@@ -6466,6 +6466,9 @@ dependencies = [
 openlineage = [
     { name = "apache-airflow-providers-openlineage" },
 ]
+sqlalchemy = [
+    { name = "sqlalchemy" },
+]
 
 [package.dev-dependencies]
 dev = [
@@ -6473,6 +6476,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" },
 ]
@@ -6496,8 +6500,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 = [
@@ -6505,6 +6510,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" },
 ]
@@ -8707,12 +8713,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 = [
@@ -8724,9 +8736,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 = [
@@ -8734,6 +8748,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" }]