Generalize caching of connection in DbApiHook to improve performance (#…

…40751)

airflow/providers/common/sql/CHANGELOG.rst

@@ -25,6 +25,14 @@
 Changelog
 ---------
 
+1.17.0
+......
+
+Features
+~~~~~~~~
+
+* ``Connection in DB Hook is now cached to avoid multiple lookups when properties from extras have to be resolved``
+
 1.16.0
 ......
 

airflow/providers/common/sql/hooks/sql.py

@@ -53,6 +53,7 @@
     from pandas import DataFrame
     from sqlalchemy.engine import URL
 
+    from airflow.models import Connection
     from airflow.providers.openlineage.extractors import OperatorLineage
     from airflow.providers.openlineage.sqlparser import DatabaseInfo
 
@@ -183,14 +184,14 @@ def __init__(self, *args, schema: str | None = None, log_sql: bool = True, **kwa
         self._replace_statement_format: str = kwargs.get(
             "replace_statement_format", "REPLACE INTO {} {} VALUES ({})"
         )
+        self._connection: Connection | None = kwargs.pop("connection", None)
 
     def get_conn_id(self) -> str:
         return getattr(self, self.conn_name_attr)
 
     @cached_property
     def placeholder(self):
-        conn = self.get_connection(self.get_conn_id())
-        placeholder = conn.extra_dejson.get("placeholder")
+        placeholder = self.connection_extra.get("placeholder")
         if placeholder:
             if placeholder in SQL_PLACEHOLDERS:
                 return placeholder
@@ -203,9 +204,28 @@ def placeholder(self):
             )
         return self._placeholder
 
+    @property
+    def connection(self) -> Connection:
+        if self._connection is None:
+            self._connection = self.get_connection(self.get_conn_id())
+        return self._connection
+
+    @property
+    def connection_extra(self) -> dict:
+        return self.connection.extra_dejson
+
+    @cached_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 extra params.
+        """
+        return {k.lower(): v for k, v in self.connection_extra.items()}
+
     def get_conn(self):
         """Return a connection object."""
-        db = self.get_connection(self.get_conn_id())
+        db = self.connection
         return self.connector.connect(host=db.host, port=db.port, username=db.login, schema=db.schema)
 
     def get_uri(self) -> str:

airflow/providers/common/sql/hooks/sql.pyi

@@ -38,6 +38,7 @@ from airflow.exceptions import (
     AirflowProviderDeprecationWarning as AirflowProviderDeprecationWarning,
 )
 from airflow.hooks.base import BaseHook as BaseHook
+from airflow.models import Connection as Connection
 from airflow.providers.openlineage.extractors import OperatorLineage as OperatorLineage
 from airflow.providers.openlineage.sqlparser import DatabaseInfo as DatabaseInfo
 from functools import cached_property as cached_property
@@ -67,6 +68,12 @@ class DbApiHook(BaseHook):
     def get_conn_id(self) -> str: ...
     @cached_property
     def placeholder(self): ...
+    @property
+    def connection(self) -> Connection: ...
+    @property
+    def connection_extra(self) -> dict: ...
+    @cached_property
+    def connection_extra_lower(self) -> dict: ...
     def get_conn(self): ...
     def get_uri(self) -> str: ...
     @property

airflow/providers/common/sql/provider.yaml

@@ -25,6 +25,7 @@ state: ready
 source-date-epoch: 1723970051
 # note that those versions are maintained by release manager - do not update them manually
 versions:
+  - 1.17.0
   - 1.16.0
   - 1.15.0
   - 1.14.2

airflow/providers/elasticsearch/hooks/elasticsearch.py

@@ -93,12 +93,11 @@ class ElasticsearchSQLHook(DbApiHook):
     def __init__(self, schema: str = "http", connection: AirflowConnection | None = None, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.schema = schema
-        self.connection = connection
+        self._connection = connection
 
     def get_conn(self) -> ESConnection:
         """Return an elasticsearch connection object."""
-        conn_id = self.get_conn_id()
-        conn = self.connection or self.get_connection(conn_id)
+        conn = self.connection
 
         conn_args = {
             "host": conn.host,
@@ -117,8 +116,7 @@ def get_conn(self) -> ESConnection:
         return connect(**conn_args)
 
     def get_uri(self) -> str:
-        conn_id = self.get_conn_id()
-        conn = self.connection or self.get_connection(conn_id)
+        conn = self.connection
 
         login = ""
         if conn.login:

airflow/providers/elasticsearch/provider.yaml

@@ -68,7 +68,7 @@ versions:
 
 dependencies:
   - apache-airflow>=2.8.0
-  - apache-airflow-providers-common-sql>=1.14.1
+  - apache-airflow-providers-common-sql>=1.17.0
   - elasticsearch>=8.10,<9
 
 integrations:

airflow/providers/jdbc/hooks/jdbc.py

@@ -106,16 +106,6 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]:
             "relabeling": {"host": "Connection URL"},
         }
 
-    @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 jdbc params.
-        """
-        conn = self.get_connection(self.get_conn_id())
-        return {k.lower(): v for k, v in conn.extra_dejson.items()}
-
     @property
     def driver_path(self) -> str | None:
         from airflow.configuration import conf

