Add reset-all command for pgconsul_util

src/cli.py

@@ -15,7 +15,7 @@
 from . import read_config, init_logging, zk as zookeeper
 from . import helpers
 from . import utils
-from .exceptions import SwitchoverException, FailoverException
+from .exceptions import SwitchoverException, FailoverException, ResetException
 
 
 class ParseHosts(argparse.Action):
@@ -178,6 +178,34 @@ def failover(opts, conf):
         sys.exit(1)
 
 
+def reset_all(opts, conf):
+    """
+    Resets all nodes in ZK, except for zk.MEMBERS_PATH
+    """
+    conf.set('global', 'iteration_timeout', 5)
+    zk = zookeeper.Zookeeper(config=conf, plugins=None)
+    logging.debug("resetting all ZK nodes")
+    all_nodes = zk.get_children("")
+    if all_nodes is None:
+        logging.error("Could not get nodes to reset")
+        all_nodes = []
+    nodes_to_delete = [x for x in all_nodes if x != zk.MEMBERS_PATH]
+    if not opts.force:
+        prompt = f'Nodes to delete: {", ".join(nodes_to_delete)}\n' \
+                 f'This is a potentially dangerous action. Proceed [y/n]?\n'
+        ans = input(prompt)
+        if ans == 'n':
+            return
+        elif ans != 'y':
+            print('Incorrect value, please type "y" or "n"')
+            return
+    for node in nodes_to_delete:
+        logging.debug(f'resetting path "{node}"')
+        if not zk.delete(node, recursive=True):
+            raise ResetException(f'Could not reset node "{node}" in ZK')
+    logging.debug("ZK structures are reset")
+
+
 def show_info(opts, conf):
     """
     Show cluster's information
@@ -387,6 +415,16 @@ def parse_args():
     )
     fail_arg.set_defaults(action=failover)
 
+    reset_all_arg = subarg.add_parser('reset-all', help='reset all nodes except members')
+    reset_all_arg.add_argument(
+        '-f',
+        '--force',
+        help='do not prompt',
+        default=False,
+        action='store_true',
+    )
+    reset_all_arg.set_defaults(action=reset_all)
+
     try:
         return arg.parse_args()
     except ValueError as err:

src/exceptions.py

@@ -42,3 +42,11 @@ class PGIsStartingUp(pgconsulException):
     """
 
     pass
+
+
+class ResetException(pgconsulException):
+    """
+    Exception for fatal errors during reset-all command
+    """
+
+    pass

tests/features/pgconsul_util.feature

@@ -884,7 +884,55 @@ Feature: Check pgconsul-util features
         | zookeeper | zookeeper1 |
 
 
-      @pgconsul_util_failover_reset
+    @pgconsul_util_reset_all
+    Scenario Outline: Check pgconsul-util reset-all command works as expected
+        Given a "pgconsul" container common config
+        """
+            pgconsul.conf:
+                global:
+                    priority: 0
+                    use_replication_slots: 'yes'
+                    quorum_commit: 'yes'
+                primary:
+                    change_replication_type: 'yes'
+                    primary_switch_checks: 1
+                replica:
+                    allow_potential_data_loss: 'no'
+                    primary_unavailability_timeout: 1
+                    primary_switch_checks: 1
+                    min_failover_timeout: 1
+                    primary_unavailability_timeout: 2
+                commands:
+                    generate_recovery_conf: /usr/local/bin/gen_rec_conf_with_slot.sh %m %p
+        """
+        Given a following cluster with "<lock_type>" with replication slots
+        """
+            postgresql1:
+                role: primary
+            postgresql2:
+                role: replica
+        """
+        Then <lock_type> "<lock_host>" has key "/pgconsul/postgresql/all_hosts/pgconsul_postgresql1_1.pgconsul_pgconsul_net"
+        When we set value "some_value" for key "/pgconsul/postgresql/some_key" in <lock_type> "<lock_host>"
+        When we set value "other_value" for key "/pgconsul/other_cluster/other_key" in <lock_type> "<lock_host>"
+        And we run following command on host "postgresql1"
+        """
+        pgconsul-util reset-all --force
+        """
+        Then command exit with return code "0"
+        And command result contains following output
+        """
+        resetting all ZK nodes
+        """
+        Then <lock_type> "<lock_host>" has key "/pgconsul/postgresql/all_hosts/pgconsul_postgresql1_1.pgconsul_pgconsul_net"
+        And <lock_type> "<lock_host>" doesn't have key "/pgconsul/postgresql/some_key"
+        Then <lock_type> "<lock_host>" has key "/pgconsul/other_cluster/other_key"
+    Examples: <lock_type>, <lock_host>
+        | lock_type | lock_host  |
+        | zookeeper | zookeeper1 |
+
+
+    @pgconsul_util_failover_reset
     Scenario Outline: Check pgconsul-util failover reset works as expected
         Given a "pgconsul" container common config
         """

tests/steps/helpers.py

@@ -322,6 +322,18 @@ def get_zk_value(context, zk_name, key):
     return None
 
 
+def zk_has_key(context, zk_name, key):
+    with contextlib.suppress(Exception):
+        zk = get_zk(context, zk_name)
+        zk.start()
+        try:
+            return zk.exists(key)
+        finally:
+            zk.stop()
+            zk.close()
+    return False
+
+
 def exec(container, cmd):
     """
     Execute command inside of given container

tests/steps/zk.py

@@ -101,6 +101,22 @@ def step_zk_value(context, name, value, key):
     )
 
 
+@then('zookeeper "{name}" has key "{key}"')
+@helpers.retry_on_assert
+def step_zk_key(context, name, key):
+    assert helpers.zk_has_key(context, name, key), '{time}: key "{key}" is missing'.format(
+        time=datetime.now().strftime("%H:%M:%S"), key=key
+    )
+
+
+@then('zookeeper "{name}" doesn\'t have key "{key}"')
+@helpers.retry_on_assert
+def step_zk_key(context, name, key):
+    assert not helpers.zk_has_key(context, name, key), '{time}: key "{key}" is present'.format(
+        time=datetime.now().strftime("%H:%M:%S"), key=key
+    )
+
+
 @then('zookeeper "{name}" has "{n}" values for key "{key}"')
 @helpers.retry_on_assert
 def step_zk_key_has_n_values(context, name, n, key):