This is an automated email from the ASF dual-hosted git repository.
dheerajturaga pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 154399c443e Add worker-name filter to `airflow edge list-workers`
(#70095)
154399c443e is described below
commit 154399c443e778d746997a25601ed72c30dce8c7
Author: Shubham Raj <[email protected]>
AuthorDate: Tue Jul 21 01:11:20 2026 +0530
Add worker-name filter to `airflow edge list-workers` (#70095)
* get the name pattern for the list-workers
* Rename list-workers name filter flag to --worker-name-pattern
Match the existing worker_name_pattern naming already used by the edge
REST/UI routes for consistency.
---
providers/edge3/docs/deployment.rst | 4 +-
.../src/airflow/providers/edge3/cli/definition.py | 6 +++
.../airflow/providers/edge3/cli/edge_command.py | 4 +-
.../airflow/providers/edge3/models/edge_worker.py | 22 +++++++-
.../edge3/tests/unit/edge3/cli/test_definition.py | 5 +-
.../edge3/tests/unit/edge3/cli/test_worker.py | 18 +++++++
.../tests/unit/edge3/models/test_edge_worker.py | 60 +++++++++++++++++++++-
7 files changed, 113 insertions(+), 6 deletions(-)
diff --git a/providers/edge3/docs/deployment.rst
b/providers/edge3/docs/deployment.rst
index 68350bfc9d1..ebbdef2e5fb 100644
--- a/providers/edge3/docs/deployment.rst
+++ b/providers/edge3/docs/deployment.rst
@@ -248,7 +248,9 @@ on all workers in the cluster or to check the status of all
workers in the clust
These set of commands need database access, and can only be called on the
central Airflow
instance. The commands are:
-- ``airflow edge list-workers``: List all workers in the cluster
+- ``airflow edge list-workers``: List all workers in the cluster. Accepts an
optional
+ ``--worker-name-pattern`` glob (e.g. ``'prod-*'``) to filter workers by name,
+ and ``-s``/``--state`` to filter by worker state.
- ``airflow edge remote-edge-worker-request-maintenance``: Request a remote
edge worker to enter maintenance mode
- ``airflow edge remote-edge-worker-update-maintenance-comment``: Updates the
maintenance comment for a remote edge worker
- ``airflow edge remote-edge-worker-exit-maintenance``: Request a remote edge
worker to exit maintenance mode
diff --git a/providers/edge3/src/airflow/providers/edge3/cli/definition.py
b/providers/edge3/src/airflow/providers/edge3/cli/definition.py
index f047e7dc808..97a609ad1a9 100644
--- a/providers/edge3/src/airflow/providers/edge3/cli/definition.py
+++ b/providers/edge3/src/airflow/providers/edge3/cli/definition.py
@@ -52,6 +52,11 @@ ARG_EDGE_HOSTNAME = Arg(
("-H", "--edge-hostname"),
help="Set the hostname of worker if you have multiple workers on a single
machine",
)
+ARG_WORKER_NAME_PATTERN = Arg(
+ ("--worker-name-pattern",),
+ metavar="PATTERN",
+ help="Optional glob pattern to filter workers by name, e.g. 'prod-*'.
Lists all workers if omitted.",
+)
ARG_REQUIRED_EDGE_HOSTNAME = Arg(
("-H", "--edge-hostname"),
help="Set the hostname of worker if you have multiple workers on a single
machine",
@@ -183,6 +188,7 @@ EDGE_COMMANDS: list[ActionCommand] = [
args=(
ARG_OUTPUT,
ARG_STATE,
+ ARG_WORKER_NAME_PATTERN,
),
),
ActionCommand(
diff --git a/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
b/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
index 3b3fb6de23d..7a9186f2f57 100644
--- a/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
+++ b/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
@@ -260,7 +260,9 @@ def list_edge_workers(args) -> None:
_check_valid_db_connection()
from airflow.providers.edge3.models.edge_worker import
get_registered_edge_hosts
- all_hosts_iter = get_registered_edge_hosts(states=args.state)
+ all_hosts_iter = get_registered_edge_hosts(
+ states=args.state, worker_name_pattern=args.worker_name_pattern
+ )
# Format and print worker info on the screen
fields = [
"worker_name",
diff --git a/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py
b/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py
index ffc6e20383c..e20ca7db058 100644
--- a/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py
+++ b/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py
@@ -239,11 +239,23 @@ def get_query_filter_by_worker_name(worker_name: str):
return select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name ==
worker_name)
+def _glob_to_like_pattern(pattern: str) -> str:
+ r"""
+ Convert a shell-style glob (``*``, ``?``) to a SQL ``LIKE`` pattern.
+
+ Literal ``LIKE`` metacharacters in the input are escaped with a backslash
so
+ only glob wildcards are treated as special; the caller must pass
``escape="\"``.
+ """
+ escaped = pattern.replace("\\", "\\\\").replace("%", "\\%").replace("_",
"\\_")
+ return escaped.replace("*", "%").replace("?", "_")
+
+
@providers_configuration_loaded
@provide_session
def _fetch_edge_hosts_from_db(
hostname: str | None = None,
states: list | None = None,
+ worker_name_pattern: str | None = None,
*,
session: Session = NEW_SESSION,
) -> Sequence[EdgeWorkerModel]:
@@ -252,14 +264,20 @@ def _fetch_edge_hosts_from_db(
query = query.where(EdgeWorkerModel.state.in_(states))
if hostname:
query = query.where(EdgeWorkerModel.worker_name == hostname)
+ if worker_name_pattern:
+ query = query.where(
+
EdgeWorkerModel.worker_name.like(_glob_to_like_pattern(worker_name_pattern),
escape="\\")
+ )
query = query.order_by(EdgeWorkerModel.worker_name)
return session.scalars(query).all()
@providers_configuration_loaded
@provide_session
-def get_registered_edge_hosts(*, states: list | None = None, session: Session
= NEW_SESSION):
- return _fetch_edge_hosts_from_db(states=states, session=session)
+def get_registered_edge_hosts(
+ *, states: list | None = None, worker_name_pattern: str | None = None,
session: Session = NEW_SESSION
+):
+ return _fetch_edge_hosts_from_db(states=states,
worker_name_pattern=worker_name_pattern, session=session)
@provide_session
diff --git a/providers/edge3/tests/unit/edge3/cli/test_definition.py
b/providers/edge3/tests/unit/edge3/cli/test_definition.py
index f64b7ca281a..cb225e89ffd 100644
--- a/providers/edge3/tests/unit/edge3/cli/test_definition.py
+++ b/providers/edge3/tests/unit/edge3/cli/test_definition.py
@@ -143,7 +143,7 @@ class TestEdgeCliDefinition:
assert args.pid == "/path/to/pid"
def test_list_workers_command_args(self):
- """Test list-workers command with output format and state filter."""
+ """Test list-workers command with output format, state, and
name-pattern filters."""
params = [
"edge",
"list-workers",
@@ -152,10 +152,13 @@ class TestEdgeCliDefinition:
"--state",
"running",
"maintenance",
+ "--worker-name-pattern",
+ "prod-*",
]
args = self.arg_parser.parse_args(params)
assert args.output == "json"
assert args.state == ["running", "maintenance"]
+ assert args.worker_name_pattern == "prod-*"
def test_remote_edge_worker_request_maintenance_args(self):
"""Test remote-edge-worker-request-maintenance command with required
arguments."""
diff --git a/providers/edge3/tests/unit/edge3/cli/test_worker.py
b/providers/edge3/tests/unit/edge3/cli/test_worker.py
index 0e292a7b038..f959237e5ad 100644
--- a/providers/edge3/tests/unit/edge3/cli/test_worker.py
+++ b/providers/edge3/tests/unit/edge3/cli/test_worker.py
@@ -1114,6 +1114,24 @@ class TestEdgeWorker:
assert key in edge_workers[0]
assert any("test_edge_worker" in h["worker_name"] for h in
edge_workers)
+ @pytest.mark.db_test
+ def test_list_edge_workers_passes_name_pattern(self, mock_edgeworker:
EdgeWorkerModel):
+ args = self.parser.parse_args(
+ ["edge", "list-workers", "--output", "json",
"--worker-name-pattern", "prod-*"]
+ )
+ with contextlib.redirect_stdout(StringIO()):
+ with (
+ patch(
+
"airflow.providers.edge3.cli.edge_command._check_valid_db_connection",
+ ),
+ patch(
+
"airflow.providers.edge3.models.edge_worker.get_registered_edge_hosts",
+ return_value=[mock_edgeworker],
+ ) as mock_get_hosts,
+ ):
+ edge_command.list_edge_workers(args)
+ mock_get_hosts.assert_called_once_with(states=None,
worker_name_pattern="prod-*")
+
class TestSignalHandling:
"""Regression tests for the SIGTERM-storm fix (Bug B)."""
diff --git a/providers/edge3/tests/unit/edge3/models/test_edge_worker.py
b/providers/edge3/tests/unit/edge3/models/test_edge_worker.py
index 05c6e16a3e3..85ddefdb02e 100644
--- a/providers/edge3/tests/unit/edge3/models/test_edge_worker.py
+++ b/providers/edge3/tests/unit/edge3/models/test_edge_worker.py
@@ -16,13 +16,26 @@
# under the License.
from __future__ import annotations
+from typing import TYPE_CHECKING
from unittest import mock
+import pytest
+from sqlalchemy import delete
+
from airflow.providers.common.compat.sdk import Stats
-from airflow.providers.edge3.models.edge_worker import EdgeWorkerState,
set_metrics
+from airflow.providers.edge3.models.edge_worker import (
+ EdgeWorkerModel,
+ EdgeWorkerState,
+ _glob_to_like_pattern,
+ get_registered_edge_hosts,
+ set_metrics,
+)
from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+
stats_reference = f"{Stats.__module__}.Stats"
@@ -62,3 +75,48 @@ def test_set_metrics():
legacy_metric_name = f"edge_worker.status.{worker_name}"
assert legacy_metric_name in metric_names
+
+
[email protected](
+ ("glob", "expected"),
+ [
+ ("prod-*", "prod-%"),
+ ("worker-?", "worker-_"),
+ ("*gpu*", "%gpu%"),
+ # Literal LIKE metacharacters are escaped so they are not treated as
wildcards.
+ ("50%_worker", "50\\%\\_worker"),
+ ("back\\slash", "back\\\\slash"),
+ ],
+)
+def test_glob_to_like_pattern(glob, expected):
+ assert _glob_to_like_pattern(glob) == expected
+
+
[email protected]_test
+class TestGetRegisteredEdgeHosts:
+ @pytest.fixture(autouse=True)
+ def setup_test_cases(self, session: Session):
+ session.execute(delete(EdgeWorkerModel))
+ for name in ("prod-worker-1", "prod-worker-2", "dev-worker-1"):
+ session.add(EdgeWorkerModel(worker_name=name, queues=["default"],
state=EdgeWorkerState.RUNNING))
+ session.commit()
+
+ def test_no_pattern_returns_all(self, session: Session):
+ hosts = get_registered_edge_hosts(session=session)
+ assert {h.worker_name for h in hosts} == {
+ "prod-worker-1",
+ "prod-worker-2",
+ "dev-worker-1",
+ }
+
+ def test_star_glob_filters_by_prefix(self, session: Session):
+ hosts = get_registered_edge_hosts(worker_name_pattern="prod-*",
session=session)
+ assert {h.worker_name for h in hosts} == {"prod-worker-1",
"prod-worker-2"}
+
+ def test_question_mark_glob_matches_single_char(self, session: Session):
+ hosts = get_registered_edge_hosts(worker_name_pattern="prod-worker-?",
session=session)
+ assert {h.worker_name for h in hosts} == {"prod-worker-1",
"prod-worker-2"}
+
+ def test_no_match_returns_empty(self, session: Session):
+ hosts = get_registered_edge_hosts(worker_name_pattern="nonexistent-*",
session=session)
+ assert list(hosts) == []
