fix: Condtionally import from airflow core

Author: tomasfarias

airflow_dbt_python/hooks/fs/local.py

@@ -10,7 +10,12 @@
 from pathlib import Path
 from typing import Optional
 
-from airflow.hooks.filesystem import FSHook
+from airflow_dbt_python.utils.version import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.providers.standard.hooks.filesystem import FSHook
+else:
+    from airflow.hooks.filesystem import FSHook
 
 from airflow_dbt_python.hooks.fs import DbtFSHook
 from airflow_dbt_python.utils.url import URL

airflow_dbt_python/hooks/target.py

@@ -24,7 +24,13 @@
     Union,
 )
 
-from airflow.hooks.base import BaseHook
+from airflow_dbt_python.utils.version import AIRFLOW_V_3_1_PLUS
+
+if AIRFLOW_V_3_1_PLUS:
+    from airflow.sdk.bases.hook import BaseHook
+else:
+    from airflow.hooks.base import BaseHook
+
 from airflow.models.connection import Connection
 
 

airflow_dbt_python/operators/dbt.py

@@ -10,10 +10,15 @@
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 from airflow.exceptions import AirflowException
-from airflow.models.baseoperator import BaseOperator
 from airflow.models.xcom import XCOM_RETURN_KEY
 
 from airflow_dbt_python.utils.enums import LogFormat
+from airflow_dbt_python.utils.version import AIRFLOW_V_3_1_PLUS
+
+if AIRFLOW_V_3_1_PLUS:
+    from airflow.sdk import BaseOperator
+else:
+    from airflow.models import BaseOperator
 
 if TYPE_CHECKING:
     from dbt.contracts.results import RunResult

airflow_dbt_python/utils/version.py

@@ -28,3 +28,15 @@
 
 DBT_INSTALLED_1_8 = DBT_1_8 < installed < DBT_1_9
 DBT_INSTALLED_1_9 = DBT_1_9 < installed < DBT_2_0
+
+
+def _get_base_airflow_version_tuple() -> tuple[int, int, int]:
+    from airflow import __version__
+    from packaging.version import Version
+
+    airflow_version = Version(__version__)
+    return airflow_version.major, airflow_version.minor, airflow_version.micro
+
+
+AIRFLOW_V_3_0_PLUS = _get_base_airflow_version_tuple() >= (3, 0, 0)
+AIRFLOW_V_3_1_PLUS = _get_base_airflow_version_tuple() >= (3, 1, 0)

tests/dags/test_dbt_dags.py

@@ -10,9 +10,15 @@
 from dbt.contracts.results import RunStatus, TestStatus
 
 airflow = pytest.importorskip("airflow", minversion="2.2")
+from airflow_dbt_python.utils.version import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG
+else:
+    from airflow import DAG
 
-from airflow import DAG, settings
 from airflow import __version__ as airflow_version
+from airflow import settings
 from airflow.models import DagBag, DagRun
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.utils.state import DagRunState, TaskInstanceState