88import psutil
99
1010from ..exceptions import ExecUtilException
11- from .os_ops import ConnectionParams ,OsOperations
12- from .os_ops import pglib
11+ from .os_ops import ConnectionParams ,OsOperations ,pglib ,get_default_encoding
1312
1413try :
1514from shutil import which as find_executable
2221error_markers = [b'error' ,b'Permission denied' ,b'fatal' ]
2322
2423
24+ def has_errors (output ):
25+ if output :
26+ if isinstance (output ,str ):
27+ output = output .encode (get_default_encoding ())
28+ return any (marker in output for marker in error_markers )
29+ return False
30+
31+
2532class LocalOperations (OsOperations ):
2633def __init__ (self ,conn_params = None ):
2734if conn_params is None :
@@ -33,72 +40,80 @@ def __init__(self, conn_params=None):
3340self .remote = False
3441self .username = conn_params .username or self .get_user ()
3542
36- # Command execution
37- def exec_command (self ,cmd ,wait_exit = False ,verbose = False ,
38- expect_error = False ,encoding = None ,shell = False ,text = False ,
39- input = None ,stdin = subprocess .PIPE ,stdout = subprocess .PIPE ,stderr = subprocess .PIPE ,
40- get_process = None ,timeout = None ):
41- """
42- Execute a command in a subprocess.
43-
44- Args:
45- - cmd: The command to execute.
46- - wait_exit: Whether to wait for the subprocess to exit before returning.
47- - verbose: Whether to return verbose output.
48- - expect_error: Whether to raise an error if the subprocess exits with an error status.
49- - encoding: The encoding to use for decoding the subprocess output.
50- - shell: Whether to use shell when executing the subprocess.
51- - text: Whether to return str instead of bytes for the subprocess output.
52- - input: The input to pass to the subprocess.
53- - stdout: The stdout to use for the subprocess.
54- - stderr: The stderr to use for the subprocess.
55- - proc: The process to use for subprocess creation.
56- :return: The output of the subprocess.
57- """
58- if os .name == 'nt' :
59- with tempfile .NamedTemporaryFile ()as buf :
60- process = subprocess .Popen (cmd ,stdout = buf ,stderr = subprocess .STDOUT )
61- process .communicate ()
62- buf .seek (0 )
63- result = buf .read ().decode (encoding )
64- return result
65- else :
43+ @staticmethod
44+ def _raise_exec_exception (message ,command ,exit_code ,output ):
45+ """Raise an ExecUtilException."""
46+ raise ExecUtilException (message = message .format (output ),
47+ command = command ,
48+ exit_code = exit_code ,
49+ out = output )
50+
51+ @staticmethod
52+ def _process_output (encoding ,temp_file_path ):
53+ """Process the output of a command from a temporary file."""
54+ with open (temp_file_path ,'rb' )as temp_file :
55+ output = temp_file .read ()
56+ if encoding :
57+ output = output .decode (encoding )
58+ return output ,None # In Windows stderr writing in stdout
59+
60+ def _run_command (self ,cmd ,shell ,input ,stdin ,stdout ,stderr ,get_process ,timeout ,encoding ):
61+ """Execute a command and return the process and its output."""
62+ if os .name == 'nt' and stdout is None :# Windows
63+ with tempfile .NamedTemporaryFile (mode = 'w+b' ,delete = False )as temp_file :
64+ stdout = temp_file
65+ stderr = subprocess .STDOUT
66+ process = subprocess .Popen (
67+ cmd ,
68+ shell = shell ,
69+ stdin = stdin or subprocess .PIPE if input is not None else None ,
70+ stdout = stdout ,
71+ stderr = stderr ,
72+ )
73+ if get_process :
74+ return process ,None ,None
75+ temp_file_path = temp_file .name
76+
77+ # Wait process finished
78+ process .wait ()
79+
80+ output ,error = self ._process_output (encoding ,temp_file_path )
81+ return process ,output ,error
82+ else :# Other OS
6683process = subprocess .Popen (
6784cmd ,
6885shell = shell ,
69- stdout = stdout ,
70- stderr = stderr ,
86+ stdin = stdin or subprocess .PIPE if input is not None else None ,
87+ stdout = stdout or subprocess .PIPE ,
88+ stderr = stderr or subprocess .PIPE ,
7189 )
7290if get_process :
73- return process
74-
91+ return process ,None ,None
7592try :
76- result ,error = process .communicate (input ,timeout = timeout )
93+ output ,error = process .communicate (input = input .encode (encoding )if input else None ,timeout = timeout )
94+ if encoding :
95+ output = output .decode (encoding )
96+ error = error .decode (encoding )
97+ return process ,output ,error
7798except subprocess .TimeoutExpired :
7899process .kill ()
79100raise ExecUtilException ("Command timed out after {} seconds." .format (timeout ))
80- exit_status = process .returncode
81-
82- error_found = exit_status != 0 or any (marker in error for marker in error_markers )
83101
84- if encoding :
85- result = result .decode (encoding )
86- error = error .decode (encoding )
87-
88- if expect_error :
89- raise Exception (result ,error )
90-
91- if exit_status != 0 or error_found :
92- if exit_status == 0 :
93- exit_status = 1
94- raise ExecUtilException (message = 'Utility exited with non-zero code. Error `{}`' .format (error ),
95- command = cmd ,
96- exit_code = exit_status ,
97- out = result )
98- if verbose :
99- return exit_status ,result ,error
100- else :
101- return result
102+ def exec_command (self ,cmd ,wait_exit = False ,verbose = False ,expect_error = False ,encoding = None ,shell = False ,
103+ text = False ,input = None ,stdin = None ,stdout = None ,stderr = None ,get_process = False ,timeout = None ):
104+ """
105+ Execute a command in a subprocess and handle the output based on the provided parameters.
106+ """
107+ process ,output ,error = self ._run_command (cmd ,shell ,input ,stdin ,stdout ,stderr ,get_process ,timeout ,encoding )
108+ if get_process :
109+ return process
110+ if process .returncode != 0 or (has_errors (error )and not expect_error ):
111+ self ._raise_exec_exception ('Utility exited with non-zero code. Error `{}`' ,cmd ,process .returncode ,error )
112+
113+ if verbose :
114+ return process .returncode ,output ,error
115+ else :
116+ return output
102117
103118# Environment setup
104119def environ (self ,var_name ):
@@ -210,7 +225,7 @@ def read(self, filename, encoding=None, binary=False):
210225if binary :
211226return content
212227if isinstance (content ,bytes ):
213- return content .decode (encoding or 'utf-8' )
228+ return content .decode (encoding or get_default_encoding () )
214229return content
215230
216231def readlines (self ,filename ,num_lines = 0 ,binary = False ,encoding = None ):