Modify Node::repair method for nodetool cluster repair

Since scylladb/scylladb#22905 nodetool repair
repairs only vnode keyspaces. To repair tablet keyspaces, you need
to use nodetool cluster repair.

Modify Node::repair method respectively. Node::repair method gets repair
parameters as arguments. It creates an object of a newly added RepairOptions
class to serialize the params. An old way of passing repair options
may still be used to repair vnode keyspaces.

Author: Deexie

ccmlib/cluster.py

@@ -616,6 +616,13 @@ def drain(self):
     def repair(self):
         self.nodetool("repair")
 
+        stdout, _ = self.nodetool('help')
+        if " cluster:" in stdout:
+            for node in list(self.nodes.values()):
+                if node.is_running():
+                    node.nodetool("cluster repair")
+                    break
+
     def cleanup(self):
         self.nodetool("cleanup")
 

ccmlib/node.py

@@ -1,10 +1,12 @@
 # ccm node
 
 
+from enum import Enum
 import errno
 import glob
 import itertools
 import os
+from pathlib import Path
 import re
 import shutil
 import signal
@@ -16,8 +18,9 @@
 from datetime import datetime
 import locale
 from collections import defaultdict, namedtuple
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
+import requests
 from ruamel.yaml import YAML
 
 from ccmlib import common
@@ -85,6 +88,124 @@ def __decode(self, value):
 _active_tasks_pattern = re.compile(r'\s*([\w-]+)\s+\w+\s+(?P<ks>\w+)\s+(?P<cf>\w+)\s+\d+\s+\d+\s+\w+\s+\d+\.\d+%')
 
 
