From 26b6e005757abed7323601443c6b9c86cf47dc26 Mon Sep 17 00:00:00 2001
From: Igor Yanchenko <1504692+yanchenko-igor@users.noreply.github.com>
Date: Fri, 20 Dec 2019 13:05:39 +0200
Subject: [PATCH] wait option for patronictl reinit implemented (#1339)

Wait to finish `reinit` if `--wait` option is used.
Every 2 seconds it pulls the status from Patroni REST API and reports to console.
---
 patroni/ctl.py    | 23 +++++++++++++++++++++--
 tests/test_ctl.py | 11 +++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/patroni/ctl.py b/patroni/ctl.py
index ead8173d..10e14ef6 100644
--- a/patroni/ctl.py
+++ b/patroni/ctl.py
@@ -558,20 +558,39 @@ def restart(obj, cluster_name, member_names, force, role, p_any, scheduled, vers
 @click.argument('cluster_name')
 @click.argument('member_names', nargs=-1)
 @option_force
+@click.option('--wait', help='Wait until reinitialization completes', is_flag=True)
 @click.pass_obj
-def reinit(obj, cluster_name, member_names, force):
+def reinit(obj, cluster_name, member_names, force, wait):
     cluster = get_dcs(obj, cluster_name).get_cluster()
     members = get_members(cluster, cluster_name, member_names, None, force, 'reinitialize')
 
+    wait_on_members = []
     for member in members:
         body = {'force': force}
         while True:
             r = request_patroni(member, 'post', 'reinitialize', body)
-            if not check_response(r, member.name, 'reinitialize') and r.data.endswith(b' already in progress') \
+            started = check_response(r, member.name, 'reinitialize')
+            if not started and r.data.endswith(b' already in progress') \
                     and not force and click.confirm('Do you want to cancel it and reinitialize anyway?'):
                 body['force'] = True
                 continue
             break
+        if started and wait:
+            wait_on_members.append(member)
+
+    last_display = []
+    while wait_on_members:
+        if wait_on_members != last_display:
+            click.echo('Waiting for reinitialize to complete on: {0}'.format(
+                ", ".join(member.name for member in wait_on_members))
+            )
+            last_display[:] = wait_on_members
+        time.sleep(2)
+        for member in wait_on_members:
+            data = json.loads(request_patroni(member, 'get', 'patroni').data.decode('utf-8'))
+            if data.get('state') != 'creating replica':
+                click.echo('Reinitialize is completed on: {0}'.format(member.name))
+                wait_on_members.remove(member)
 
 
 def _do_failover_or_switchover(obj, action, cluster_name, master, candidate, force, scheduled=None):
diff --git a/tests/test_ctl.py b/tests/test_ctl.py
index c035a4ea..46759ff7 100644
--- a/tests/test_ctl.py
+++ b/tests/test_ctl.py
@@ -590,3 +590,14 @@ class TestCtl(unittest.TestCase):
         mock_get_dcs.return_value.get_cluster = get_cluster_not_initialized_without_leader
         result = self.runner.invoke(ctl, ['reinit', 'dummy'])
         assert "cluster doesn\'t have any members" in result.output
+
+    @patch('time.sleep', Mock())
+    @patch('patroni.ctl.get_dcs')
+    def test_reinit_wait(self, mock_get_dcs):
+        mock_get_dcs.return_value.get_cluster = get_cluster_initialized_with_leader
+        with patch.object(PoolManager, 'request') as mocked:
+            mocked.side_effect = [Mock(data=s, status=200) for s in
+                                  [b"reinitialize", b'{"state":"creating replica"}', b'{"state":"running"}']]
+            result = self.runner.invoke(ctl, ['reinit', 'alpha', 'other', '--wait'], input='y\ny')
+        self.assertIn("Waiting for reinitialize to complete on: other", result.output)
+        self.assertIn("Reinitialize is completed on: other", result.output)