apache-airflow-providers-common-sql 1.30.1__tar.gz → 1.30.4__tar.gz

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/NOTICE

@@ -1,5 +1,5 @@
 Apache Airflow
-Copyright 2016-2025 The Apache Software Foundation
+Copyright 2016-2026 The Apache Software Foundation
 
 This product includes software developed at
 The Apache Software Foundation (http://www.apache.org/).

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-common-sql
-Version: 1.30.1
+Version: 1.30.4
 Summary: Provider package apache-airflow-providers-common-sql for Apache Airflow
 Keywords: airflow-provider,common.sql,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -23,7 +23,7 @@ Classifier: Topic :: System :: Monitoring
 License-File: LICENSE
 License-File: NOTICE
 Requires-Dist: apache-airflow>=2.11.0
-Requires-Dist: apache-airflow-providers-common-compat>=1.10.1
+Requires-Dist: apache-airflow-providers-common-compat>=1.12.0
 Requires-Dist: sqlparse>=0.5.1
 Requires-Dist: more-itertools>=9.0.0
 Requires-Dist: methodtools>=0.4.7
@@ -31,9 +31,10 @@ Requires-Dist: apache-airflow-providers-openlineage ; extra == "openlineage"
 Requires-Dist: pandas[sql-other]>=2.1.2 ; extra == "pandas" and ( python_version <"3.13")
 Requires-Dist: pandas>=2.2.3 ; extra == "pandas" and ( python_version >="3.13")
 Requires-Dist: polars>=1.26.0 ; extra == "polars"
+Requires-Dist: sqlalchemy>=1.4.49 ; extra == "sqlalchemy"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.1/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.1
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.4/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.4
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
@@ -41,6 +42,7 @@ Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 Provides-Extra: openlineage
 Provides-Extra: pandas
 Provides-Extra: polars
+Provides-Extra: sqlalchemy
 
 
 .. Licensed to the Apache Software Foundation (ASF) under one
@@ -67,7 +69,7 @@ Provides-Extra: polars
 
 Package ``apache-airflow-providers-common-sql``
 
-Release: ``1.30.1``
+Release: ``1.30.4``
 
 
 `Common SQL Provider <https://en.wikipedia.org/wiki/SQL>`__
@@ -80,7 +82,7 @@ This is a provider package for ``common.sql`` provider. All classes for this pro
 are in ``airflow.providers.common.sql`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.4/>`_.
 
 Installation
 ------------
@@ -98,7 +100,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  ==================
 ``apache-airflow``                          ``>=2.11.0``
-``apache-airflow-providers-common-compat``  ``>=1.10.1``
+``apache-airflow-providers-common-compat``  ``>=1.12.0``
 ``sqlparse``                                ``>=0.5.1``
 ``more-itertools``                          ``>=9.0.0``
 ``methodtools``                             ``>=0.4.7``
@@ -133,8 +135,9 @@ Extra            Dependencies
 ``pandas``       ``pandas[sql-other]>=2.1.2; python_version <"3.13"``, ``pandas>=2.2.3; python_version >="3.13"``
 ``openlineage``  ``apache-airflow-providers-openlineage``
 ``polars``       ``polars>=1.26.0``
+``sqlalchemy``   ``sqlalchemy>=1.4.49``
 ===============  ================================================================================================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.4/changelog.html>`_.
 

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/README.rst

@@ -23,7 +23,7 @@
 
 Package ``apache-airflow-providers-common-sql``
 
-Release: ``1.30.1``
+Release: ``1.30.4``
 
 
 `Common SQL Provider <https://en.wikipedia.org/wiki/SQL>`__
@@ -36,7 +36,7 @@ This is a provider package for ``common.sql`` provider. All classes for this pro
 are in ``airflow.providers.common.sql`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.4/>`_.
 
 Installation
 ------------
@@ -54,7 +54,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  ==================
 ``apache-airflow``                          ``>=2.11.0``
-``apache-airflow-providers-common-compat``  ``>=1.10.1``
+``apache-airflow-providers-common-compat``  ``>=1.12.0``
 ``sqlparse``                                ``>=0.5.1``
 ``more-itertools``                          ``>=9.0.0``
 ``methodtools``                             ``>=0.4.7``
@@ -89,7 +89,8 @@ Extra            Dependencies
 ``pandas``       ``pandas[sql-other]>=2.1.2; python_version <"3.13"``, ``pandas>=2.2.3; python_version >="3.13"``
 ``openlineage``  ``apache-airflow-providers-openlineage``
 ``polars``       ``polars>=1.26.0``
+``sqlalchemy``   ``sqlalchemy>=1.4.49``
 ===============  ================================================================================================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.4/changelog.html>`_.

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/docs/changelog.rst

@@ -25,6 +25,50 @@
 Changelog
 ---------
 
+1.30.4
+......
+
+Misc
+~~~~
+
+* ``Add SQLAlchemy to the dev‑dependency group in providers (#60472)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+
+1.30.3
+......
+
+Misc
+~~~~
+
+* ``Consume ''AirflowOptionalProviderFeatureException'' from compat sdk in providers (#60335)``
+* ``New year means updated Copyright notices (#60344)``
+* ``Assorted SQLA2/Mypy fixes (#60094)``
+* ``Add proper dependencies for sqlalchemy - common.sql (#59945)``
+* ``Migrate common provider to use airflow.sdk.configuration.conf (#59989)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+
+1.30.2
+......
+
+Misc
+~~~~
+
+* ``Remove top-level SDK reference in Core (#59817)``
+* ``Extract shared "module_loading" distribution (#59139)``
+
+Doc-only
+~~~~~~~~
+
+* ``Add missing fields to SQLValueCheckOperator docstring (#59409)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+   * ``TaskInstance unused method cleanup (#59835)``
+
 1.30.1
 ......
 

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/docs/index.rst

@@ -79,7 +79,7 @@ apache-airflow-providers-common-sql package
 `Common SQL Provider <https://en.wikipedia.org/wiki/SQL>`__
 
 
-Release: 1.30.1
+Release: 1.30.4
 
 Provider package
 ----------------
@@ -103,7 +103,7 @@ The minimum Apache Airflow version supported by this provider distribution is ``
 PIP package                                 Version required
 ==========================================  ==================
 ``apache-airflow``                          ``>=2.11.0``
-``apache-airflow-providers-common-compat``  ``>=1.10.1``
+``apache-airflow-providers-common-compat``  ``>=1.12.0``
 ``sqlparse``                                ``>=0.5.1``
 ``more-itertools``                          ``>=9.0.0``
 ``methodtools``                             ``>=0.4.7``
@@ -135,5 +135,5 @@ Downloading official packages
 You can download officially released packages and verify their checksums and signatures from the
 `Official Apache Download site <https://downloads.apache.org/airflow/providers/>`_
 
-* `The apache-airflow-providers-common-sql 1.30.1 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.1.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.1.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.1.tar.gz.sha512>`__)
-* `The apache-airflow-providers-common-sql 1.30.1 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.1-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.1-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.1-py3-none-any.whl.sha512>`__)
+* `The apache-airflow-providers-common-sql 1.30.4 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.4.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.4.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.4.tar.gz.sha512>`__)
+* `The apache-airflow-providers-common-sql 1.30.4 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.4-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.4-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-1.30.4-py3-none-any.whl.sha512>`__)

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/docs/operators.rst

