Make SQLAlchemy an optional dependency for JDBC and YDB providers

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

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"}),

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,17 +20,20 @@
 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
 
 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",

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