snowflake-stats/exec.py

66 lines
1.9 KiB
Python
Raw Normal View History

2022-04-06 14:42:27 +02:00
#!/usr/bin/env python3
from typing import List
import sys
import subprocess
import shlex
class NonZeroExitCode(Exception):
def __init__(self, exit_code, stdout, stderr):
super(NonZeroExitCode, self).__init__(f'exit_code: {exit_code}, stdout: {stdout}, stderr: {stderr}')
def wrap_command_with_ssh(command: List[str], ssh_host: str) -> List[str]:
return ['ssh', ssh_host, ' '.join([shlex.quote(arg) for arg in command])]
def execute(command: List[str], ssh_host: str = None) -> None:
"""
https://docs.python.org/3.10/library/subprocess.html#frequently-used-arguments
Executes the given command using stdin, stdout, stderr of the parent process.
The output is not captured.
:raises subprocess.CalledProcessError: In case of non-zero exit code.
"""
if ssh_host is not None:
command = wrap_command_with_ssh(command, ssh_host)
print(f'[DEBUG] execute: {command}')
subprocess.run(
command,
stdin=sys.stdin,
stdout=sys.stdout,
stderr=sys.stderr,
check=True,
)
def execute_and_capture(command: List[str], ssh_host: str = None, file='stdout') -> str:
"""
Executes the given command and captures it stdout.
:return: Stdout of subprocess.
:raises NonZeroExitCode: In case of non-zero exit code. This exception includes the exit_code, stdout and stderr in it's message.
"""
if ssh_host is not None:
command = wrap_command_with_ssh(command, ssh_host)
print(f'[DEBUG] execute_and_capture: {command}')
completed: subprocess.CompletedProcess = subprocess.run(
command,
capture_output=True,
text=True,
)
if completed.returncode != 0:
raise NonZeroExitCode(completed.returncode, completed.stdout, completed.stderr)
if file == 'stdout':
return completed.stdout
if file == 'stderr':
return completed.stderr
raise ValueError('invalid argument')