apache-airflow-providers-edge3 2.0.0rc1__py3-none-any.whl → 3.0.1rc1__py3-none-any.whl

airflow/providers/edge3/executors/edge_executor.py

@@ -23,17 +23,15 @@ from copy import deepcopy
 from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any
 
-from sqlalchemy import delete, inspect, text
+from sqlalchemy import delete, inspect, select, text
 from sqlalchemy.exc import NoSuchTableError
 from sqlalchemy.orm import Session
 
-from airflow.cli.cli_config import GroupCommand
 from airflow.configuration import conf
 from airflow.executors import workloads
 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.edge3.cli.edge_command import EDGE_COMMANDS
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
 from airflow.providers.edge3.models.edge_logs import EdgeLogsModel
 from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState, reset_metrics
@@ -42,10 +40,9 @@ from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
-    import argparse
-
     from sqlalchemy.engine.base import Engine
 
+    from airflow.cli.cli_config import GroupCommand
     from airflow.models.taskinstancekey import TaskInstanceKey
 
     # TODO: Airflow 2 type hints; remove when Airflow 2 support is removed
@@ -140,17 +137,15 @@ class EdgeExecutor(BaseExecutor):
         key = task_instance.key
 
         # Check if job already exists with same dag_id, task_id, run_id, map_index, try_number
