dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067725414
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,87 @@ def stop(args):
logger.info("Worker has been shut down.")
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(
+ hostname: str | None = None, states: list | None = None, session: Session
= NEW_SESSION
+) -> list:
+ from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
+
+ query = select(EdgeWorkerModel)
+ if states:
+ query = query.where(EdgeWorkerModel.state.in_(states))
+ if hostname:
+ query = query.where(EdgeWorkerModel.worker_name == hostname)
+ query = query.order_by(EdgeWorkerModel.worker_name)
+ return session.scalars(query).all()
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+@provide_session
+def list_edge_workers(args, session: Session = NEW_SESSION) -> None:
+ """Query the db to list all registered edge workers."""
+ all_hosts_iter = _fetch_edge_hosts_from_db(states=args.state,
session=session)
+ # Format and print worker info on the screen
+ fields = [
+ "worker_name",
+ "state",
+ "queues",
+ "maintenance_comment",
+ ]
+ all_hosts = [{f: str(host.__getattribute__(f)) for f in fields} for host
in all_hosts_iter]
+ AirflowConsole().print_as(data=all_hosts, output=args.output)
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def put_remote_worker_on_maintenance(args) -> None:
+ """Put remote edge worker on maintenance."""
+ if not _fetch_edge_hosts_from_db(hostname=args.edge_hostname):
+ raise SystemExit(f"Error: Edge Worker {args.edge_hostname} is
unknown!")
+ from airflow.providers.edge3.models.edge_worker import request_maintenance
+
+ request_maintenance(args.edge_hostname, args.comments)
+ logger.info("%s has been put on maintenance by %s.", args.edge_hostname,
getuser())
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def remove_remote_worker_from_maintenance(args) -> None:
+ """Remove remote edge worker from maintenance."""
+ if not _fetch_edge_hosts_from_db(hostname=args.edge_hostname):
+ raise SystemExit(f"Error: Edge Worker {args.edge_hostname} is
unknown!")
+ from airflow.providers.edge3.models.edge_worker import exit_maintenance
+
+ exit_maintenance(args.edge_hostname)
+ logger.info("%s has been removed from maintenance by %s.",
args.edge_hostname, getuser())
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def remote_worker_update_maintenance_comment(args) -> None:
+ """Update maintenance comments of the remote edge worker."""
+ if not _fetch_edge_hosts_from_db(hostname=args.edge_hostname):
+ raise SystemExit(f"Error: Edge Worker {args.edge_hostname} is
unknown!")
+ from airflow.providers.edge3.models.edge_worker import
change_maintenance_comment
+
+ change_maintenance_comment(args.edge_hostname, args.comments)
+ logger.info("Maintenance comments updated for %s by %s.",
args.edge_hostname, getuser())
Review Comment:
I have addressed this. Let me know if you like the implementation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