+class RepairOptions:
+    class KeyspaceReplication(Enum):
+        TABLETS = "tablets",
+        VNODES = "vnodes",
+        ALL = "all"
+
+    class TabletRepairOptions:
+        def __init__(self, tablet_tokens: Optional[list[str]] = None):
+            self.tablet_tokens = tablet_tokens
+
+        def get_options(self):
+            options = []
+            if self.tablet_tokens:
+                options.append("--tablet-tokens")
+                options.append(",".join(self.tablet_tokens))
+
+            return " ".join(options)
+
+    class VnodeRepairOptions:
+        def __init__(self, dc_parallel: bool = False, end_token: Optional[int] = None, full: bool = False, local: bool = False,
+                     ignore_unreplicated_keyspaces: bool = False, job_threads: Optional[int] = None, partitioner_range: bool = False,
+                     pull: bool = False, sequential: bool = False, start_token: Optional[int] = None, trace: bool = False):
+            self.dc_parallel = dc_parallel
+            self.end_token = end_token
+            self.full = full
+            self.local = local
+            self.ignore_unreplicated_keyspaces = ignore_unreplicated_keyspaces
+            self.job_threads = job_threads
+            self.partitioner_range = partitioner_range
+            self.pull = pull
+            self.sequential = sequential
+            self.start_token = start_token
+            self.trace = trace
+
+        def get_options(self):
+            options = []
+            if self.dc_parallel:
+                options.append("--dc-parallel")
+
+            if self.end_token:
+                options.append("--end-token")
+                options.append(f"{self.end_token}")
+
+            if self.full:
+                options.append("--full")
+
+            if self.local:
+                options.append("--in-local-dc")
+
+            if self.ignore_unreplicated_keyspaces:
+                options.append("--ignore-unreplicated-keyspaces")
+
+            if self.job_threads:
+                options.append("--job-threads")
+                options.append(f"{self.job_threads}")
+
+            if self.partitioner_range:
+                options.append("--partitioner-range")
+
+            if self.pull:
+                options.append("--pull")
+
+            if self.sequential:
+                options.append("--sequential")
+
+            if self.start_token:
+                options.append("--start-token")
+                options.append(f"{self.start_token}")
+
+            if self.trace:
+                options.append("--trace")
+
+            return " ".join(options)
+
+    def __init__(self, replication: KeyspaceReplication = KeyspaceReplication.ALL, keyspace: Optional[str] = None, tables: Optional[list[str]] = None,
+                 hosts: Optional[list[str]] = None, dcs: Optional[list[str]] = None, dc_parallel: bool = False, end_token: Optional[int] = None,
+                 full: bool = False, local: bool = False, ignore_unreplicated_keyspaces: bool = False, job_threads: Optional[int] = None,
+                 partitioner_range: bool = False, pull: bool = False, sequential: bool = False, start_token: Optional[int] = None, trace: bool = False,
+                 tablet_tokens: Optional[list[str]] = None):
+        self.replication = replication
+        self.keyspace = keyspace
+        self.tables = tables
+        self.hosts = hosts
+        self.dcs = dcs
+        self.vnode_repair_options = self.VnodeRepairOptions(dc_parallel, end_token, full, local, ignore_unreplicated_keyspaces,
+                                                            job_threads, partitioner_range, pull, sequential, start_token, trace)
+        self.tablets_repair_options = self.TabletRepairOptions(tablet_tokens)
+
+    def _get_repair_options(self):
+        options = []
+        if self.keyspace:
+            options.append(self.keyspace)
+            if self.tables:
+                options.append(" ".join(self.tables))
+
+        if self.hosts:
+            options.append("--in-hosts")
+            options.append(",".join(self.hosts))
+
+        if self.dcs:
+            options.append("--in-dc")
+            options.append(",".join(self.dcs))
+
+        return " ".join(options)
+
+    def get_tablet_repair_options(self):
+        return self._get_repair_options() + " " + self.tablets_repair_options.get_options()
+
+    def get_vnode_repair_options(self):
+        return self._get_repair_options() + " " + self.vnode_repair_options.get_options()
+
+    def repair_tablets(self):
+        return self.replication in [self.KeyspaceReplication.TABLETS, self.KeyspaceReplication.ALL]
+
+    def repair_vnodes(self):
+        return self.replication in [self.KeyspaceReplication.VNODES, self.KeyspaceReplication.ALL]
+
+
 class Node(object):
 
     """
@@ -1418,11 +1539,69 @@ def drain(self, block_on_log=False):
         if block_on_log:
             self.watch_log_for("DRAINED", from_mark=mark)
 
-    def repair(self, options=[], **kwargs):
-        args = ["repair"] + options
-        cmd = ' '.join(args)
+    @property
+    def scylla_yaml(self) -> Dict[str, Any]:
+        return YAML().load(Path(self.get_conf_dir()) / common.SCYLLA_CONF)
+
+    @property
+    def api_port(self) -> int:
+        return self.scylla_yaml.get('api_port', 10000)
+
+    def cluster_command_available(self):
+        stdout, _ = self.nodetool('help')
+        return "  cluster  " in stdout
+
+    def run_repair(self, options, **kwargs):
+        cmd = "repair " + options
+        return self.nodetool(cmd, **kwargs)
+
+    def run_cluster_repair(self, options, **kwargs):
+        cmd = "cluster repair " + options
         return self.nodetool(cmd, **kwargs)
 
+    def repair(self, options=[], replication: RepairOptions.KeyspaceReplication = RepairOptions.KeyspaceReplication.ALL, keyspace: Optional[str] = None,
+               tables: Optional[list[str]] = None, hosts: Optional[list[str]] = None, dcs: Optional[list[str]] = None, dc_parallel: bool = False,
+               end_token: Optional[int] = None, full: bool = False, local: bool = False, ignore_unreplicated_keyspaces: bool = False,
+               job_threads: Optional[int] = None, partitioner_range: bool = False, pull: bool = False, sequential: bool = False, start_token: Optional[int] = None,
+               trace: bool = False, tablet_tokens: Optional[list[str]] = None, **kwargs):
+        if not options:
+            repair_options = RepairOptions(replication, keyspace, tables, hosts, dcs, dc_parallel, end_token, full, local, ignore_unreplicated_keyspaces,
+                                           job_threads, partitioner_range, pull, sequential, start_token, trace, tablet_tokens)
+            outs_and_errs = []
+            if repair_options.keyspace:
+                url = f"http://{self.address()}:{self.api_port}/storage_service/keyspaces"
+                resp = requests.get(url=url, params={"replication": "vnodes"})
+
+                if self.cluster_command_available() and resp.status_code == requests.codes.ok:
+                    if repair_options.keyspace in resp.json():
+                        res = self.run_repair(repair_options.get_vnode_repair_options(), **kwargs)
+                        outs_and_errs.append(res)
+                    else:
+                        res = self.run_cluster_repair(repair_options.get_tablet_repair_options(), **kwargs)
+                        outs_and_errs.append(res)
+                else:
+                    res = self.run_repair(repair_options.get_vnode_repair_options(), **kwargs)
+                    outs_and_errs.append(res)
+            else:
+                if self.cluster_command_available():
+                    if repair_options.repair_vnodes():
+                        res = self.run_repair(repair_options.get_vnode_repair_options(), **kwargs)
+                        outs_and_errs.append(res)
+
+                    if repair_options.repair_tablets() and self.cluster_command_available():
+                        res = self.run_cluster_repair(repair_options.get_tablet_repair_options(), **kwargs)
+                        outs_and_errs.append(res)
+                else:
+                    res = self.run_repair(repair_options.get_vnode_repair_options(), **kwargs)
+                    outs_and_errs.append(res)
+
+            outs, errs = zip(*outs_and_errs)
+            return outs, errs
+        else:
+            args = ["repair"] + options
+            cmd = ' '.join(args)
+            return self.nodetool(cmd, **kwargs)
+
     def move(self, new_token):
         self.nodetool("move " + str(new_token))