diff --git a/airflow/providers/microsoft/mssql/hooks/mssql.py b/airflow/providers/microsoft/mssql/hooks/mssql.py index 91952f5b37e40..d4deb609486d4 100644 --- a/airflow/providers/microsoft/mssql/hooks/mssql.py +++ b/airflow/providers/microsoft/mssql/hooks/mssql.py @@ -19,12 +19,16 @@ from __future__ import annotations -from typing import Any +from functools import cached_property +from typing import TYPE_CHECKING, Any import pymssql from airflow.providers.common.sql.hooks.sql import DbApiHook +if TYPE_CHECKING: + from airflow.models import Connection + class MsSqlHook(DbApiHook): """ @@ -53,6 +57,15 @@ def __init__( self.schema = kwargs.pop("schema", None) self._sqlalchemy_scheme = sqlalchemy_scheme + @cached_property + def connection(self) -> Connection: + """ + Get the airflow connection object. + + :return: The connection object. + """ + return self.get_connection(getattr(self, self.conn_name_attr)) + @property def connection_extra_lower(self) -> dict: """ @@ -60,8 +73,7 @@ def connection_extra_lower(self) -> dict: This is used internally for case-insensitive access of mssql params. """ - conn = self.get_connection(self.mssql_conn_id) # type: ignore[attr-defined] - return {k.lower(): v for k, v in conn.extra_dejson.items()} + return {k.lower(): v for k, v in self.connection.extra_dejson.items()} @property def sqlalchemy_scheme(self) -> str: @@ -94,9 +106,8 @@ def get_sqlalchemy_connection( def get_conn(self) -> pymssql.connect: """Return ``pymssql`` connection object.""" - conn = self.get_connection(self.mssql_conn_id) # type: ignore[attr-defined] - - conn = pymssql.connect( + conn = self.connection + return pymssql.connect( server=conn.host, user=conn.login, password=conn.password,