@@ -148,6 +148,29 @@ The below example demonstrates how to instantiate the SQLTableCheckOperator task
     :end-before: [END howto_operator_sql_table_check]
 
 
+.. _howto/operator:SQLValueCheckOperator:
+
+Check value against expected
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use the :class:`~airflow.providers.common.sql.operators.sql.SQLValueCheckOperator` to compare a SQL query result
+against an expected value, with some optionally specified tolerance for numeric results.
+The parameters for this operator are:
+
+- ``sql`` - the sql query to be executed, as a templated string.
+- ``pass_value`` - the expected value to compare the query result against.
+- ``tolerance`` (optional) - numerical tolerance for comparisons involving numeric values.
+- ``conn_id`` (optional) - the connection ID used to connect to the database.
+- ``database`` (optional) - name of the database which overwrites the name defined in the connection.
+
+The below example demonstrates how to instantiate the SQLValueCheckOperator task.
+
+.. exampleinclude:: /../tests/system/common/sql/example_sql_value_check.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_sql_value_check]
+    :end-before: [END howto_operator_sql_value_check]
+
 .. _howto/operator:SQLThresholdCheckOperator:
 
 Check values against a threshold

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/provider.yaml

@@ -22,12 +22,15 @@ description: |
     `Common SQL Provider <https://en.wikipedia.org/wiki/SQL>`__
 
 state: ready
