airflow icon indicating copy to clipboard operation
airflow copied to clipboard

Add on kill to ssh

Open MRLab12 opened this issue 1 year ago • 6 comments

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


MRLab12 avatar Jun 21 '24 20:06 MRLab12

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

boring-cyborg[bot] avatar Jun 21 '24 20:06 boring-cyborg[bot]

@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.

MRLab12 avatar Jun 24 '24 17:06 MRLab12

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.

potiuk avatar Jun 24 '24 17:06 potiuk

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.

MRLab12 avatar Jun 24 '24 18:06 MRLab12

Might be a race condition, yes.

potiuk avatar Jun 24 '24 19:06 potiuk

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.

MRLab12 avatar Jun 26 '24 20:06 MRLab12

@potiuk Leaving a reminder here that this PR is ready for review.

MRLab12 avatar Jul 03 '24 12:07 MRLab12

@potiuk Leaving a reminder here that this PR is ready for review.

Leaving a reminder that tests are failing.

potiuk avatar Jul 03 '24 12:07 potiuk

@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.

MRLab12 avatar Jul 03 '24 14:07 MRLab12

Nope. Seems they are still failing

potiuk avatar Jul 04 '24 13:07 potiuk

Awesome work, congrats on your first merged pull request! You are invited to check our Issue Tracker for additional contributions.

boring-cyborg[bot] avatar Jul 24 '24 18:07 boring-cyborg[bot]