[AIRFLOW-5936] Allow explicit get_pty in SSHOperator

[22quinn]

Make sure you have checked all steps below.

Jira

• My PR addresses the following Airflow Jira issues and references them in the PR title. For example, "[AIRFLOW-XXX] My Airflow PR"
  • https://issues.apache.org/jira/browse/AIRFLOW-5936
  • In case you are fixing a typo in the documentation you can prepend your commit with [AIRFLOW-XXX], code changes always need a Jira issue.
  • In case you are proposing a fundamental code change, you need to create an Airflow Improvement Proposal (AIP).
  • In case you are adding a dependency, check if the license complies with the ASF 3rd Party License Policy.

Description

• Here are some details about my PR, including screenshots of any UI changes:

Currently when execution_timeout is reached for a SSHOperator task, the ssh connection will be closed but the remote process continues to run. In many scenarios, users might want the process to be killed upon task timeout. Giving users an explicit get_pty option achieves this goal.

Tests

• My PR adds the following unit tests OR does not need testing for this extremely good reason:

Commits

• My commits all reference Jira issues in their subject lines, and I have squashed multiple commits if they address the same issue. In addition, my commits follow the guidelines from "How to write a good git commit message":
  1. Subject is separated from body by a blank line
  2. Subject is limited to 50 characters (not including Jira issue reference)
  3. Subject does not end with a period
  4. Subject uses the imperative mood ("add", not "adding")
  5. Body wraps at 72 characters
  6. Body explains "what" and "why", not "how"

Documentation

• In case of new functionality, my PR adds documentation that describes how to use it.
  • All the public functions and the classes in the PR contain docstrings that explain what it does
  • If you implement backwards incompatible changes, please leave a note in the Updating.md so we can assign it to a appropriate release

[codecov-io]

Codecov Report

Merging #6586 into master will decrease coverage by 0.33%.
The diff coverage is 100%.

@@            Coverage Diff             @@
##           master    #6586      +/-   ##
==========================================
- Coverage   83.83%   83.49%   -0.34%     
==========================================
  Files         647      647              
  Lines       37409    37407       -2     
==========================================
- Hits        31360    31234     -126     
- Misses       6049     6173     +124

Impacted Files	Coverage Δ
airflow/contrib/operators/ssh_operator.py	83.33% <100%> (+0.83%)	⬆️
airflow/kubernetes/volume_mount.py	44.44% <0%> (-55.56%)	⬇️
airflow/kubernetes/volume.py	52.94% <0%> (-47.06%)	⬇️
airflow/kubernetes/pod_launcher.py	45.25% <0%> (-46.72%)	⬇️
airflow/kubernetes/refresh_config.py	50.98% <0%> (-23.53%)	⬇️
...rflow/contrib/operators/kubernetes_pod_operator.py	77.14% <0%> (-21.43%)	⬇️
airflow/configuration.py	89.13% <0%> (-3.63%)	⬇️
airflow/models/taskinstance.py	93.28% <0%> (-0.51%)	⬇️
airflow/utils/dag_processing.py	57.99% <0%> (-0.5%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update e50abb8...661f782. Read the comment docs.

[22quinn]

Just want to make it clearer - by timeout I mean the execution_timeout of the task, not the channel timeout parameter in ssh_client.exec_command(). In many use cases, when the execution_timeout is reached, we want not only to close the ssh connection but also to kill the remote process.

[potiuk]

Thanks for the first-time contribution @zzk-dev !

I think it would also be great if all the get_pty logic - even if it is simple - is also tested in TestSSHOperator. if you move the "if" to the constructor you will be able to have a set of tests where you will test the simple logic in the Test itself - you could pass a command with/without sudo and with different get_pty values and test if self.get_pty is set properly.

[potiuk]

I think it might make sense to move this to constructor. command is set only in the constructor so we can keep all the pty logic there and it could be one-liner
self.get_pty = self.start_command.starstwith('sudo') or self.get_pty

[22quinn]

Thanks a lot @potiuk for the review! I have made the suggested changes and added a set of tests.

[22quinn]

NICE!. Just one nit. I think this could be much nicer if you write it as "@parameterized" test - much less code , you specify an array of input-> expected values.

See for example here: https://github.com/apache/airflow/blob/master/tests/contrib/hooks/test_sftp_hook.py#L205

Good suggestion! But I got a question - how do I parameterize the case where get_pty is not provided and the default is applied? I tried to use None but it does not work. Or should I just remove the two test cases with no get_pty input?

[22quinn]

I decided to remove the two cases where get_pty is not provided because it was like unnecessarily testing the Python functionality of applying default.

I also added two constants - COMMAND, COMMAND_WITH_SUDO because other tests use the same command echo -n airflow.

[potiuk]

Great!. Waiting for Travis. It's ready to merge! Thank you!

[22quinn]

Thanks @potiuk ! It's a great learning experience!

[potiuk]

Thanks @zzk-dev -> welcome as contributor! The code is merged. Happy that you've learned something and looking forward for future contributions :)