airflow/providers/jdbc/provider.yaml

@@ -53,7 +53,7 @@ versions:
 
 dependencies:
   - apache-airflow>=2.8.0
-  - apache-airflow-providers-common-sql>=1.14.1
+  - apache-airflow-providers-common-sql>=1.17.0
   - jaydebeapi>=1.1.1
 
 integrations:

airflow/providers/microsoft/mssql/hooks/mssql.py

@@ -19,18 +19,14 @@
 
 from __future__ import annotations
 
-from functools import cached_property
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 import pymssql
 from methodtools import lru_cache
 from pymssql import Connection as PymssqlConnection
 
 from airflow.providers.common.sql.hooks.sql import DbApiHook, fetch_all_handler
 
-if TYPE_CHECKING:
-    from airflow.models import Connection
-
 
 class MsSqlHook(DbApiHook):
     """
@@ -59,24 +55,6 @@ 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(self.get_conn_id())
-
-    @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.
-        """
-        return {k.lower(): v for k, v in self.connection.extra_dejson.items()}
-
     @property
     def sqlalchemy_scheme(self) -> str:
         """Sqlalchemy scheme either from constructor, connection extras or default."""

airflow/providers/microsoft/mssql/provider.yaml

@@ -55,7 +55,7 @@ versions:
 
 dependencies:
   - apache-airflow>=2.8.0
-  - apache-airflow-providers-common-sql>=1.14.1
+  - apache-airflow-providers-common-sql>=1.17.0
   - pymssql>=2.3.0
   # The methodtools dependency can be removed with min airflow version >=2.9.1
   # as it was added in https://github.com/apache/airflow/pull/37757

airflow/providers/mysql/hooks/mysql.py

@@ -72,7 +72,6 @@ class MySqlHook(DbApiHook):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.schema = kwargs.pop("schema", None)
-        self.connection = kwargs.pop("connection", None)
         self.local_infile = kwargs.pop("local_infile", False)
         self.init_command = kwargs.pop("init_command", None)
 

airflow/providers/mysql/provider.yaml

@@ -66,7 +66,7 @@ versions:
 
 dependencies:
   - apache-airflow>=2.8.0
-  - apache-airflow-providers-common-sql>=1.14.1
+  - apache-airflow-providers-common-sql>=1.17.0
   - mysqlclient>=1.4.0
   - mysql-connector-python>=8.0.29
 

airflow/providers/odbc/hooks/odbc.py

