diff --git a/patroni/postgresql/__init__.py b/patroni/postgresql/__init__.py index a80b76e3..c88568a7 100644 --- a/patroni/postgresql/__init__.py +++ b/patroni/postgresql/__init__.py @@ -233,6 +233,17 @@ class Postgresql(object): pg_ctl = [self.pgcommand('pg_ctl'), cmd] return subprocess.call(pg_ctl + ['-D', self._data_dir] + list(args), **kwargs) == 0 + def initdb(self, *args: str, **kwargs: Any) -> bool: + """Builds and executes the initdb command. + + :param args: List of arguments to be joined into the initdb command. + :param kwargs: Keyword arguments to pass to ``subprocess.call``. + + :returns: ``True`` if the result of ``subprocess.call`, the exit code, is ``0``. + """ + initdb = [self.pgcommand('initdb')] + list(args) + [self.data_dir] + return subprocess.call(initdb, **kwargs) == 0 + def pg_isready(self): """Runs pg_isready to see if PostgreSQL is accepting connections. diff --git a/patroni/postgresql/bootstrap.py b/patroni/postgresql/bootstrap.py index 42c24c9c..e21e1202 100644 --- a/patroni/postgresql/bootstrap.py +++ b/patroni/postgresql/bootstrap.py @@ -3,10 +3,11 @@ import os import shlex import tempfile import time +from typing import List, Dict, Union, Callable, Tuple from ..dcs import RemoteMember from ..psycopg import quote_ident, quote_literal -from ..utils import deep_compare +from ..utils import deep_compare, unquote logger = logging.getLogger(__name__) @@ -26,7 +27,56 @@ class Bootstrap(object): return self._running_custom_bootstrap and self._keep_existing_recovery_conf @staticmethod - def process_user_options(tool, options, not_allowed_options, error_handler): + def process_user_options(tool: str, + options: Union[Dict[str, str], List[Union[str, Dict[str, str]]]], + not_allowed_options: Tuple[str, ...], + error_handler: Callable[[str], None]) -> List: + """Format *options* in a list or dictionary format into command line long form arguments. + + .. note:: + The format of the output of this method is to prepare arguments for use in the ``initdb`` + method of `self._postgres`. + + :Example: + + The *options* can be defined as a dictionary of key, values to be converted into arguments: + >>> Bootstrap.process_user_options('foo', {'foo': 'bar'}, (), print) + ['--foo=bar'] + + Or as a list of single string arguments + >>> Bootstrap.process_user_options('foo', ['yes'], (), print) + ['--yes'] + + Or as a list of key, value options + >>> Bootstrap.process_user_options('foo', [{'foo': 'bar'}], (), print) + ['--foo=bar'] + + Or a combination of single and key, values + >>> Bootstrap.process_user_options('foo', ['yes', {'foo': 'bar'}], (), print) + ['--yes', '--foo=bar'] + + Options that contain spaces are passed as is to ``subprocess.call`` + >>> Bootstrap.process_user_options('foo', [{'foo': 'bar baz'}], (), print) + ['--foo=bar baz'] + + Options that are quoted will be unquoted, so the quotes aren't interpreted + literally by the postgres command + >>> Bootstrap.process_user_options('foo', [{'foo': '"bar baz"'}], (), print) + ['--foo=bar baz'] + + .. note:: + The *error_handler* is called when any of these conditions are met: + + * Key, value dictionaries in the list form contains multiple keys. + * If a key is listed in *not_allowed_options*. + * If the options list is not in the required structure. + + :param tool: The name of the tool used in error reports to *error_handler* + :param options: Options to parse as a list of key, values or single values, or a dictionary + :param not_allowed_options: List of keys that cannot be used in the list of key, value formatted options + :param error_handler: A function which will be called when an error condition is encountered + :returns: List of long form arguments to pass to the named tool + """ user_options = [] def option_is_allowed(name): @@ -36,9 +86,9 @@ class Bootstrap(object): return ret if isinstance(options, dict): - for k, v in options.items(): - if k and v: - user_options.append('--{0}={1}'.format(k, v)) + for key, val in options.items(): + if key and val: + user_options.append('--{0}={1}'.format(key, unquote(val))) elif isinstance(options, list): for opt in options: if isinstance(opt, str) and option_is_allowed(opt): @@ -48,7 +98,7 @@ class Bootstrap(object): if len(keys) != 1 or not isinstance(opt[keys[0]], str) or not option_is_allowed(keys[0]): error_handler('Error when parsing {0} key-value option {1}: only one key-value is allowed' ' and value should be a string'.format(tool, opt[keys[0]])) - user_options.append('--{0}={1}'.format(keys[0], opt[keys[0]])) + user_options.append('--{0}={1}'.format(keys[0], unquote(opt[keys[0]]))) else: error_handler('Error when parsing {0} option {1}: value should be string value' ' or a single key-value pair'.format(tool, opt)) @@ -74,9 +124,8 @@ class Bootstrap(object): os.write(fd, self._postgresql.config.superuser['password'].encode('utf-8')) os.close(fd) options.append('--pwfile={0}'.format(pwfile)) - options = ['-o', ' '.join(options)] if options else [] - ret = self._postgresql.pg_ctl('initdb', *options) + ret = self._postgresql.initdb(*options) if pwfile: os.remove(pwfile) if ret: diff --git a/patroni/utils.py b/patroni/utils.py index 1939185a..f8f9e2fd 100644 --- a/patroni/utils.py +++ b/patroni/utils.py @@ -19,6 +19,7 @@ import socket import sys import tempfile import time +from shlex import split from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union, TYPE_CHECKING @@ -911,3 +912,32 @@ def enable_keepalive(sock: socket.socket, timeout: int, idle: int, cnt: Optional for opt in keepalive_socket_options(timeout, idle, cnt): sock.setsockopt(*opt) + + +def unquote(string: str) -> str: + """Unquote a fully quoted *string*. + + :Examples: + + A *string* with quotes will have those quotes removed + >>> unquote('"a quoted string"') + 'a quoted string' + + A *string* with multiple quotes will be returned as is + >>> unquote('"a multi" "quoted string"') + '"a multi" "quoted string"' + + So will a *string* with unbalanced quotes + >>> unquote('unbalanced "quoted string') + 'unbalanced "quoted string' + + :param string: The string to be checked for quoting. + :returns: The string with quotes removed, if it is a fully quoted single string, + or the original string if quoting is not detected, or unquoting was not possible. + """ + try: + ret = split(string) + ret = ret[0] if len(ret) == 1 else string + except ValueError: + ret = string + return ret diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py index 1cce58b9..bdc284e5 100644 --- a/tests/test_bootstrap.py +++ b/tests/test_bootstrap.py @@ -1,4 +1,5 @@ import os +import sys from mock import Mock, PropertyMock, patch @@ -99,6 +100,48 @@ class TestBootstrap(BaseTestPostgresql): self.assertRaises(Exception, self.b.bootstrap, {'initdb': [1]}) self.assertRaises(Exception, self.b.bootstrap, {'initdb': 1}) + def test__process_user_options(self): + def error_handler(msg): + raise Exception(msg) + + self.assertEqual(self.b.process_user_options('initdb', ['string'], (), error_handler), ['--string']) + self.assertEqual( + self.b.process_user_options( + 'initdb', + [{'key': 'value'}], + (), error_handler + ), + ['--key=value']) + if sys.platform != 'win32': + self.assertEqual( + self.b.process_user_options( + 'initdb', + [{'key': 'value with spaces'}], + (), error_handler + ), + ["--key=value with spaces"]) + self.assertEqual( + self.b.process_user_options( + 'initdb', + [{'key': "'value with spaces'"}], + (), error_handler + ), + ["--key=value with spaces"]) + self.assertEqual( + self.b.process_user_options( + 'initdb', + {'key': 'value with spaces'}, + (), error_handler + ), + ["--key=value with spaces"]) + self.assertEqual( + self.b.process_user_options( + 'initdb', + {'key': "'value with spaces'"}, + (), error_handler + ), + ["--key=value with spaces"]) + @patch.object(CancellableSubprocess, 'call', Mock()) @patch.object(Postgresql, 'is_running', Mock(return_value=True)) @patch.object(Postgresql, 'data_directory_empty', Mock(return_value=False)) diff --git a/tests/test_utils.py b/tests/test_utils.py index 0c6b21e4..36925041 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,8 +1,9 @@ import unittest from mock import Mock, patch + from patroni.exceptions import PatroniException -from patroni.utils import Retry, RetryFailedError, enable_keepalive, polling_loop, validate_directory +from patroni.utils import Retry, RetryFailedError, enable_keepalive, polling_loop, validate_directory, unquote class TestUtils(unittest.TestCase): @@ -41,6 +42,31 @@ class TestUtils(unittest.TestCase): with patch('sys.platform', platform): self.assertIsNone(enable_keepalive(Mock(), 10, 5)) + def test_unquote(self): + self.assertEqual(unquote('value'), 'value') + self.assertEqual(unquote('value with spaces'), "value with spaces") + self.assertEqual(unquote( + '"double quoted value"'), + 'double quoted value') + self.assertEqual(unquote( + '\'single quoted value\''), + 'single quoted value') + self.assertEqual(unquote( + 'value "with" double quotes'), + 'value "with" double quotes') + self.assertEqual(unquote( + '"value starting with" double quotes'), + '"value starting with" double quotes') + self.assertEqual(unquote( + '\'value starting with\' single quotes'), + '\'value starting with\' single quotes') + self.assertEqual(unquote( + 'value with a \' single quote'), + 'value with a \' single quote') + self.assertEqual(unquote( + '\'value with a \'"\'"\' single quote\''), + 'value with a \' single quote') + @patch('time.sleep', Mock()) class TestRetrySleeper(unittest.TestCase):