Fix initdb error on Windows

Author: vshepard

testgres/operations/local_ops.py

@@ -8,8 +8,7 @@
 import psutil
 
 from ..exceptions import ExecUtilException
-from .os_ops import ConnectionParams, OsOperations
-from .os_ops import pglib
+from .os_ops import ConnectionParams, OsOperations, pglib, get_default_encoding
 
 try:
     from shutil import which as find_executable
@@ -22,6 +21,12 @@
 error_markers = [b'error', b'Permission denied', b'fatal']
 
 
+def has_errors(output):
+    if isinstance(output, str):
+        output = output.encode(get_default_encoding())
+    return any(marker in output for marker in error_markers)
+
+
 class LocalOperations(OsOperations):
     def __init__(self, conn_params=None):
         if conn_params is None:
@@ -33,7 +38,38 @@ def __init__(self, conn_params=None):
         self.remote = False
         self.username = conn_params.username or self.get_user()
 
-    # Command execution
+    @staticmethod
+    def _run_command(cmd, shell, input, timeout, encoding, temp_file=None):
+        """Execute a command and return the process."""
+        if temp_file is not None:
+            stdout = temp_file
+            stderr = subprocess.STDOUT
+        else:
+            stdout = subprocess.PIPE
+            stderr = subprocess.PIPE
+
+        process = subprocess.Popen(
+            cmd,
+            shell=shell,
+            stdin=subprocess.PIPE if input is not None else None,
+            stdout=stdout,
+            stderr=stderr,
+        )
+
+        try:
+            return process.communicate(input=input.encode(encoding) if input else None, timeout=timeout), process
+        except subprocess.TimeoutExpired:
+            process.kill()
+            raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
+
+    @staticmethod
+    def _raise_exec_exception(message, command, exit_code, output):
+        """Raise an ExecUtilException."""
+        raise ExecUtilException(message=message.format(output),
+                                command=command,
+                                exit_code=exit_code,
+                                out=output)
+
     def exec_command(self, cmd, wait_exit=False, verbose=False,
                      expect_error=False, encoding=None, shell=False, text=False,
                      input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
@@ -56,16 +92,15 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
         :return: The output of the subprocess.
         """
         if os.name == 'nt':
-            with tempfile.NamedTemporaryFile() as buf:
-                process = subprocess.Popen(cmd, stdout=buf, stderr=subprocess.STDOUT)
-                process.communicate()
-                buf.seek(0)
-                result = buf.read().decode(encoding)
-            return result
+            self._exec_command_windows(cmd, wait_exit=wait_exit, verbose=verbose,
+                                       expect_error=expect_error, encoding=encoding, shell=shell, text=text,
+                                       input=input, stdin=stdin, stdout=stdout, stderr=stderr,
+                                       get_process=get_process, timeout=timeout)
         else:
             process = subprocess.Popen(
                 cmd,
                 shell=shell,
+                stdin=stdin,
                 stdout=stdout,
                 stderr=stderr,
             )
@@ -79,7 +114,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
                 raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
             exit_status = process.returncode
 
-            error_found = exit_status != 0 or any(marker in error for marker in error_markers)
+            error_found = exit_status != 0 or has_errors(error)
 
             if encoding:
                 result = result.decode(encoding)
@@ -91,15 +126,49 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
             if exit_status != 0 or error_found:
                 if exit_status == 0:
                     exit_status = 1
-                raise ExecUtilException(message='Utility exited with non-zero code. Error `{}`'.format(error),
-                                        command=cmd,
-                                        exit_code=exit_status,
-                                        out=result)
+                self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, exit_status, result)
             if verbose:
                 return exit_status, result, error
             else:
                 return result
 
+    @staticmethod
+    def _process_output(process, encoding, temp_file=None):
+        """Process the output of a command."""
+        if temp_file is not None:
+            temp_file.seek(0)
+            output = temp_file.read()
+        else:
+            output = process.stdout.read()
+
+        if encoding:
+            output = output.decode(encoding)
+
+        return output
+    def _exec_command_windows(self, cmd, wait_exit=False, verbose=False,
+                              expect_error=False, encoding=None, shell=False, text=False,
+                              input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                              get_process=None, timeout=None):
+        with tempfile.NamedTemporaryFile(mode='w+b') as temp_file:
+            _, process = self._run_command(cmd, shell, input, timeout, encoding, temp_file)
+            if get_process:
+                return process
+            output = self._process_output(process, encoding, temp_file)
+
+            if process.returncode != 0 or has_errors(output):
+                if process.returncode == 0:
+                    process.returncode = 1
+                if expect_error:
+                    if verbose:
+                        return process.returncode, output, output
+                    else:
+                        return output
+                else:
+                    self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, process.returncode,
+                                               output)
+
+            return (process.returncode, output, output) if verbose else output
+
     # Environment setup
     def environ(self, var_name):
         return os.environ.get(var_name)
@@ -210,7 +279,7 @@ def read(self, filename, encoding=None, binary=False):
             if binary:
                 return content
             if isinstance(content, bytes):
-                return content.decode(encoding or 'utf-8')
+                return content.decode(encoding or get_default_encoding())
             return content
 
     def readlines(self, filename, num_lines=0, binary=False, encoding=None):

testgres/operations/os_ops.py

@@ -1,3 +1,5 @@
+import locale
+
 try:
     import psycopg2 as pglib  # noqa: F401
 except ImportError:
@@ -14,6 +16,10 @@ def __init__(self, host='127.0.0.1', ssh_key=None, username=None):
         self.username = username
 
 
+def get_default_encoding():
+    return locale.getdefaultlocale()[1] or 'UTF-8'
+
+
 class OsOperations:
     def __init__(self, username=None):
         self.ssh_key = None

testgres/operations/remote_ops.py

@@ -1,4 +1,3 @@
-import locale
 import logging
 import os
 import subprocess
@@ -15,12 +14,7 @@
         raise ImportError("You must have psycopg2 or pg8000 modules installed")
 
 from ..exceptions import ExecUtilException
-
-from .os_ops import OsOperations, ConnectionParams
-
-ConsoleEncoding = locale.getdefaultlocale()[1]
-if not ConsoleEncoding:
-    ConsoleEncoding = 'UTF-8'
+from .os_ops import OsOperations, ConnectionParams, get_default_encoding
 
 error_markers = [b'error', b'Permission denied', b'fatal', b'No such file or directory']
 
@@ -283,7 +277,9 @@ def copytree(self, src, dst):
         return self.exec_command("cp -r {} {}".format(src, dst))
 
     # Work with files
-    def write(self, filename, data, truncate=False, binary=False, read_and_write=False, encoding=ConsoleEncoding):
+    def write(self, filename, data, truncate=False, binary=False, read_and_write=False, encoding=None):
+        if not encoding:
+            encoding = get_default_encoding()
         mode = "wb" if binary else "w"
         if not truncate:
             mode = "ab" if binary else "a"

testgres/utils.py

@@ -3,14 +3,18 @@
 from __future__ import division
 from __future__ import print_function
 
+import locale
 import os
-import port_for
+import random
+import socket
+
 import sys
 
 from contextlib import contextmanager
 from packaging.version import Version, InvalidVersion
 import re
 
+from port_for import PortForException
 from six import iteritems
 
 from .exceptions import ExecUtilException
@@ -37,13 +41,49 @@ def reserve_port():
     """
     Generate a new port and add it to 'bound_ports'.
     """
-
-    port = port_for.select_random(exclude_ports=bound_ports)
+    port = select_random(exclude_ports=bound_ports)
     bound_ports.add(port)
 
     return port
 
 
+def select_random(
+        ports=None,
+        exclude_ports=None,
+) -> int:
+    """
+    Return random unused port number.
+    Standard function from port_for does not work on Windows 
+        - an error 'port_for.exceptions.PortForException: Can't select a port'
+    We should update it.
+    """
+    if ports is None:
+        ports = set(range(1024, 65535))
+
+    if exclude_ports is None:
+        exclude_ports = set()
+
+    ports.difference_update(set(exclude_ports))
+
+    sampled_ports = random.sample(tuple(ports), min(len(ports), 100))
+
+    for port in sampled_ports:
+        if is_port_free(port):
+            return port
+
+    raise PortForException("Can't select a port")
+
+
+def is_port_free(port: int) -> bool:
+    """Check if a port is free to use."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        try:
+            s.bind(("", port))
+            return True
+        except OSError:
+            return False
+
+
 def release_port(port):
     """
     Free port provided by reserve_port().
@@ -80,7 +120,8 @@ def execute_utility(args, logfile=None, verbose=False):
                 lines = [u'\n'] + ['# ' + line for line in out.splitlines()] + [u'\n']
                 tconf.os_ops.write(filename=logfile, data=lines)
         except IOError:
-            raise ExecUtilException("Problem with writing to logfile `{}` during run command `{}`".format(logfile, args))
+            raise ExecUtilException(
+                "Problem with writing to logfile `{}` during run command `{}`".format(logfile, args))
     if verbose:
         return exit_status, out, error
     else:

tests/test_simple.py

@@ -74,6 +74,24 @@ def good_properties(f):
             return True
 
 
+def rm_carriage_returns(out):
+    """
+    In Windows we have additional '\r' symbols in output.
+    Let's get rid of them.
+    """
+    if os.name == 'nt':
+        if isinstance(out, (int, float, complex)):
+            return out
+        elif isinstance(out, tuple):
+            return tuple(rm_carriage_returns(item) for item in out)
+        elif isinstance(out, bytes):
+            return out.replace(b'\r', b'')
+        else:
+            return out.replace('\r', '')
+    else:
+        return out
+
+
 @contextmanager
 def removing(f):
     try:
@@ -254,34 +272,34 @@ def test_psql(self):
 
             # check returned values (1 arg)
             res = node.psql('select 1')
-            self.assertEqual(res, (0, b'1\n', b''))
+            self.assertEqual(rm_carriage_returns(res), (0, b'1\n', b''))
 
             # check returned values (2 args)
             res = node.psql('postgres', 'select 2')
-            self.assertEqual(res, (0, b'2\n', b''))
+            self.assertEqual(rm_carriage_returns(res), (0, b'2\n', b''))
 
             # check returned values (named)
             res = node.psql(query='select 3', dbname='postgres')
-            self.assertEqual(res, (0, b'3\n', b''))
+            self.assertEqual(rm_carriage_returns(res), (0, b'3\n', b''))
 
             # check returned values (1 arg)
             res = node.safe_psql('select 4')
-            self.assertEqual(res, b'4\n')
+            self.assertEqual(rm_carriage_returns(res), b'4\n')
 
             # check returned values (2 args)
             res = node.safe_psql('postgres', 'select 5')
-            self.assertEqual(res, b'5\n')
+            self.assertEqual(rm_carriage_returns(res), b'5\n')
 
             # check returned values (named)
             res = node.safe_psql(query='select 6', dbname='postgres')
-            self.assertEqual(res, b'6\n')
+            self.assertEqual(rm_carriage_returns(res), b'6\n')
 
             # check feeding input
             node.safe_psql('create table horns (w int)')
             node.safe_psql('copy horns from stdin (format csv)',
                            input=b"1\n2\n3\n\\.\n")
             _sum = node.safe_psql('select sum(w) from horns')
-            self.assertEqual(_sum, b'6\n')
+            self.assertEqual(rm_carriage_returns(_sum), b'6\n')
 
             # check psql's default args, fails
             with self.assertRaises(QueryException):
@@ -455,7 +473,7 @@ def test_synchronous_replication(self):
                     master.safe_psql(
                         'insert into abc select generate_series(1, 1000000)')
                     res = standby1.safe_psql('select count(*) from abc')
-                    self.assertEqual(res, b'1000000\n')
+                    self.assertEqual(rm_carriage_returns(res), b'1000000\n')
 
     @unittest.skipUnless(pg_version_ge('10'), 'requires 10+')
     def test_logical_replication(self):
@@ -589,7 +607,7 @@ def test_promotion(self):
                 # make standby becomes writable master
                 replica.safe_psql('insert into abc values (1)')
                 res = replica.safe_psql('select * from abc')
-                self.assertEqual(res, b'1\n')
+                self.assertEqual(rm_carriage_returns(res), b'1\n')
 
     def test_dump(self):
         query_create = 'create table test as select generate_series(1, 2) as val'
@@ -614,6 +632,7 @@ def test_users(self):
         with get_new_node().init().start() as node:
             node.psql('create role test_user login')
             value = node.safe_psql('select 1', username='test_user')
+            value = rm_carriage_returns(value)
             self.assertEqual(value, b'1\n')
 
     def test_poll_query_until(self):
@@ -977,7 +996,9 @@ def test_child_pids(self):
 
     def test_child_process_dies(self):
         # test for FileNotFound exception during child_processes() function
-        with subprocess.Popen(["sleep", "60"]) as process:
+        cmd = ["timeout", "60"] if os.name == 'nt' else ["sleep", "60"]
+
+        with subprocess.Popen(cmd, shell=True) as process:  # shell=True might be needed on Windows
             self.assertEqual(process.poll(), None)
             # collect list of processes currently running
             children = psutil.Process(os.getpid()).children()