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 :
1514 from 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 ):
2633 def __init__ (self , conn_params = None ):
2734 if conn_params is None :
@@ -33,72 +40,80 @@ def __init__(self, conn_params=None):
3340 self .remote = False
3441 self .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
6683 process = subprocess .Popen (
6784 cmd ,
6885 shell = 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 )
7290 if get_process :
73- return process
74-
91+ return process , None , None
7592 try :
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
7798 except subprocess .TimeoutExpired :
7899 process .kill ()
79100 raise 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
104119 def environ (self , var_name ):
@@ -210,7 +225,7 @@ def read(self, filename, encoding=None, binary=False):
210225 if binary :
211226 return content
212227 if isinstance (content , bytes ):
213- return content .decode (encoding or 'utf-8' )
228+ return content .decode (encoding or get_default_encoding () )
214229 return content
215230
216231 def readlines (self , filename , num_lines = 0 , binary = False , encoding = None ):
0 commit comments