airflow
airflow copied to clipboard
Add on kill to ssh
This is my first contribution here, so still learning the codebase. The SSHOperator does not have on_kill method that can be used to close a connection when needed. This PR adds a on_kill method to the SSHOperator class.
Not sure how extensive the changes should be. The new method should most likely be used in different parts of the codebase that the ssh hook is used.
closes: #40343
Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide (https://github.com/apache/airflow/blob/main/contributing-docs/README.rst) Here are some useful points:
- Pay attention to the quality of your code (ruff, mypy and type annotations). Our pre-commits will help you with that.
- In case of a new feature add useful documentation (in docstrings or in
docs/directory). Adding a new operator? Check this short guide Consider adding an example DAG that shows how users should use it. - Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
- Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
- Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
- Be sure to read the Airflow Coding style.
- Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits. Apache Airflow is a community-driven project and together we are making it better 🚀. In case of doubts contact the developers at: Mailing List: [email protected] Slack: https://s.apache.org/airflow-slack
@potiuk @uranusjr Sorry, I got misguided there with the cached_property. I found the if/else use in the sftp hook file:
def get_conn(self) -> paramiko.SFTPClient: # type: ignore[override]
"""Open an SFTP connection to the remote host."""
if self.conn is None:
# TODO: remove support for ssh_hook when it is removed from SFTPOperator
if self.ssh_hook is not None:
self.conn = self.ssh_hook.get_conn().open_sftp()
else:
self.conn = super().get_conn().open_sftp()
return self.conn
I can do the same in the ssh hook and return the existing self.client. @potiuk Let me know if this is what you're referring to.
I can do the same in the ssh hook and return the existing self.client. @potiuk Let me know if this is what you're referring to.
Yes.
I can do the same in the ssh hook and return the existing self.client. @potiuk Let me know if this is what you're referring to.
Yes.
Made the change. Wondering if there could be a case where the on_kill method is used, but there wasn't any active connection. A new connection would be created. Let me know your thoughts.
Might be a race condition, yes.
Might be a race condition, yes.
Looks like to close the connection we need to do it in execute from SSHOperator. I made the change. This should close all shh connections opened to run a command.
@potiuk Leaving a reminder here that this PR is ready for review.
@potiuk Leaving a reminder here that this PR is ready for review.
Leaving a reminder that tests are failing.
@potiuk Leaving a reminder here that this PR is ready for review.
Leaving a reminder that tests are failing.
Apologies, thought I had seen tests passing before I merged main. Had to do some formatting fixes. Looks like tests are passing now.
Nope. Seems they are still failing
Awesome work, congrats on your first merged pull request! You are invited to check our Issue Tracker for additional contributions.