jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2064458723
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +623,110 @@ def stop(args):
logger.info("Worker has been shut down.")
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(session: Session = NEW_SESSION) -> list:
+ from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
+
+ return
session.scalars(select(EdgeWorkerModel).order_by(EdgeWorkerModel.worker_name)).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(session)
+ # Format and print worker info on the screen
+ fields = [
+ "worker_name",
+ "state",
+ "queues",
+ "jobs_active",
+ "jobs_success",
+ "jobs_failed",
+ "jobs_taken",
Review Comment:
These fields have been planned in DB scheme but are not filled. You can
leave this out.
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +623,110 @@ def stop(args):
logger.info("Worker has been shut down.")
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(session: Session = NEW_SESSION) -> list:
+ from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
+
+ return
session.scalars(select(EdgeWorkerModel).order_by(EdgeWorkerModel.worker_name)).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(session)
+ # Format and print worker info on the screen
+ fields = [
+ "worker_name",
+ "state",
+ "queues",
+ "jobs_active",
+ "jobs_success",
+ "jobs_failed",
+ "jobs_taken",
+ "maintenance_comment",
+ ]
+ all_hosts = [{f: str(host.__getattribute__(f)) for f in fields} for host
in all_hosts_iter]
+ # Filter hosts by state
+ if args.state:
+ all_hosts = [host for host in all_hosts if host["state"] in args.state]
Review Comment:
Can you please pass the filter directly into the DB query?
##########
providers/edge3/pyproject.toml:
##########
@@ -61,6 +61,8 @@ dependencies = [
"apache-airflow-providers-fab>=1.5.3",
"pydantic>=2.11.0",
"retryhttp>=1.2.0,!=1.3.0",
+ "psycopg2-binary>=2.9.10",
+ "asyncpg>=0.30.0"
Review Comment:
Interesting... why are the additional dependencies needed? Are they missing
really (today already)?
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +623,110 @@ def stop(args):
logger.info("Worker has been shut down.")
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(session: Session = NEW_SESSION) -> list:
+ from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
+
+ return
session.scalars(select(EdgeWorkerModel).order_by(EdgeWorkerModel.worker_name)).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(session)
+ # Format and print worker info on the screen
+ fields = [
+ "worker_name",
+ "state",
+ "queues",
+ "jobs_active",
+ "jobs_success",
+ "jobs_failed",
+ "jobs_taken",
+ "maintenance_comment",
+ ]
+ all_hosts = [{f: str(host.__getattribute__(f)) for f in fields} for host
in all_hosts_iter]
+ # Filter hosts by state
+ if args.state:
+ all_hosts = [host for host in all_hosts if host["state"] in args.state]
+ AirflowConsole().print_as(data=all_hosts, output=args.output)
+
+
+@providers_configuration_loaded
+@provide_session
+def _valid_host(hostname, session: Session = NEW_SESSION) -> bool:
+ """
+ Query the database to check if provided hostname exists
+ """
+ all_hosts_iter = _fetch_edge_hosts_from_db(session)
+ for host in all_hosts_iter:
+ if hostname == host.worker_name:
+ return True
+ return False
Review Comment:
I think it would be better to filter directly on DB.
But second thought also... as currently there might not be a very good
validation, how about pushing the validation of existence into the EdgeWorker
model class (and not checking on CLI)? Because in
providers/edge3/src/airflow/providers/edge3/models/edge_worker.py:203 the
validation for existance could be made directly there.
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +624,98 @@ def stop(args):
logger.info("Worker has been shut down.")
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(session: Session = NEW_SESSION) -> list:
+ from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
+
+ return
session.scalars(select(EdgeWorkerModel).order_by(EdgeWorkerModel.worker_name)).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(session)
+ # Format and print worker info on the screen
+ fields = [
+ "worker_name",
+ "state",
+ "queues",
+ "jobs_active",
+ "jobs_success",
+ "jobs_failed",
+ "jobs_taken",
+ "maintenance_comment",
+ ]
+ all_hosts = [{f: str(host.__getattribute__(f)) for f in fields} for host
in all_hosts_iter]
+ # Filter hosts by state
+ if args.state:
+ all_hosts = [host for host in all_hosts if host["state"] in args.state]
+ AirflowConsole().print_as(data=all_hosts, output=args.output)
+
+
+@providers_configuration_loaded
+@provide_session
+def _valid_host(hostname, session: Session = NEW_SESSION) -> bool:
+ """Query the database to check if provided hostname exists."""
+ all_hosts_iter = _fetch_edge_hosts_from_db(session)
+ for host in all_hosts_iter:
+ if hostname == host.worker_name:
+ return True
+ return False
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def put_remote_worker_on_maintainance(args) -> None:
+ """Put remote edge worker on maintenance."""
+ if not _valid_host(args.edge_hostname):
+ raise SystemExit(f"Error: Edge Worker {args.edge_hostname} is
unknown!")
+ from airflow.providers.edge3.models.edge_worker import request_maintenance
+
+ maintenance_comment = f"[{datetime.now().strftime('%Y-%m-%d %H:%M')}] -
{getuser()} put node into maintenance mode\nComment: {args.comments}"
+ request_maintenance(args.edge_hostname, args.comments)
+ logger.info(f"{args.edge_hostname} has been put on maintenance by
{getuser()}!")
Review Comment:
I'd propose w/o exclamation mark. It is no error or warning.
```suggestion
logger.info(f"{args.edge_hostname} has been put on maintenance by
{getuser()}")
```
(and same for the other info logs)
--
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]
