Add queue and remove queue cli commands for EdgeExecutor (#53505)

* Add queue and remove queue cli commands for EdgeExecutor

* Add queue cli commands to docs

Author: dheerajturaga

providers/edge3/docs/deployment.rst

@@ -190,3 +190,5 @@ instance. The commands are:
 - ``airflow edge remote-edge-worker-exit-maintenance``: Request a remote edge worker to exit maintenance mode
 - ``airflow edge shutdown-remote-edge-worker``: Shuts down a remote edge worker gracefully
 - ``airflow edge remove-remote-edge-worker``: Remove a worker instance from the cluster
+- ``airflow edge add-worker-queues``: Add queues to an edge worker
+- ``airflow edge remove-worker-queues``: Remove queues from an edge worker

providers/edge3/src/airflow/providers/edge3/cli/edge_command.py

@@ -351,6 +351,48 @@ def remote_worker_request_shutdown(args) -> None:
     logger.info("Requested shutdown of Edge Worker host %s by %s.", args.edge_hostname, getuser())
 
 
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def add_worker_queues(args) -> None:
+    """Add queues to an edge worker."""
+    _check_valid_db_connection()
+    _check_if_registered_edge_host(hostname=args.edge_hostname)
+    from airflow.providers.edge3.models.edge_worker import add_worker_queues
+
+    queues = args.queues.split(",") if args.queues else []
+    if not queues:
+        raise SystemExit("Error: No queues specified to add.")
+
+    try:
+        add_worker_queues(args.edge_hostname, queues)
+        logger.info("Added queues %s to Edge Worker host %s by %s.", queues, args.edge_hostname, getuser())
+    except TypeError as e:
+        logger.error(str(e))
+        raise SystemExit
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def remove_worker_queues(args) -> None:
+    """Remove queues from an edge worker."""
+    _check_valid_db_connection()
+    _check_if_registered_edge_host(hostname=args.edge_hostname)
+    from airflow.providers.edge3.models.edge_worker import remove_worker_queues
+
+    queues = args.queues.split(",") if args.queues else []
+    if not queues:
+        raise SystemExit("Error: No queues specified to remove.")
+
+    try:
+        remove_worker_queues(args.edge_hostname, queues)
+        logger.info(
+            "Removed queues %s from Edge Worker host %s by %s.", queues, args.edge_hostname, getuser()
+        )
+    except TypeError as e:
+        logger.error(str(e))
+        raise SystemExit
+
+
 ARG_CONCURRENCY = Arg(
     ("-c", "--concurrency"),
     type=int,
@@ -380,6 +422,11 @@ def remote_worker_request_shutdown(args) -> None:
     help="Maintenance comments to report reason. Required if enabling maintenance",
     required=True,
 )
+ARG_QUEUES_MANAGE = Arg(
+    ("-q", "--queues"),
+    help="Comma delimited list of queues to add or remove.",
+    required=True,
+)
 ARG_WAIT_MAINT = Arg(
     ("-w", "--wait"),
     default=False,
@@ -516,4 +563,22 @@ def remote_worker_request_shutdown(args) -> None:
         func=remote_worker_request_shutdown,
         args=(ARG_REQUIRED_EDGE_HOSTNAME,),
     ),
+    ActionCommand(
+        name="add-worker-queues",
+        help=add_worker_queues.__doc__,
+        func=add_worker_queues,
+        args=(
+            ARG_REQUIRED_EDGE_HOSTNAME,
+            ARG_QUEUES_MANAGE,
+        ),
+    ),
+    ActionCommand(
+        name="remove-worker-queues",
+        help=remove_worker_queues.__doc__,
+        func=remove_worker_queues,
+        args=(
+            ARG_REQUIRED_EDGE_HOSTNAME,
+            ARG_QUEUES_MANAGE,
+        ),
+    ),
 ]

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

@@ -283,3 +283,37 @@ def request_shutdown(worker_name: str, session: Session = NEW_SESSION) -> None:
         EdgeWorkerState.UNKNOWN,
     ):
         worker.state = EdgeWorkerState.SHUTDOWN_REQUEST
+
+
+@provide_session
+def add_worker_queues(worker_name: str, queues: list[str], session: Session = NEW_SESSION) -> None:
+    """Add queues to an edge worker."""
+    query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
+    worker: EdgeWorkerModel = session.scalar(query)
+    if worker.state in (
+        EdgeWorkerState.OFFLINE,
+        EdgeWorkerState.OFFLINE_MAINTENANCE,
+        EdgeWorkerState.UNKNOWN,
+    ):
+        error_message = f"Cannot add queues to edge worker {worker_name} as it is in {worker.state} state!"
+        logger.error(error_message)
+        raise TypeError(error_message)
+    worker.add_queues(queues)
+
+
+@provide_session
+def remove_worker_queues(worker_name: str, queues: list[str], session: Session = NEW_SESSION) -> None:
+    """Remove queues from an edge worker."""
+    query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
+    worker: EdgeWorkerModel = session.scalar(query)
+    if worker.state in (
+        EdgeWorkerState.OFFLINE,
+        EdgeWorkerState.OFFLINE_MAINTENANCE,
+        EdgeWorkerState.UNKNOWN,
+    ):
+        error_message = (
+            f"Cannot remove queues from edge worker {worker_name} as it is in {worker.state} state!"
+        )
+        logger.error(error_message)
+        raise TypeError(error_message)
+    worker.remove_queues(queues)