Implement missing tests and add pg-10 support to wale_restore(#446)

in addition to that get rid from two modules and fix formatting of tests

patroni/ctl.py

@@ -945,11 +945,10 @@ def apply_yaml_file(data, filename):
     return format_config_for_editing(changed_data), changed_data
 
 
-def invoke_editor(before_editing, data, cluster_name):
+def invoke_editor(before_editing, cluster_name):
     """Starts editor command to edit configuration in human readable format
 
     :param before_editing: human representation before editing
-    :param data: configuration datastructure
     :returns tuple of human readable and parsed datastructure after changes
     """
     editor_cmd = os.environ.get('EDITOR')
@@ -1003,7 +1002,7 @@ def edit_config(obj, cluster_name, force, quiet, kvpairs, pgkvpairs, apply_filen
 
     # If no changes were specified on the command line invoke editor
     if after_editing is None:
-        after_editing, changed_data = invoke_editor(before_editing, cluster.config.data, cluster_name)
+        after_editing, changed_data = invoke_editor(before_editing, cluster_name)
 
     if cluster.config.data == changed_data:
         if not quiet:

patroni/ha.py

@@ -538,7 +538,6 @@ class Ha(object):
                 self._async_executor.schedule('starting after demotion')
                 self._async_executor.run_async(self.state_handler.follow, (node_to_follow,))
             else:
-                logger.info('cluster.leader = %s', cluster.leader)
                 if self.state_handler.rewind_needed_and_possible(cluster.leader):
                     return False  # do not start postgres, but run pg_rewind on the next iteration
                 return self.state_handler.follow(node_to_follow)

patroni/scripts/wale_restore.py

@@ -23,35 +23,29 @@
 # currently also requires that you configure the restore_command to use wal_e, example:
 #       recovery_conf:
 #               restore_command: envdir /etc/wal-e.d/env wal-e wal-fetch "%f" "%p" -p 1
+import argparse
 import csv
-from collections import namedtuple
-import humanize
 import logging
 import os
-from enum import IntEnum
-
 import psycopg2
 import subprocess
 import sys
 import time
-import argparse
+
+from collections import namedtuple
 
 logger = logging.getLogger(__name__)
 
 RETRY_SLEEP_INTERVAL = 1
+si_prefixes = ['K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y']
 
 
-class ExitCode(IntEnum):
-    """
-    Gives meaningful names to the exit codes used by WALERestore
-    """
-
-    #: Succeeded
-    SUCCESS = 0
-    #: External issue, retry later
-    RETRY_LATER = 1
-    #: Don't try again unless configuration changes
-    FAIL = 2
+# Meaningful names to the exit codes used by WALERestore
+ExitCode = type('Enum', (), {
+    'SUCCESS': 0,  #: Succeeded
+    'RETRY_LATER': 1,  #: External issue, retry later
+    'FAIL': 2  #: Don't try again unless configuration changes
+})
 
 
 # We need to know the current PG version in order to figure out the correct WAL directory name
@@ -67,19 +61,33 @@ def get_major_version(data_dir):
 
 
 def repr_size(n_bytes):
-    return humanize.naturalsize(n_bytes, binary=True)
+    """
+    >>> repr_size(1000)
+    '1000 Bytes'
+    >>> repr_size(8257332324597)
+    '7.5 TiB'
+    """
+    if n_bytes < 1024:
+        return '{0} Bytes'.format(n_bytes)
+    i = -1
+    while n_bytes > 1023:
+        n_bytes /= 1024.0
+        i += 1
+    return '{0} {1}iB'.format(round(n_bytes, 1), si_prefixes[i])
 
 
 def size_as_bytes(size_, prefix):
-    si_prefixes = ['K', 'M', 'G', 'T'', P', 'E', 'Z', 'Y']
-
+    """
+    >>> size_as_bytes(7.5, 'T')
+    8246337208320
+    """
     prefix = prefix.upper()
 
     assert prefix in si_prefixes
 
     exponent = si_prefixes.index(prefix) + 1
 
-    return size_ * 1024.0 ** -exponent
+    return int(size_ * (1024.0 ** exponent))
 
 
 WALEConfig = namedtuple(
@@ -147,9 +155,6 @@ class WALERestore(object):
                 return ExitCode.FAIL
         except Exception:
             logger.exception("Unhandled exception when running WAL-E restore")
-            return ExitCode.FAIL
-
-        logger.warning('Missing exit code', stack_info=True)
         return ExitCode.FAIL
 
     def should_use_s3_to_create_replica(self):
@@ -211,17 +216,22 @@ class WALERestore(object):
                 try:
                     # get the difference in bytes between the current WAL location and the backup start offset
                     with psycopg2.connect(self.master_connection) as con:
+                        if con.server_version >= 100000:
+                            wal_name = 'wal'
+                            lsn_name = 'lsn'
+                        else:
+                            wal_name = 'xlog'
+                            lsn_name = 'location'
                         con.autocommit = True
                         with con.cursor() as cur:
                             cur.execute("""SELECT CASE WHEN pg_is_in_recovery()
                                                        THEN GREATEST(
-                                                                pg_xlog_location_diff(COALESCE(
-                                                                  pg_last_xlog_receive_location(), '0/0'), %s)::bigint,
-                                                                pg_xlog_location_diff(
-                                                                  pg_last_xlog_replay_location(), %s)::bigint)
-                                                       ELSE pg_xlog_location_diff(
-                                                                pg_current_xlog_location(), %s)::bigint
-                                                   END""", (backup_start_lsn, backup_start_lsn, backup_start_lsn))
+                                                                pg_{0}_{1}_diff(COALESCE(
+                                                                  pg_last_{0}_receive_{1}(), '0/0'), %s)::bigint,
+                                                                pg_{0}_{1}_diff(pg_last_{0}_replay_{1}(), %s)::bigint)
+                                                       ELSE pg_{0}_{1}_diff(pg_current_{0}_{1}(), %s)::bigint
+                                                   END""".format(wal_name, lsn_name),
+                                        (backup_start_lsn, backup_start_lsn, backup_start_lsn))
 
                             diff_in_bytes = int(cur.fetchone()[0])
                 except psycopg2.Error:
@@ -268,7 +278,7 @@ class WALERestore(object):
                 return ', '.join('{}={!r}'.format(key, value)
                                  for key, value in self.items)
 
-        human_context = HumanContext([
+        human_context = repr(HumanContext([
             ('threshold_size', Size(threshold_megabytes, 'M')),
             ('threshold_percent', threshold_percent),
             ('threshold_percent_size', Size(threshold_pct_bytes)),
@@ -276,19 +286,13 @@ class WALERestore(object):
             ('backup_diff', Size(diff_in_bytes)),
             ('is_size_thresh_ok', is_size_thresh_ok),
             ('is_percentage_thresh_ok', is_percentage_thresh_ok),
-        ])
+        ]))
 
         if not are_thresholds_ok:
-            logger.error(
-                'wal-e backup size diff is over threshold, falling back '
-                'to other means of restore. %r',
-                human_context
-            )
+            logger.info('wal-e backup size diff is over threshold, falling back '
+                        'to other means of restore: %s', human_context)
         else:
-            logger.info(
-                'Thresholds are OK, using wal-e basebackup. %r',
-                human_context
-            )
+            logger.info('Thresholds are OK, using wal-e basebackup: %s', human_context)
         return are_thresholds_ok
 
     def fix_subdirectory_path_if_broken(self, dirname):
@@ -322,7 +326,7 @@ class WALERestore(object):
             return ExitCode.RETRY_LATER
 
         if (exit_code == 0 and not
-           self.fix_subdirectory_path_if_broken('pg_xlog' if get_major_version(self.data_dir) < 10.0 else 'pg_wal')):
+           self.fix_subdirectory_path_if_broken('pg_xlog' if get_major_version(self.data_dir) < 10 else 'pg_wal')):
             return ExitCode.FAIL
         return exit_code
 

requirements.txt

@@ -11,6 +11,4 @@ click>=4.1
 prettytable>=0.7
 tzlocal
 python-dateutil
-enum34
 cdiff
-humanize==0.5.1

tests/test_ctl.py

@@ -8,7 +8,7 @@ from click.testing import CliRunner
 from mock import patch, Mock
 from patroni.ctl import ctl, members, store_config, load_config, output_members, request_patroni, get_dcs, parse_dcs, \
     get_all_members, get_any_member, get_cursor, query_member, configure, PatroniCtlException, apply_config_changes, \
-    format_config_for_editing, show_diff
+    format_config_for_editing, show_diff, invoke_editor
 from patroni.dcs.etcd import Client
 from psycopg2 import OperationalError
 from test_etcd import etcd_read, requests_get, socket_getaddrinfo, MockResponse
@@ -447,25 +447,27 @@ class TestCtl(unittest.TestCase):
 
         # Spaces are allowed and stripped, numbers and booleans are interpreted
         after_editing, changed_config = apply_config_changes(before_editing, config,
-            ["postgresql.parameters.work_mem = 5MB", "ttl=15", "postgresql.use_pg_rewind=off"])
-        self.assertEquals(changed_config,
-            {"postgresql": {"parameters": {"work_mem": "5MB"}, "use_pg_rewind": False}, "ttl": 15})
+                                                             ["postgresql.parameters.work_mem = 5MB",
+                                                              "ttl=15", "postgresql.use_pg_rewind=off", 'a.b=c'])
+        self.assertEquals(changed_config, {"a": {"b": "c"}, "postgresql": {"parameters": {"work_mem": "5MB"},
+                                                                           "use_pg_rewind": False}, "ttl": 15})
 
         # postgresql.parameters namespace is flattened
         after_editing, changed_config = apply_config_changes(before_editing, config,
-            ["postgresql.parameters.work_mem.sub = x"])
-        self.assertEquals(changed_config,
-            {"postgresql": {"parameters": {"work_mem": "4MB", "work_mem.sub": "x"}, "use_pg_rewind": True}, "ttl": 30})
+                                                             ["postgresql.parameters.work_mem.sub = x"])
+        self.assertEquals(changed_config, {"postgresql": {"parameters": {"work_mem": "4MB", "work_mem.sub": "x"},
+                                                          "use_pg_rewind": True}, "ttl": 30})
 
         # Setting to null deletes
         after_editing, changed_config = apply_config_changes(before_editing, config,
-            ["postgresql.parameters.work_mem=null"])
-        self.assertEquals(changed_config,
-            {"postgresql": {"use_pg_rewind": True}, "ttl": 30})
+                                                             ["postgresql.parameters.work_mem=null"])
+        self.assertEquals(changed_config, {"postgresql": {"use_pg_rewind": True}, "ttl": 30})
         after_editing, changed_config = apply_config_changes(before_editing, config,
-            ["postgresql.use_pg_rewind=null", "postgresql.parameters.work_mem=null"])
-        self.assertEquals(changed_config,
-            {"ttl": 30})
+                                                             ["postgresql.use_pg_rewind=null",
+                                                              "postgresql.parameters.work_mem=null"])
+        self.assertEquals(changed_config, {"ttl": 30})
+
+        self.assertRaises(PatroniCtlException, apply_config_changes, before_editing, config, ['a'])
 
     @patch('sys.stdout.isatty', return_value=False)
     @patch('cdiff.markup_to_pager')
@@ -480,3 +482,27 @@ class TestCtl(unittest.TestCase):
         # Test that unicode handling doesn't fail with an exception
         show_diff(b"foo:\n  bar: \xc3\xb6\xc3\xb6\n".decode('utf-8'),
                   b"foo:\n  bar: \xc3\xbc\xc3\xbc\n".decode('utf-8'))
+
+    def test_invoke_editor(self):
+        for e in ('', 'false'):
+            os.environ['EDITOR'] = e
+            self.assertRaises(PatroniCtlException, invoke_editor, 'foo: bar\n', 'test')
+
+    @patch('patroni.ctl.get_dcs')
+    def test_show_config(self, mock_get_dcs):
+        mock_get_dcs.return_value = self.e
+        mock_get_dcs.return_value.get_cluster = get_cluster_initialized_with_leader
+        self.runner.invoke(ctl, ['show-config', 'dummy'])
+
+    @patch('patroni.ctl.get_dcs')
+    def test_edit_config(self, mock_get_dcs):
+        mock_get_dcs.return_value = self.e
+        mock_get_dcs.return_value.get_cluster = get_cluster_initialized_with_leader
+        os.environ['EDITOR'] = 'true'
+        self.runner.invoke(ctl, ['edit-config', 'dummy'])
+        self.runner.invoke(ctl, ['edit-config', 'dummy', '-s', 'foo=bar'])
+        self.runner.invoke(ctl, ['edit-config', 'dummy', '--replace', 'postgres0.yml'])
+        self.runner.invoke(ctl, ['edit-config', 'dummy', '--apply', '-'], input='foo: bar')
+        self.runner.invoke(ctl, ['edit-config', 'dummy', '--force', '--apply', '-'], input='foo: bar')
+        mock_get_dcs.return_value.set_config_value = Mock(return_value=True)
+        self.runner.invoke(ctl, ['edit-config', 'dummy', '--force', '--apply', '-'], input='foo: bar')

tests/test_postgresql.py

@@ -30,7 +30,7 @@ class MockCursor(object):
         elif sql.startswith('SELECT slot_name'):
             self.results = [('blabla',), ('foobar',)]
         elif sql.startswith('SELECT CASE WHEN pg_is_in_recovery()'):
-            self.results = [(0,)]
+            self.results = [(2,)]
         elif sql == 'SELECT pg_is_in_recovery()':
             self.results = [(False, )]
         elif sql.startswith('WITH replication_info AS ('):
@@ -453,7 +453,7 @@ class TestPostgresql(unittest.TestCase):
         self.assertTrue(self.p.promote())
 
     def test_last_operation(self):
-        self.assertEquals(self.p.last_operation(), '0')
+        self.assertEquals(self.p.last_operation(), '2')
         Thread(target=self.p.last_operation).start()
 
     @patch('os.path.isfile', Mock(return_value=True))

tests/test_wale_restore.py

@@ -1,13 +1,12 @@
 import psycopg2
 import subprocess
 import unittest
-import pytest
 
-from mock import Mock, MagicMock, patch, mock_open
+from mock import Mock, PropertyMock, patch, mock_open
 from patroni.scripts import wale_restore
-from patroni.scripts.wale_restore import WALERestore, main as _main, \
-    get_major_version, ExitCode
+from patroni.scripts.wale_restore import WALERestore, main as _main, get_major_version
 from six.moves import builtins
+from test_postgresql import MockConnect, psycopg2_connect
 
 
 wale_output_header = (
@@ -17,7 +16,7 @@ wale_output_header = (
     b'wal_segment_backup_stop\twal_segment_offset_backup_stop\n'
 )
 
-wale_output_values =  (
+wale_output_values = (
     b'base_00000001000000000000007F_00000040\t2015-05-18T10:13:25.000Z\t'
     b'167772160\t'
     b'00000001000000000000007F\t00000040\t'
@@ -26,124 +25,29 @@ wale_output_values =  (
 
 wale_output = wale_output_header + wale_output_values
 
-wale_restore.RETRY_SLEEP_INTERVAL = 0.1  # Speed up retries
+wale_restore.RETRY_SLEEP_INTERVAL = 0.001  # Speed up retries
 WALE_TEST_RETRIES = 2
 
 
-def make_wale_restore():
-    return WALERestore(
-        scope="batman",
-        datadir="/data",
-        connstring="host=batman port=5432 user=batman",
-        env_dir="/etc",
-        threshold_mb=100,
-        threshold_pct=100,
-        use_iam=1,
-        no_master=0,
-        retries=WALE_TEST_RETRIES,
-    )
-
-
-@pytest.fixture(params=[
-        # Nn space
-        wale_output,
-        # Space
-        wale_output.replace(
-            b'\t2015-05-18T10:13:25.000Z',
-            b'\t2015-05-18 10:13:25.000Z'),
-    ])
-def fx_wale_spaces(request):
-    return request.param
-
-
-@pytest.fixture()
-def fx_wale_restore(request):
-    patches = [
-        patch('psycopg2.extensions.cursor', Mock(autospec=True)),
-        patch('psycopg2.extensions.connection', Mock(autospec=True)),
-        patch('psycopg2.connect', MagicMock(autospec=True)),
-    ]
-    for patch_ in patches:
-        patch_.start()
-
-    def _finalize():
-        for patch_ in patches:
-            patch_.stop()
-
-    request.addfinalizer(_finalize)
-
-    return make_wale_restore()
-
-
-@pytest.mark.parametrize('exit_code_int,exit_code', [
-    (0, ExitCode.SUCCESS),
-    (1, ExitCode.RETRY_LATER),
-    (2, ExitCode.FAIL),
-])
-def test_exit_code_enum_members_are_int_compatible(exit_code_int, exit_code):
-    assert exit_code_int == exit_code
-
-
-@pytest.mark.parametrize('mock,exit_code', [
-    (Mock(return_value=True), ExitCode.SUCCESS),
-    (Mock(return_value=False), ExitCode.FAIL),
-    (Mock(return_value=None), ExitCode.RETRY_LATER),  # Handled exception
-    (Mock(side_effect=Exception('Unhandled exception')), ExitCode.FAIL)
-])
-def test_run_exit_codes_by_should_use_s3(mock, exit_code, fx_wale_restore):
-    """
-    Verify that WALERestore.run() returns the correct values based on the
-    results of WALERestore.should_use_s3t_to_create_replica().
-    """
-    with patch.object(fx_wale_restore, 'should_use_s3_to_create_replica',
-                      mock),\
-            patch.object(fx_wale_restore, 'create_replica_with_s3',
-                         Mock(return_value=ExitCode.SUCCESS)):
-        assert fx_wale_restore.run() == exit_code
-
-
-def test_should_use_s3_too_many_rows(fx_wale_restore):
-    with patch('subprocess.check_output',
-               Mock(return_value=wale_output_header +
-                       wale_output_values +
-                       wale_output_values)):
-        assert not fx_wale_restore.should_use_s3_to_create_replica()
-
-
-def test_should_use_s3_handles_space_in_date(fx_wale_restore, fx_wale_spaces):
-    with patch('subprocess.check_output',
-               Mock(return_value=fx_wale_spaces)):
-
-        assert fx_wale_restore.should_use_s3_to_create_replica()
-
-
-def test_should_use_s3_missing_unused_field(fx_wale_restore):
-    with patch('subprocess.check_output',
-               Mock(return_value=wale_output.replace(b'\twal_segment_offset_backup_stop', b''))):
-        assert fx_wale_restore.should_use_s3_to_create_replica()
-
-
-def test_should_use_s3_missing_used_field(fx_wale_restore):
-    with patch('subprocess.check_output',
-               Mock(return_value=wale_output.replace(b'expanded_size_bytes', b'expanded_size_foo'))):
-        assert fx_wale_restore.should_use_s3_to_create_replica() is None
-
-
 @patch('os.access', Mock(return_value=True))
 @patch('os.makedirs', Mock(return_value=True))
 @patch('os.path.exists', Mock(return_value=True))
 @patch('os.path.isdir', Mock(return_value=True))
-@patch('psycopg2.extensions.cursor', Mock(autospec=True))
-@patch('psycopg2.extensions.connection', Mock(autospec=True))
-@patch('psycopg2.connect', MagicMock(autospec=True))
+@patch('psycopg2.connect', psycopg2_connect)
 @patch('subprocess.check_output', Mock(return_value=wale_output))
 class TestWALERestore(unittest.TestCase):
 
     def setUp(self):
-        self.wale_restore = make_wale_restore()
+        self.wale_restore = WALERestore('batman', '/data', 'host=batman port=5432 user=batman',
+                                        '/etc', 100, 100, 1, 0, WALE_TEST_RETRIES)
 
     def test_should_use_s3_to_create_replica(self):
         self.assertTrue(self.wale_restore.should_use_s3_to_create_replica())
+        with patch.object(MockConnect, 'server_version', PropertyMock(return_value=100000)):
+            self.assertTrue(self.wale_restore.should_use_s3_to_create_replica())
+
+        with patch('subprocess.check_output', Mock(return_value=wale_output.replace(b'167772160', b'1'))):
+            self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
 
         with patch('psycopg2.connect', Mock(side_effect=psycopg2.Error("foo"))):
             save_no_master = self.wale_restore.no_master
@@ -167,7 +71,12 @@ class TestWALERestore(unittest.TestCase):
 
         with patch('subprocess.check_output', Mock(side_effect=subprocess.CalledProcessError(1, "cmd", "foo"))):
             self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
-        with patch('subprocess.check_output', Mock(return_value=wale_output.split(b'\n')[0])):
+        with patch('subprocess.check_output', Mock(return_value=wale_output_header)):
+            self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
+        with patch('subprocess.check_output', Mock(return_value=wale_output + wale_output_values)):
+            self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
+        with patch('subprocess.check_output',
+                   Mock(return_value=wale_output.replace(b'expanded_size_bytes', b'expanded_size_foo'))):
             self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
 
     def test_create_replica_with_s3(self):
@@ -186,10 +95,12 @@ class TestWALERestore(unittest.TestCase):
         with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', Mock(return_value=True)):
             with patch.object(self.wale_restore, 'create_replica_with_s3', Mock(return_value=0)):
                 self.assertEqual(self.wale_restore.run(), 0)
-            with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', Mock(return_value=None)):
-                self.assertEqual(self.wale_restore.run(), 1)
-            with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', Mock(side_effect=Exception)):
-                self.assertEqual(self.wale_restore.run(), 2)
+        with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', Mock(return_value=False)):
+            self.assertEqual(self.wale_restore.run(), 2)
+        with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', Mock(return_value=None)):
+            self.assertEqual(self.wale_restore.run(), 1)
+        with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', Mock(side_effect=Exception)):
+            self.assertEqual(self.wale_restore.run(), 2)
 
     @patch('sys.exit', Mock())
     def test_main(self):
@@ -197,7 +108,7 @@ class TestWALERestore(unittest.TestCase):
             self.assertEqual(_main(), 0)
 
         with patch.object(WALERestore, 'run', Mock(return_value=1)), \
-             patch('time.sleep', Mock(return_value=None)) as mock_sleep:
+                patch('time.sleep', Mock(return_value=None)) as mock_sleep:
             self.assertEqual(_main(), 1)
             assert mock_sleep.call_count == WALE_TEST_RETRIES