This is an automated email from the ASF dual-hosted git repository.
taragolis pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 08f4b923ab Add `sqlalchemy_url` property to `DbApiHook` class (#38871)
08f4b923ab is described below
commit 08f4b923ab6fe63aad72e3a9da9507ed5b9c6932
Author: Kalyan <[email protected]>
AuthorDate: Sun Apr 14 15:32:58 2024 +0530
Add `sqlalchemy_url` property to `DbApiHook` class (#38871)
* Add sqlalchemy_url property to DbApiHook
Signed-off-by: kalyanr <[email protected]>
* 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 <[email protected]>
* Update sql.py
Co-authored-by: Andrey Anshin <[email protected]>
* fix error message
---------
Signed-off-by: kalyanr <[email protected]>
Co-authored-by: Andrey Anshin <[email protected]>
---
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 332bd996fd..070bf0647c 100644
--- a/airflow/providers/common/sql/hooks/sql.py
+++ b/airflow/providers/common/sql/hooks/sql.py
@@ -49,6 +49,7 @@ from airflow.hooks.base import BaseHook
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 @@ class DbApiHook(BaseHook):
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 85edd625f9..16c3d6592a 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
