replace try-except import pattern with AIRFLOW_V_3_2_PLUS checks

Author: xBis7

providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py

@@ -29,6 +29,7 @@
 from airflow.executors.base_executor import BaseExecutor
 from airflow.models.taskinstance import TaskInstance
 from airflow.providers.common.compat.sdk import Stats, timezone
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS
 from airflow.providers.edge3.models.db import EdgeDBManager, check_db_manager_config
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
 from airflow.providers.edge3.models.edge_logs import EdgeLogsModel
@@ -37,11 +38,6 @@
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.state import TaskInstanceState
 
-try:
-    from airflow.sdk.observability.stats import DualStatsManager
-except ImportError:
-    DualStatsManager = None  # type: ignore[assignment,misc]  # Airflow < 3.2 compat
-
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
@@ -197,7 +193,9 @@ def _update_orphaned_jobs(self, session: Session) -> bool:
                     "queue": job.queue,
                     "state": str(TaskInstanceState.FAILED),
                 }
-                if DualStatsManager is not None:
+                if AIRFLOW_V_3_2_PLUS:
+                    from airflow.sdk.observability.stats import DualStatsManager
+
                     DualStatsManager.incr("edge_worker.ti.finish", tags={}, legacy_name_tags=tags)
                 else:
                     Stats.incr(

providers/edge3/src/airflow/providers/edge3/models/edge_worker.py

@@ -27,13 +27,9 @@
 from sqlalchemy.orm import Mapped
 
 from airflow.providers.common.compat.sdk import AirflowException, Stats, timezone
-from airflow.providers.edge3.models.edge_base import Base
-
-try:
-    from airflow.sdk.observability.stats import DualStatsManager
-except ImportError:
-    DualStatsManager = None  # type: ignore[assignment,misc]  # Airflow < 3.2 compat
 from airflow.providers.common.compat.sqlalchemy.orm import mapped_column
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS
+from airflow.providers.edge3.models.edge_base import Base
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.providers_configuration_loader import providers_configuration_loaded
 from airflow.utils.session import NEW_SESSION, provide_session
@@ -178,7 +174,9 @@ def set_metrics(
         EdgeWorkerState.OFFLINE_MAINTENANCE,
     )
 
-    if DualStatsManager is not None:
+    if AIRFLOW_V_3_2_PLUS:
+        from airflow.sdk.observability.stats import DualStatsManager
+
         DualStatsManager.gauge(
             "edge_worker.connected",
             int(connected),

providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py

@@ -27,11 +27,7 @@
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.executors.workloads import ExecuteTask
 from airflow.providers.common.compat.sdk import Stats, timezone
-
-try:
-    from airflow.sdk.observability.stats import DualStatsManager
-except ImportError:
-    DualStatsManager = None  # type: ignore[assignment,misc]  # Airflow < 3.2 compat
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
 from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest
 from airflow.providers.edge3.worker_api.datamodels import (
@@ -41,6 +37,9 @@
 )
 from airflow.utils.state import TaskInstanceState
 
+if AIRFLOW_V_3_2_PLUS:
+    from airflow.sdk.observability.stats import DualStatsManager
+
 jobs_router = AirflowRouter(tags=["Jobs"], prefix="/jobs")
 
 
@@ -91,7 +90,7 @@ def fetch(
     session.commit()
     # Edge worker does not backport emitted Airflow metrics, so export some metrics
     tags = {"dag_id": job.dag_id, "task_id": job.task_id, "queue": job.queue}
-    if DualStatsManager is not None:
+    if AIRFLOW_V_3_2_PLUS:
         DualStatsManager.incr("edge_worker.ti.start", tags={}, legacy_name_tags=tags)
     else:
         Stats.incr(f"edge_worker.ti.start.{job.queue}.{job.dag_id}.{job.task_id}", tags=tags)
@@ -148,7 +147,7 @@ def state(
                 "queue": job.queue,
                 "state": state,
             }
-            if DualStatsManager is not None:
+            if AIRFLOW_V_3_2_PLUS:
                 DualStatsManager.incr("edge_worker.ti.finish", tags={}, legacy_name_tags=tags)
             else:
                 Stats.incr(

providers/edge3/src/airflow/providers/edge3/worker_api/routes/worker.py

@@ -27,11 +27,7 @@
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.providers.common.compat.sdk import Stats, timezone
-
-try:
-    from airflow.sdk.observability.stats import DualStatsManager
-except ImportError:
-    DualStatsManager = None  # type: ignore[assignment,misc]  # Airflow < 3.2 compat
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS
 from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState, set_metrics
 from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest
 from airflow.providers.edge3.worker_api.datamodels import (
@@ -217,7 +213,9 @@ def set_state(
     worker.sysinfo = json.dumps(body.sysinfo)
     worker.last_update = timezone.utcnow()
     session.commit()
-    if DualStatsManager is not None:
+    if AIRFLOW_V_3_2_PLUS:
+        from airflow.sdk.observability.stats import DualStatsManager
+
         DualStatsManager.incr(
             "edge_worker.heartbeat_count",
             1,

providers/edge3/tests/unit/edge3/executors/test_edge_executor.py

@@ -33,8 +33,9 @@
 from airflow.utils.state import TaskInstanceState
 
 from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS
 
-try:
+if AIRFLOW_V_3_2_PLUS:
     from airflow.sdk._shared.observability.metrics.dual_stats_manager import DualStatsManager  # noqa: F401
 
     stats_reference = "airflow.sdk._shared.observability.metrics.dual_stats_manager.DualStatsManager"
@@ -48,7 +49,7 @@
         },
     }
     expected_call_count = 1
-except ImportError:
+else:
     from airflow.providers.common.compat.sdk import Stats
 
     stats_reference = f"{Stats.__module__}.Stats"

providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py

@@ -27,6 +27,8 @@
 from airflow.utils.session import create_session
 from airflow.utils.state import TaskInstanceState
 
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS
+
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
@@ -35,7 +37,7 @@
 RUN_ID = "manual__2024-11-24T21:03:01+01:00"
 QUEUE = "test"
 
-try:
+if AIRFLOW_V_3_2_PLUS:
     from airflow.sdk._shared.observability.metrics.dual_stats_manager import DualStatsManager  # noqa: F401
 
     stats_reference = "airflow.sdk._shared.observability.metrics.dual_stats_manager.DualStatsManager"
@@ -49,7 +51,7 @@
         },
     }
     expected_call_count = 1
-except ImportError:
+else:
     from airflow.providers.common.compat.sdk import Stats
 
     stats_reference = f"{Stats.__module__}.Stats"

providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py

@@ -36,11 +36,7 @@
 )
 
 from airflow.providers.common.compat.sdk import Stats, conf as airflow_conf
-
-try:
-    from airflow.sdk.observability.stats import DualStatsManager
-except ImportError:
-    DualStatsManager = None  # type: ignore[assignment,misc]  # Airflow < 3.2 compat
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS
 from airflow.providers.openlineage import __version__ as OPENLINEAGE_PROVIDER_VERSION, conf
 from airflow.providers.openlineage.utils.utils import (
     OpenLineageRedactor,
@@ -163,7 +159,9 @@ def emit(self, event: RunEvent):
 
         try:
             with ExitStack() as stack:
-                if DualStatsManager is not None:
+                if AIRFLOW_V_3_2_PLUS:
+                    from airflow.sdk.observability.stats import DualStatsManager
+
                     stack.enter_context(
                         DualStatsManager.timer(
                             "ol.emit.attempts",

providers/openlineage/tests/unit/openlineage/plugins/test_adapter.py

@@ -60,11 +60,11 @@
 from tests_common.test_utils.taskinstance import create_task_instance
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS
 
-try:
+if AIRFLOW_V_3_2_PLUS:
     from airflow.sdk._shared.observability.metrics.dual_stats_manager import DualStatsManager  # noqa: F401
 
     stats_reference = "airflow.sdk._shared.observability.metrics.dual_stats_manager.DualStatsManager"
-except ImportError:
+else:
     stats_reference = "airflow.providers.openlineage.plugins.adapter.Stats"