-source-date-epoch: 1765298925
+source-date-epoch: 1769460904
 # Note that those versions are maintained by release manager - do not update them manually
 # with the exception of case where other provider in sources has >= new provider version.
 # In such case adding >= NEW_VERSION and bumping to NEW_VERSION in a provider have
 # to be done in the same PR
 versions:
+  - 1.30.4
+  - 1.30.3
+  - 1.30.2
   - 1.30.1
   - 1.30.0
   - 1.29.0

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-common-sql"
-version = "1.30.1"
+version = "1.30.4"
 description = "Provider package apache-airflow-providers-common-sql for Apache Airflow"
 readme = "README.rst"
 license = "Apache-2.0"
@@ -59,7 +59,7 @@ requires-python = ">=3.10"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "apache-airflow>=2.11.0",
-    "apache-airflow-providers-common-compat>=1.10.1",
+    "apache-airflow-providers-common-compat>=1.12.0",
     "sqlparse>=0.5.1",
     "more-itertools>=9.0.0",
     # The methodtools dependency is necessary since the introduction of dialects:
@@ -82,6 +82,9 @@ dependencies = [
 "polars" = [
     "polars>=1.26.0"
 ]
+"sqlalchemy" = [
+    "sqlalchemy>=1.4.49",
+]
 
 [dependency-groups]
 dev = [
@@ -96,6 +99,7 @@ dev = [
     "apache-airflow-providers-mysql",
     "apache-airflow-providers-postgres",
     "apache-airflow-providers-odbc",
+    "apache-airflow-providers-common-sql[sqlalchemy]"
 ]
 
 # To build docs:
@@ -124,8 +128,8 @@ apache-airflow-providers-common-sql = {workspace = true}
 apache-airflow-providers-standard = {workspace = true}
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.1"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.1/changelog.html"
+"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.4"
+"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-common-sql/1.30.4/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/src/airflow/providers/common/sql/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "1.30.1"
+__version__ = "1.30.4"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.11.0"

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/src/airflow/providers/common/sql/hooks/sql.py

@@ -29,16 +29,29 @@ import sqlparse
 from deprecated import deprecated
 from methodtools import lru_cache
 from more_itertools import chunked
-from sqlalchemy import create_engine, inspect
-from sqlalchemy.engine import make_url
-from sqlalchemy.exc import ArgumentError, NoSuchModuleError
 
-from airflow.configuration import conf
-from airflow.exceptions import AirflowOptionalProviderFeatureException, AirflowProviderDeprecationWarning
-from airflow.providers.common.compat.sdk import AirflowException, BaseHook
+try:
+    from sqlalchemy import create_engine, inspect
+    from sqlalchemy.engine import make_url
+    from sqlalchemy.exc import ArgumentError, NoSuchModuleError
+except ImportError:
+    create_engine = None  # type: ignore[assignment]
+    inspect = None  # type: ignore[assignment]
+    make_url = None  # type: ignore[assignment]
+    ArgumentError = Exception  # type: ignore[misc,assignment]
+    NoSuchModuleError = Exception  # type: ignore[misc,assignment]
+
+
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.module_loading import import_string
+from airflow.providers.common.compat.sdk import (
+    AirflowException,
+    AirflowOptionalProviderFeatureException,
+    BaseHook,
+    conf,
+)
 from airflow.providers.common.sql.dialects.dialect import Dialect
 from airflow.providers.common.sql.hooks import handlers
-from airflow.utils.module_loading import import_string
 
 if TYPE_CHECKING:
     from pandas import DataFrame as PandasDataFrame
@@ -304,6 +317,12 @@ class DbApiHook(BaseHook):
         :param engine_kwargs: Kwargs used in :func:`~sqlalchemy.create_engine`.
         :return: the created engine.
         """
+        if create_engine is None:
+            raise AirflowOptionalProviderFeatureException(
+                "SQLAlchemy is required to generate the connection URI. "
+                "Install it with: pip install 'apache-airflow-providers-common-sql[sqlalchemy]'"
+            )
+
         if engine_kwargs is None:
             engine_kwargs = {}
 
@@ -318,22 +337,30 @@ class DbApiHook(BaseHook):
 
     @property
     def inspector(self) -> Inspector:
+        if inspect is None:
+            raise AirflowOptionalProviderFeatureException(
+                "SQLAlchemy is required for database inspection. "
+                "Install it with: pip install 'apache-airflow-providers-common-sql[sqlalchemy]'"
+            )
         return inspect(self.get_sqlalchemy_engine())
 
     @cached_property
     def dialect_name(self) -> str:
-        try:
-            return make_url(self.get_uri()).get_dialect().name
-        except (ArgumentError, NoSuchModuleError, ValueError):
-            config = self.connection_extra
-            sqlalchemy_scheme = config.get("sqlalchemy_scheme")
-            if sqlalchemy_scheme:
-                return sqlalchemy_scheme.split("+")[0] if "+" in sqlalchemy_scheme else sqlalchemy_scheme
-            return config.get("dialect", "default")
+        if make_url is not None:
+            try:
+                return make_url(self.get_uri()).get_dialect().name
+            except (ArgumentError, NoSuchModuleError, ValueError):
+                pass
+
+        config = self.connection_extra
+        sqlalchemy_scheme = config.get("sqlalchemy_scheme")
+        if sqlalchemy_scheme:
+            return sqlalchemy_scheme.split("+")[0] if "+" in sqlalchemy_scheme else sqlalchemy_scheme
+        return config.get("dialect", "default")
 
     @cached_property
     def dialect(self) -> Dialect:
-        from airflow.utils.module_loading import import_string
+        from airflow.providers.common.compat.module_loading import import_string
 
         dialect_info = self._dialects.get(self.dialect_name)
 

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/src/airflow/providers/common/sql/operators/generic_transfer.pyi

@@ -38,7 +38,7 @@ from _typeshed import Incomplete as Incomplete
 
 from airflow.models import BaseOperator
 from airflow.providers.common.sql.hooks.sql import DbApiHook as DbApiHook
-from airflow.utils.context import Context as Context
+from airflow.sdk import Context
 
 class GenericTransfer(BaseOperator):
     template_fields: Sequence[str]

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/src/airflow/providers/common/sql/operators/sql.py

@@ -862,6 +862,9 @@ class SQLValueCheckOperator(BaseSQLOperator):
     :param sql: the sql to be executed. (templated)
     :param conn_id: the connection ID used to connect to the database.
     :param database: name of database which overwrite the defined one in connection
+    :param pass_value: the value to check against
+    :param tolerance: (optional) the tolerance allowed to pass records within for
+        numeric queries
     """
 
     __mapper_args__ = {"polymorphic_identity": "SQLValueCheckOperator"}

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/src/airflow/providers/common/sql/sensors/sql.pyi

@@ -38,7 +38,7 @@ from typing import Any
 from _typeshed import Incomplete as Incomplete
 
 from airflow.providers.common.compat.sdk import BaseSensorOperator
-from airflow.utils.context import Context as Context
+from airflow.sdk import Context
 
 class SqlSensor(BaseSensorOperator):
     template_fields: Sequence[str]

apache_airflow_providers_common_sql-1.30.4/tests/system/common/sql/example_sql_value_check.py

@@ -0,0 +1,62 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from airflow import DAG
+from airflow.providers.common.sql.operators.sql import SQLValueCheckOperator
+from airflow.sdk.timezone import datetime
+
+connection_args = {
+    "conn_id": "sales_db",
+    "conn_type": "Postgres",
+    "host": "postgres",
+    "schema": "postgres",
+    "login": "postgres",
+    "password": "postgres",
+    "port": 5432,
+}
+
+with DAG(
+    "example_sql_value_check_query",
+    description="Example DAG for SQLValueCheckOperator.",
+    default_args=connection_args,
+    start_date=datetime(2025, 12, 15),
+    schedule=None,
+    catchup=False,
+) as dag:
+    """
+    ### Example SQL value check DAG
+
+    Runs the SQLValueCheckOperator against the Airflow metadata DB.
+    """
+
+    # [START howto_operator_sql_value_check]
+    value_check = SQLValueCheckOperator(
+        task_id="threshhold_check",
+        conn_id="sales_db",
+        sql="SELECT count(distinct(customer_id)) FROM sales LIMIT 50;",
+        pass_value=40,
+        tolerance=5,
+    )
+    # [END howto_operator_sql_value_check]
+
+
+from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)

{apache_airflow_providers_common_sql-1.30.1 → apache_airflow_providers_common_sql-1.30.4}/tests/unit/common/sql/operators/test_sql.py

@@ -18,21 +18,13 @@
 from __future__ import annotations
 
 import datetime
+import importlib.util
 import inspect
 from unittest import mock
 from unittest.mock import MagicMock
 
 import pytest
 
-try:
-    import importlib.util
-
-    if not importlib.util.find_spec("airflow.sdk.bases.hook"):
-        raise ImportError
-
-    BASEHOOK_PATCH_PATH = "airflow.sdk.bases.hook.BaseHook"
-except ImportError:
-    BASEHOOK_PATCH_PATH = "airflow.hooks.base.BaseHook"
 from airflow import DAG
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.models import Connection
@@ -60,8 +52,19 @@ from tests_common.test_utils.dag import sync_dag_to_db
 from tests_common.test_utils.db import clear_db_dag_bundles, clear_db_dags, clear_db_runs, clear_db_xcom
 from tests_common.test_utils.markers import skip_if_force_lowest_dependencies_marker
 from tests_common.test_utils.providers import get_provider_min_airflow_version
+from tests_common.test_utils.taskinstance import TaskInstanceWrapper
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_1, AIRFLOW_V_3_0_PLUS
 
+try:
+    import importlib.util
+
+    if not importlib.util.find_spec("airflow.sdk.bases.hook"):
+        raise ImportError
+
+    BASEHOOK_PATCH_PATH = "airflow.sdk.bases.hook.BaseHook"
+except ImportError:
+    BASEHOOK_PATCH_PATH = "airflow.hooks.base.BaseHook"
+
 if AIRFLOW_V_3_0_PLUS:
     from airflow.utils.types import DagRunTriggeredByType
 
@@ -1151,7 +1154,13 @@ class TestSqlBranch:
             )
 
         ti = dr.get_task_instance(task_id)
-        ti.task = self.dag.get_task(ti.task_id)
+        task = self.dag.get_task(ti.task_id)
+
+        if AIRFLOW_V_3_0_PLUS:
+            ti.task = self.scheduler_dag.get_task(ti.task_id)
+            ti = TaskInstanceWrapper(ti, task)
+        else:
+            ti.task = task
 
         return ti