You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "dheerajturaga (via GitHub)" <gi...@apache.org> on 2025/04/28 19:52:36 UTC
[PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
dheerajturaga opened a new pull request, #49915:
URL: https://github.com/apache/airflow/pull/49915
EdgeExecutor is missing EdgeWorkerHosts html pages in Airflow 3. This PR attempts to bridge the gap by providing those utilities via CLI options. These options are intended to run on the server side ( or have access to the server and db). They aim to control the remote edge workers.
This PR enables the following functionality:
1. List all registered edge workers
2. Put an edge worker on maintenance mode
3. Update maintenance mode comments for a worker who is in maintenance mode
4. Remove edge worker from maintenance mode
5. Remove edge worker from the pool
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2072016454
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,93 @@ 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: 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
+
+ try:
+ change_maintenance_comment(args.edge_hostname, args.comments)
+ logger.info("Maintenance comments updated for %s by %s.", args.edge_hostname, getuser())
+ except TypeError:
+ raise SystemExit
+
+
+@cli_utils.action_cli
Review Comment:
Done!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl merged PR #49915:
URL: https://github.com/apache/airflow/pull/49915
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2065352483
##########
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:
This is done
##########
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:
This is done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067725879
##########
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)
Review Comment:
I made a note of this. I can update the status command with this in my next PR :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067726747
##########
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())
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def remove_remote_worker(args) -> None:
+ """Remove remote edge worker entry from db."""
+ 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 remove_worker
+
+ remove_worker(args.edge_hostname)
Review Comment:
I have addressed this, let me know if you like what you see. Else we can massage it further
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2069480150
##########
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:
Funny co-incidence. It was an infra-glitch and just at the exact test. But re-running made it green. If it would have been a dependency issue should have failed mysql and sqlite tests as well. So all is good!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067205815
##########
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:
Mhm, okay that is a bit strange because (1) this would onyl work for psotgres and when mysql/sqlite is used (which probably not many use... but is supported though) then it would be the same problem. But in reality the CLI is always installed with airflow-core which needs to bring DB drivers for all supported backends. So a dedicated definitions of drivers should not be needed.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067321728
##########
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())
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def remove_remote_worker(args) -> None:
+ """Remove remote edge worker entry from db."""
+ 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 remove_worker
+
+ remove_worker(args.edge_hostname)
Review Comment:
The EdgeModel should prevent removing the worker if it is not offline / maintenance offline.
Reason: When the worker is active and the record is deleted, then the (online) worker fails reporting heartbeat as the DB entry is missing.
##########
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:
Would be great - else you can easily generate inconsistencies.
##########
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",
Review Comment:
Queues are stored as string in the DB, to make it proper the strin must be re-interpreted, else it is printing a string of the queue list:
Maybe this can be generalized because in the UI the same is made. So would be best to add this to EdgeModel class and consume it in CLI only?
##########
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:
I would remove them, I checked also during istall the both packages are not re-installed - as they are installed already when airflow-core is installed.
##########
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:
Maintenance comment can also be set if not in maintenance. This might also need to be checked with the status model.
##########
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)
Review Comment:
Outsideof this PR... maybe this "standard" printing in any format needs to be used in the `airflow edge status` reporting as well, did not know about this before :-D
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067727113
##########
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:
These are removed!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067726747
##########
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())
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def remove_remote_worker(args) -> None:
+ """Remove remote edge worker entry from db."""
+ 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 remove_worker
+
+ remove_worker(args.edge_hostname)
Review Comment:
I have addressed this, would like to know what you think about this
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067838984
##########
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:
@jscheffl , postgres related ci steps are failing once I removed this. You think this is related?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2065360590
##########
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 see what you are saying. Currently I updated to query the db directly. Let me see if I can update the edge worker model
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067683390
##########
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())
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def remove_remote_worker(args) -> None:
+ """Remove remote edge worker entry from db."""
+ 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 remove_worker
+
+ remove_worker(args.edge_hostname)
Review Comment:
Sure, will do this. Im looking for a "shut down remote edge worker" aswell.
Im trying to reduce the need to "ssh" into these hosts to do the dirty work
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067726268
##########
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",
Review Comment:
I was setting the value to a string which wasnt needed. this is addressed
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067724971
##########
providers/edge3/tests/unit/edge3/cli/test_edge_command.py:
##########
@@ -395,3 +430,23 @@ def test_get_sysinfo(self, worker_with_job: _EdgeWorkerCli):
assert "edge_provider_version" in sysinfo
assert "concurrency" in sysinfo
assert sysinfo["concurrency"] == concurrency
+
+ def test_list_edge_workers(self, mock_edgeworker: EdgeWorkerModel):
Review Comment:
this is addressed
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2072015688
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,93 @@ def stop(args):
logger.info("Worker has been shut down.")
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(
Review Comment:
I moved all of the db related logic to the edge worker model as per your recommendation
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067371580
##########
providers/edge3/tests/unit/edge3/cli/test_edge_command.py:
##########
@@ -395,3 +430,23 @@ def test_get_sysinfo(self, worker_with_job: _EdgeWorkerCli):
assert "edge_provider_version" in sysinfo
assert "concurrency" in sysinfo
assert sysinfo["concurrency"] == concurrency
+
+ def test_list_edge_workers(self, mock_edgeworker: EdgeWorkerModel):
Review Comment:
To fix this test you need to add/mark it via:
`pytestmark = pytest.mark.db_test`
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2065279099
##########
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:
I was trying to run my cli commands `uv run airflow edge list-workers` when I hit the issue. My database_sql connection was a postgress+psycopg2 connection and I needed these to be included to get going
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2069486697
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,93 @@ def stop(args):
logger.info("Worker has been shut down.")
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(
Review Comment:
After a night of sleep I _think_ the playcement is okay-ish for the DB activity. Just thinking (but not forcing) it might be a bit cleaner if all DB related stuff is moced intoedge worker model and CLI just calls to retrieve - then CLI would need to DB imports (and the risk of missing them on the client). What do you think?
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,93 @@ 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: 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
+
+ try:
+ change_maintenance_comment(args.edge_hostname, args.comments)
+ logger.info("Maintenance comments updated for %s by %s.", args.edge_hostname, getuser())
+ except TypeError:
+ raise SystemExit
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def remove_remote_worker(args) -> None:
Review Comment:
Forgot to test this - don't know if you considered this? What happens if you call the commend on the edge where (most probably) DB is not configured? Will it throw a crazy stack trace? Should we handle this? (not that it attempts to setup a local sqlite in the file system)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "dheerajturaga (via GitHub)" <gi...@apache.org>.
dheerajturaga commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2072016218
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,93 @@ 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: 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
+
+ try:
+ change_maintenance_comment(args.edge_hostname, args.comments)
+ logger.info("Maintenance comments updated for %s by %s.", args.edge_hostname, getuser())
+ except TypeError:
+ raise SystemExit
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def remove_remote_worker(args) -> None:
Review Comment:
I have added a check_valid_db_connection method that should catch this
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2070249485
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,93 @@ 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: 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
+
+ try:
+ change_maintenance_comment(args.edge_hostname, args.comments)
+ logger.info("Maintenance comments updated for %s by %s.", args.edge_hostname, getuser())
+ except TypeError:
+ raise SystemExit
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def remove_remote_worker(args) -> None:
Review Comment:
Okay, just checked this. If you attempt to call this on the edge-side w/o db connection configured you generate a lot of errors:
```
root@cd4307c6ba06:/opt/airflow# airflow edge list-workers
Please confirm database initialize (or wait 4 seconds to skip it). Are you sure? [y/N]
[2025-05-01T12:48:53.345+0000] {db.py:924} INFO - Log template table does not exist (added in 2.3.0); skipping log template sync.
Traceback (most recent call last):
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1910, in _execute_context
self.dialect.do_execute(
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line 736, in do_execute
cursor.execute(statement, parameters)
sqlite3.OperationalError: no such table: edge_worker
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/usr/local/bin/airflow", line 10, in <module>
sys.exit(main())
^^^^^^
File "/opt/airflow/airflow-core/src/airflow/__main__.py", line 55, in main
args.func(args)
File "/opt/airflow/airflow-core/src/airflow/utils/cli.py", line 111, in wrapper
return f(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/opt/airflow/airflow-core/src/airflow/utils/providers_configuration_loader.py", line 55, in wrapped_function
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/airflow/airflow-core/src/airflow/utils/session.py", line 101, in wrapper
return func(*args, session=session, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/airflow/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py", line 649, in list_edge_workers
all_hosts_iter = _fetch_edge_hosts_from_db(states=args.state, session=session)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/airflow/airflow-core/src/airflow/utils/providers_configuration_loader.py", line 55, in wrapped_function
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/airflow/airflow-core/src/airflow/utils/session.py", line 99, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/airflow/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py", line 641, in _fetch_edge_hosts_from_db
return session.scalars(query).all()
^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/orm/session.py", line 1778, in scalars
return self.execute(
^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/orm/session.py", line 1717, in execute
result = conn._execute_20(statement, params or {}, execution_options)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1710, in _execute_20
return meth(self, args_10style, kwargs_10style, execution_options)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/sql/elements.py", line 334, in _execute_on_connection
return connection._execute_clauseelement(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1577, in _execute_clauseelement
ret = self._execute_context(
^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1953, in _execute_context
self._handle_dbapi_exception(
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 2134, in _handle_dbapi_exception
util.raise_(
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/util/compat.py", line 211, in raise_
raise exception
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1910, in _execute_context
self.dialect.do_execute(
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line 736, in do_execute
cursor.execute(statement, parameters)
sqlalchemy.exc.OperationalError: (sqlite3.OperationalError) no such table: edge_worker
[SQL: SELECT edge_worker.queues, edge_worker.worker_name, edge_worker.state, edge_worker.maintenance_comment, edge_worker.first_online, edge_worker.last_update, edge_worker.jobs_active, edge_worker.jobs_taken, edge_worker.jobs_success, edge_worker.jobs_failed, edge_worker.sysinfo
FROM edge_worker ORDER BY edge_worker.worker_name]
(Background on this error at: https://sqlalche.me/e/14/e3q8)
```
Can you please add a small validation preventing sqlite DB creation and the error?
Proposal as code to check whether a valid DB is configured, maybe put this into a utility?:
```
db_conn = conf.get("database", "sql_alchemy_conn")
db_default = conf.get_default_value("database", "sql_alchemy_conn")
if db_conn == db_default:
raise SystemExit(
"Error: The database connection is not set. Please set the connection in the configuration file."
)
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Extend Edge Worker CLI commands operate on remote edge workers [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2070250173
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,93 @@ 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: 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
+
+ try:
+ change_maintenance_comment(args.edge_hostname, args.comments)
+ logger.info("Maintenance comments updated for %s by %s.", args.edge_hostname, getuser())
+ except TypeError:
+ raise SystemExit
+
+
+@cli_utils.action_cli
Review Comment:
Alongside to prevent that DB is attempted to be initialized, can you change this?
```suggestion
@cli_utils.action_cli(check_db=False)
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org