Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SSH task exit code added to XCOM as 'ssh_exit' key #27370

Merged
merged 10 commits into from
Dec 4, 2022
Merged

SSH task exit code added to XCOM as 'ssh_exit' key #27370

Changes from 5 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
6dbc2ef
SSH task exit code added to XCOM as 'ssh_exit' key
kolfild26 Oct 29, 2022
2258629
Merge branch 'main' into ssh_exit_code
kolfild26 Oct 30, 2022
f4a597c
raise_for_status double-quotes added
kolfild26 Oct 31, 2022
59de9e7
self.raise_for_status double-quotes added
kolfild26 Oct 31, 2022
6ab73a2
Merge branch 'main' into ssh_exit_code
kolfild26 Oct 31, 2022
1ae2d81
Merge branch 'main' into ssh_exit_code
kolfild26 Nov 14, 2022
bf6b6e3
Merge branch 'main' into ssh_exit_code
kolfild26 Nov 19, 2022
752a918
if context and self.do_xcom_push then push ssh_exit
kolfild26 Nov 19, 2022
6627c9b
test_push_ssh_exit_to_xcom added
kolfild26 Nov 20, 2022
ea5afa9
Merge branch 'main' into ssh_exit_code
kolfild26 Nov 26, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions airflow/providers/ssh/operators/ssh.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,16 +151,18 @@ def exec_ssh_client_command(self, ssh_client: SSHClient, command: str):
ssh_client, command, timeout=self.timeout, environment=self.environment, get_pty=self.get_pty
)

def raise_for_status(self, exit_status: int, stderr: bytes) -> None:
def raise_for_status(self, exit_status: int, stderr: bytes, **kwargs) -> None:
ti=kwargs["context"].get("task_instance")
Copy link
Contributor

@bdsoha bdsoha Nov 15, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
ti=kwargs["context"].get("task_instance")
ti = kwargs["context"].get("task_instance")

ti.xcom_push(key="ssh_exit", value=exit_status)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should verify user wish to push to xcom.
users ask for this specifically with do_xcom_push parameter

Copy link
Contributor

@bdsoha bdsoha Nov 1, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@eladkal It is not optional on other operators:

airflow/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py

Lines 444 to 446 in b29ca4e

ti = context["ti"]
ti.xcom_push(key="pod_name", value=self.pod.metadata.name)
ti.xcom_push(key="pod_namespace", value=self.pod.metadata.namespace)

if exit_status != 0:
raise AirflowException(f"SSH operator error: exit status = {exit_status}")

def run_ssh_client_command(self, ssh_client: SSHClient, command: str) -> bytes:
def run_ssh_client_command(self, ssh_client: SSHClient, command: str, **kwargs) -> bytes:
assert self.ssh_hook
exit_status, agg_stdout, agg_stderr = self.ssh_hook.exec_ssh_client_command(
ssh_client, command, timeout=self.timeout, environment=self.environment, get_pty=self.get_pty
)
self.raise_for_status(exit_status, agg_stderr)
self.raise_for_status(exit_status, agg_stderr, context=kwargs["context"])
return agg_stdout

def execute(self, context=None) -> bytes | str:
Expand All @@ -172,7 +174,7 @@ def execute(self, context=None) -> bytes | str:
self.get_pty = self.command.startswith("sudo") or self.get_pty

with self.get_ssh_client() as ssh_client:
result = self.run_ssh_client_command(ssh_client, self.command)
result = self.run_ssh_client_command(ssh_client, self.command, context=context)
enable_pickling = conf.getboolean("core", "enable_xcom_pickling")
if not enable_pickling:
result = b64encode(result).decode("utf-8")
Expand Down