@@ -79,13 +79,6 @@ def __init__(
         self._connection = None
         self._connect_kwargs = connect_kwargs
 
-    @property
-    def connection(self):
-        """The Connection object with ID ``odbc_conn_id``."""
-        if not self._connection:
-            self._connection = self.get_connection(self.get_conn_id())
-        return self._connection
-
     @property
     def database(self) -> str | None:
         """Database provided in init if exists; otherwise, ``schema`` from ``Connection`` object."""
@@ -99,15 +92,6 @@ def sqlalchemy_scheme(self) -> str:
             raise RuntimeError("sqlalchemy_scheme in connection extra should not contain : or / characters")
         return self._sqlalchemy_scheme or extra_scheme or self.DEFAULT_SQLALCHEMY_SCHEME
 
-    @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 odbc params.
-        """
-        return {k.lower(): v for k, v in self.connection.extra_dejson.items()}
-
     @property
     def driver(self) -> str | None:
         """Driver from init param if given; else try to find one in connection extra."""
@@ -166,9 +150,7 @@ def odbc_connection_string(self):
                 conn_str += f"PORT={self.connection.port};"
 
             extra_exclude = {"driver", "dsn", "connect_kwargs", "sqlalchemy_scheme", "placeholder"}
-            extra_params = {
-                k: v for k, v in self.connection.extra_dejson.items() if k.lower() not in extra_exclude
-            }
+            extra_params = {k: v for k, v in self.connection_extra.items() if k.lower() not in extra_exclude}
             for k, v in extra_params.items():
                 conn_str += f"{k}={v};"
 

airflow/providers/odbc/provider.yaml

@@ -54,7 +54,7 @@ versions:
 
 dependencies:
   - apache-airflow>=2.8.0
-  - apache-airflow-providers-common-sql>=1.14.1
+  - apache-airflow-providers-common-sql>=1.17.0
   - pyodbc>=5.0.0
 
 integrations:

airflow/providers/postgres/hooks/postgres.py

@@ -88,7 +88,6 @@ def __init__(self, *args, options: str | None = None, **kwargs) -> None:
             )
             kwargs["database"] = kwargs["schema"]
         super().__init__(*args, **kwargs)
-        self.connection: Connection | None = kwargs.pop("connection", None)
         self.conn: connection = None
         self.database: str | None = kwargs.pop("database", None)
         self.options = options
@@ -142,8 +141,7 @@ def _get_cursor(self, raw_cursor: str) -> CursorType:
 
     def get_conn(self) -> connection:
         """Establish a connection to a postgres database."""
-        conn_id = self.get_conn_id()
-        conn = deepcopy(self.connection or self.get_connection(conn_id))
+        conn = deepcopy(self.connection)
 
         # check for authentication via AWS IAM
         if conn.extra_dejson.get("iam", False):

airflow/providers/postgres/provider.yaml

@@ -65,7 +65,7 @@ versions:
 
 dependencies:
   - apache-airflow>=2.8.0
-  - apache-airflow-providers-common-sql>=1.14.1
+  - apache-airflow-providers-common-sql>=1.17.0
   - psycopg2-binary>=2.9.4
 
 additional-extras:

dev/breeze/tests/test_packages.py

@@ -204,7 +204,7 @@ def test_get_documentation_package_path():
             "postgres",
             "beta0",
             """
-    "apache-airflow-providers-common-sql>=1.14.1b0",
+    "apache-airflow-providers-common-sql>=1.17.0b0",
     "apache-airflow>=2.8.0b0",
     "psycopg2-binary>=2.9.4",
     """,
@@ -214,7 +214,7 @@ def test_get_documentation_package_path():
             "postgres",
             "",
             """
-    "apache-airflow-providers-common-sql>=1.14.1",
+    "apache-airflow-providers-common-sql>=1.17.0",
     "apache-airflow>=2.8.0",
     "psycopg2-binary>=2.9.4",
     """,

generated/provider_dependencies.json

@@ -518,7 +518,7 @@
   },
   "elasticsearch": {
     "deps": [
-      "apache-airflow-providers-common-sql>=1.14.1",
+      "apache-airflow-providers-common-sql>=1.17.0",
       "apache-airflow>=2.8.0",
       "elasticsearch>=8.10,<9"
     ],
@@ -756,7 +756,7 @@
   },
   "jdbc": {
     "deps": [
-      "apache-airflow-providers-common-sql>=1.14.1",
+      "apache-airflow-providers-common-sql>=1.17.0",
       "apache-airflow>=2.8.0",
       "jaydebeapi>=1.1.1"
     ],
@@ -820,7 +820,7 @@
   },
   "microsoft.mssql": {
     "deps": [
-      "apache-airflow-providers-common-sql>=1.14.1",
+      "apache-airflow-providers-common-sql>=1.17.0",
       "apache-airflow>=2.8.0",
       "methodtools>=0.4.7",
       "pymssql>=2.3.0"
@@ -871,7 +871,7 @@
   },
   "mysql": {
     "deps": [
-      "apache-airflow-providers-common-sql>=1.14.1",
+      "apache-airflow-providers-common-sql>=1.17.0",
       "apache-airflow>=2.8.0",
       "mysql-connector-python>=8.0.29",
       "mysqlclient>=1.4.0"
@@ -902,7 +902,7 @@
   },
   "odbc": {
     "deps": [
-      "apache-airflow-providers-common-sql>=1.14.1",
+      "apache-airflow-providers-common-sql>=1.17.0",
       "apache-airflow>=2.8.0",
       "pyodbc>=5.0.0"
     ],
@@ -1047,7 +1047,7 @@
   },
   "postgres": {
     "deps": [
-      "apache-airflow-providers-common-sql>=1.14.1",
+      "apache-airflow-providers-common-sql>=1.17.0",
       "apache-airflow>=2.8.0",
       "psycopg2-binary>=2.9.4"
     ],