Skip to content

Commit

Permalink
use cached_property to reuse get_connections value in mssql pro…
Browse files Browse the repository at this point in the history
…vider (#39575)

* use self.conn instead of calling get_connections

* rename property
  • Loading branch information
rawwar authored May 12, 2024
1 parent 97e867f commit 9e61daf
Showing 1 changed file with 17 additions and 6 deletions.
23 changes: 17 additions & 6 deletions airflow/providers/microsoft/mssql/hooks/mssql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down Expand Up @@ -53,15 +57,23 @@ 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:
"""
``connection.extra_dejson`` but where keys are converted to lower case.
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:
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 9e61daf

Please sign in to comment.