From 08f4b923ab6fe63aad72e3a9da9507ed5b9c6932 Mon Sep 17 00:00:00 2001 From: Kalyan Date: Sun, 14 Apr 2024 15:32:58 +0530 Subject: [PATCH] Add `sqlalchemy_url` property to `DbApiHook` class (#38871) * Add sqlalchemy_url property to DbApiHook Signed-off-by: kalyanr * update get_sqlalchemy_url docstring * update drivername property docstring * fix spelling * add sqlalchemy_url property to DbApiHook * fix typo * precommit failure fix Signed-off-by: kalyanr * Update sql.py Co-authored-by: Andrey Anshin * fix error message --------- Signed-off-by: kalyanr Co-authored-by: Andrey Anshin --- airflow/providers/common/sql/hooks/sql.py | 17 +++++++++++++++++ airflow/providers/common/sql/hooks/sql.pyi | 3 +++ 2 files changed, 20 insertions(+) diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py index 332bd996fd7fe..070bf0647cfc6 100644 --- a/airflow/providers/common/sql/hooks/sql.py +++ b/airflow/providers/common/sql/hooks/sql.py @@ -49,6 +49,7 @@ if TYPE_CHECKING: from pandas import DataFrame + from sqlalchemy.engine import URL from airflow.providers.openlineage.extractors import OperatorLineage from airflow.providers.openlineage.sqlparser import DatabaseInfo @@ -209,6 +210,22 @@ def get_uri(self) -> str: conn.schema = self.__schema or conn.schema return conn.get_uri() + @property + def sqlalchemy_url(self) -> URL: + """ + Return a Sqlalchemy.engine.URL object from the connection. + + Needs to be implemented in the provider subclass to return the sqlalchemy.engine.URL object. + + :return: the extracted sqlalchemy.engine.URL object. + """ + qualname = f"{self.__class__.__module__}.{self.__class__.__qualname__}" + if qualname != "airflow.providers.common.sql.hooks.sql.DbApiHook": + msg = f"{qualname!r} does not implement/support built SQLAlchemy URL." + else: + msg = "`sqlalchemy_url` property should be implemented in the provider subclass." + raise NotImplementedError(msg) + def get_sqlalchemy_engine(self, engine_kwargs=None): """ Get an sqlalchemy_engine object. diff --git a/airflow/providers/common/sql/hooks/sql.pyi b/airflow/providers/common/sql/hooks/sql.pyi index 85edd625f99e0..16c3d6592a341 100644 --- a/airflow/providers/common/sql/hooks/sql.pyi +++ b/airflow/providers/common/sql/hooks/sql.pyi @@ -41,6 +41,7 @@ from airflow.hooks.base import BaseHook as BaseHook from airflow.providers.openlineage.extractors import OperatorLineage as OperatorLineage from airflow.providers.openlineage.sqlparser import DatabaseInfo as DatabaseInfo from pandas import DataFrame as DataFrame +from sqlalchemy.engine import URL as URL from typing import Any, Callable, Generator, Iterable, Mapping, Protocol, Sequence, TypeVar, overload T = TypeVar("T") @@ -66,6 +67,8 @@ class DbApiHook(BaseHook): def placeholder(self): ... def get_conn(self): ... def get_uri(self) -> str: ... + @property + def sqlalchemy_url(self) -> URL: ... def get_sqlalchemy_engine(self, engine_kwargs: Incomplete | None = None): ... def get_pandas_df( self, sql, parameters: list | tuple | Mapping[str, Any] | None = None, **kwargs