-        existing_job = (
-            session.query(EdgeJobModel)
-            .filter_by(
-                dag_id=key.dag_id,
-                task_id=key.task_id,
-                run_id=key.run_id,
-                map_index=key.map_index,
-                try_number=key.try_number,
+        existing_job = session.scalars(
+            select(EdgeJobModel).where(
+                EdgeJobModel.dag_id == key.dag_id,
+                EdgeJobModel.task_id == key.task_id,
+                EdgeJobModel.run_id == key.run_id,
+                EdgeJobModel.map_index == key.map_index,
+                EdgeJobModel.try_number == key.try_number,
             )
-            .first()
-        )
+        ).first()
 
         if existing_job:
             existing_job.state = TaskInstanceState.QUEUED
@@ -176,10 +171,10 @@ class EdgeExecutor(BaseExecutor):
         """Reset worker state if heartbeat timed out."""
         changed = False
         heartbeat_interval: int = conf.getint("edge", "heartbeat_interval")
-        lifeless_workers: list[EdgeWorkerModel] = (
-            session.query(EdgeWorkerModel)
+        lifeless_workers: Sequence[EdgeWorkerModel] = session.scalars(
+            select(EdgeWorkerModel)
             .with_for_update(skip_locked=True)
-            .filter(
+            .where(
                 EdgeWorkerModel.state.not_in(
                     [
                         EdgeWorkerState.UNKNOWN,
@@ -189,8 +184,7 @@ class EdgeExecutor(BaseExecutor):
                 ),
                 EdgeWorkerModel.last_update < (timezone.utcnow() - timedelta(seconds=heartbeat_interval * 5)),
             )
-            .all()
-        )
+        ).all()
 
         for worker in lifeless_workers:
             changed = True
@@ -212,15 +206,14 @@ class EdgeExecutor(BaseExecutor):
     def _update_orphaned_jobs(self, session: Session) -> bool:
         """Update status ob jobs when workers die and don't update anymore."""
         heartbeat_interval: int = conf.getint("scheduler", "task_instance_heartbeat_timeout")
-        lifeless_jobs: list[EdgeJobModel] = (
-            session.query(EdgeJobModel)
+        lifeless_jobs: Sequence[EdgeJobModel] = session.scalars(
+            select(EdgeJobModel)
             .with_for_update(skip_locked=True)
-            .filter(
+            .where(
                 EdgeJobModel.state == TaskInstanceState.RUNNING,
                 EdgeJobModel.last_update < (timezone.utcnow() - timedelta(seconds=heartbeat_interval)),
             )
-            .all()
-        )
+        ).all()
 
         for job in lifeless_jobs:
             ti = TaskInstance.get_task_instance(
@@ -254,10 +247,10 @@ class EdgeExecutor(BaseExecutor):
         purged_marker = False
         job_success_purge = conf.getint("edge", "job_success_purge")
         job_fail_purge = conf.getint("edge", "job_fail_purge")
-        jobs: list[EdgeJobModel] = (
-            session.query(EdgeJobModel)
+        jobs: Sequence[EdgeJobModel] = session.scalars(
+            select(EdgeJobModel)
             .with_for_update(skip_locked=True)
-            .filter(
+            .where(
                 EdgeJobModel.state.in_(
                     [
                         TaskInstanceState.RUNNING,
@@ -269,8 +262,7 @@ class EdgeExecutor(BaseExecutor):
                     ]
                 )
             )
-            .all()
-        )
+        ).all()
 
         # Sync DB with executor otherwise runs out of sync in multi scheduler deployment
         already_removed = self.running - set(job.key for job in jobs)
@@ -388,23 +380,6 @@ class EdgeExecutor(BaseExecutor):
 
     @staticmethod
     def get_cli_commands() -> list[GroupCommand]:
-        return [
-            GroupCommand(
-                name="edge",
-                help="Edge Worker components",
-                description=(
-                    "Start and manage Edge Worker. Works only when using EdgeExecutor. For more information, "
-                    "see https://airflow.apache.org/docs/apache-airflow-providers-edge3/stable/edge_executor.html"
-                ),
-                subcommands=EDGE_COMMANDS,
-            ),
-        ]
-
-
-def _get_parser() -> argparse.ArgumentParser:
-    """
-    Generate documentation; used by Sphinx.
+        from airflow.providers.edge3.cli.definition import get_edge_cli_commands
 
-    :meta private:
-    """
-    return EdgeExecutor._get_parser()
+        return get_edge_cli_commands()

airflow/providers/edge3/get_provider_info.py

@@ -32,6 +32,7 @@ def get_provider_info():
                 "plugin-class": "airflow.providers.edge3.plugins.edge_executor_plugin.EdgeExecutorPlugin",
             }
         ],
+        "cli": ["airflow.providers.edge3.cli.definition.get_edge_cli_commands"],
         "executors": ["airflow.providers.edge3.executors.EdgeExecutor"],
         "config": {
             "edge": {

airflow/providers/edge3/models/edge_job.py

@@ -27,8 +27,7 @@ from sqlalchemy import (
 from sqlalchemy.orm import Mapped
 
 from airflow.models.base import Base, StringID
-from airflow.models.taskinstancekey import TaskInstanceKey
-from airflow.providers.common.compat.sdk import timezone
+from airflow.providers.common.compat.sdk import TaskInstanceKey, timezone
 from airflow.providers.common.compat.sqlalchemy.orm import mapped_column
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.sqlalchemy import UtcDateTime

airflow/providers/edge3/models/edge_worker.py

@@ -173,27 +173,72 @@ def set_metrics(
         EdgeWorkerState.OFFLINE_MAINTENANCE,
     )
 
-    Stats.gauge(f"edge_worker.connected.{worker_name}", int(connected))
-    Stats.gauge("edge_worker.connected", int(connected), tags={"worker_name": worker_name})
+    try:
+        from airflow.sdk._shared.observability.metrics.dual_stats_manager import DualStatsManager
+
+        DualStatsManager.gauge(
+            "edge_worker.connected",
+            int(connected),
+            tags={},
+            extra_tags={"worker_name": worker_name},
+        )
 
-    Stats.gauge(f"edge_worker.maintenance.{worker_name}", int(maintenance))
-    Stats.gauge("edge_worker.maintenance", int(maintenance), tags={"worker_name": worker_name})
+        DualStatsManager.gauge(
+            "edge_worker.maintenance",
+            int(maintenance),
+            tags={},
+            extra_tags={"worker_name": worker_name},
+        )
 
-    Stats.gauge(f"edge_worker.jobs_active.{worker_name}", jobs_active)
-    Stats.gauge("edge_worker.jobs_active", jobs_active, tags={"worker_name": worker_name})
+        DualStatsManager.gauge(
+            "edge_worker.jobs_active",
+            jobs_active,
+            tags={},
+            extra_tags={"worker_name": worker_name},
+        )
 
-    Stats.gauge(f"edge_worker.concurrency.{worker_name}", concurrency)
-    Stats.gauge("edge_worker.concurrency", concurrency, tags={"worker_name": worker_name})
+        DualStatsManager.gauge(
+            "edge_worker.concurrency",
+            concurrency,
+            tags={},
+            extra_tags={"worker_name": worker_name},
+        )
 
-    Stats.gauge(f"edge_worker.free_concurrency.{worker_name}", free_concurrency)
-    Stats.gauge("edge_worker.free_concurrency", free_concurrency, tags={"worker_name": worker_name})
+        DualStatsManager.gauge(
+            "edge_worker.free_concurrency",
+            free_concurrency,
+            tags={},
+            extra_tags={"worker_name": worker_name},
+        )
 
-    Stats.gauge(f"edge_worker.num_queues.{worker_name}", len(queues))
-    Stats.gauge(
-        "edge_worker.num_queues",
-        len(queues),
-        tags={"worker_name": worker_name, "queues": ",".join(queues)},
-    )
+        DualStatsManager.gauge(
+            "edge_worker.num_queues",
+            len(queues),
+            tags={},
+            extra_tags={"worker_name": worker_name, "queues": ",".join(queues)},
+        )
+    except ImportError:
+        Stats.gauge(f"edge_worker.connected.{worker_name}", int(connected))
+        Stats.gauge("edge_worker.connected", int(connected), tags={"worker_name": worker_name})
+
+        Stats.gauge(f"edge_worker.maintenance.{worker_name}", int(maintenance))
+        Stats.gauge("edge_worker.maintenance", int(maintenance), tags={"worker_name": worker_name})
+
+        Stats.gauge(f"edge_worker.jobs_active.{worker_name}", jobs_active)
+        Stats.gauge("edge_worker.jobs_active", jobs_active, tags={"worker_name": worker_name})
+
+        Stats.gauge(f"edge_worker.concurrency.{worker_name}", concurrency)
+        Stats.gauge("edge_worker.concurrency", concurrency, tags={"worker_name": worker_name})
+
+        Stats.gauge(f"edge_worker.free_concurrency.{worker_name}", free_concurrency)
+        Stats.gauge("edge_worker.free_concurrency", free_concurrency, tags={"worker_name": worker_name})
+
+        Stats.gauge(f"edge_worker.num_queues.{worker_name}", len(queues))
+        Stats.gauge(
+            "edge_worker.num_queues",
+            len(queues),
+            tags={"worker_name": worker_name, "queues": ",".join(queues)},
+        )
 
 
 def reset_metrics(worker_name: str) -> None:

airflow/providers/edge3/plugins/edge_executor_plugin.py

@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowConfigException
-from airflow.plugins_manager import AirflowPlugin
+from airflow.providers.common.compat.sdk import AirflowPlugin
 from airflow.providers.edge3.version_compat import AIRFLOW_V_3_1_PLUS
 from airflow.utils.session import NEW_SESSION, provide_session