#!/usr/bin/env python

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See LICENSE for more details.
#
# Copyright (c) 2016 ScyllaDB

import traceback
from itertools import zip_longest, chain, count as itertools_count
import json
import random
import time
import re
from functools import wraps, cache
from typing import List
import contextlib

import cassandra
import tenacity
from argus.client.sct.types import Package
from cassandra import ConsistencyLevel
from cassandra.query import SimpleStatement

from sdcm import argus_results
from sdcm import wait
from sdcm.cluster import BaseNode
from sdcm.utils.issues import SkipPerIssues
from sdcm.fill_db_data import FillDatabaseData
from sdcm.sct_events import Severity
from sdcm.stress_thread import CassandraStressThread
from sdcm.utils.parallel_object import ParallelObject, ParallelObjectException
from sdcm.utils.decorators import retrying
from sdcm.utils.sstable.sstable_utils import get_sstable_data_dump_command
from sdcm.utils.user_profile import get_profile_content
from sdcm.utils.version_utils import (
    get_node_supported_sstable_versions,
    get_node_enabled_sstable_version,
    ComparableScyllaVersion,
)
from sdcm.sct_events.system import InfoEvent
from sdcm.sct_events.database import (
    IndexSpecialColumnErrorEvent,
    DatabaseLogEvent,
)
from sdcm.sct_events.filters import DbEventsFilter
from sdcm.sct_events.group_common_events import (
    decorate_with_context,
    ignore_abort_requested_errors,
    ignore_upgrade_schema_errors,
    ignore_ycsb_connection_refused,
    ignore_raft_topology_cmd_failing,
)
from sdcm.utils import loader_utils
from sdcm.utils.features import CONSISTENT_TOPOLOGY_CHANGES_FEATURE, is_tablets_feature_enabled
from sdcm.wait import wait_for
from sdcm.rest.raft_upgrade_procedure import RaftUpgradeProcedure
from test_lib.sla import create_sla_auth
from sdcm.exceptions import Group0LimitedVotersFeatureNotEnableOnNodes


NUMBER_OF_ROWS_FOR_TRUNCATE_TEST = 10


def truncate_entries(func):
    @wraps(func)
    def inner(self, *args, **kwargs):
        if self.do_truncates:
            node = args[0]
            with self.db_cluster.cql_connection_patient(node, keyspace="truncate_ks", connect_timeout=600) as session:
                self.actions_log.info("Start truncate simple tables")
                session.default_timeout = 60.0 * 5
                session.default_consistency_level = ConsistencyLevel.QUORUM
                try:
                    self.cql_truncate_simple_tables(session=session, rows=NUMBER_OF_ROWS_FOR_TRUNCATE_TEST)
                    self.actions_log.info("Finish truncate simple tables")
                except cassandra.DriverException as details:
                    InfoEvent(
                        message=f"Failed truncate simple tables. Error: {details!s}. Traceback: {traceback.format_exc()}",
                        severity=Severity.ERROR,
                    ).publish()
                self.validate_truncated_entries_for_table(session=session, system_truncated=True)

        func_result = func(self, *args, **kwargs)

        if self.do_truncates:
            # re-new connection
            with self.db_cluster.cql_connection_patient(node, keyspace="truncate_ks", connect_timeout=600) as session:
                session.default_timeout = 60.0 * 5
                session.default_consistency_level = ConsistencyLevel.QUORUM
                self.validate_truncated_entries_for_table(session=session, system_truncated=True)
                self.read_data_from_truncated_tables(session=session)
                self.cql_insert_data_to_simple_tables(session=session, rows=NUMBER_OF_ROWS_FOR_TRUNCATE_TEST)

        return func_result

    return inner


def check_reload_systemd_config(node):
    for i in ["scylla-server", "scylla-jmx"]:
        result = node.remoter.run("systemctl status %s" % i, ignore_status=True)
        if ".service changed on disk. Run 'systemctl daemon-reload' to reload units" in result.stderr:
            raise Exception("Systemd config is changed, but not reload automatically")


def backup_conf(node):
    if node.distro.is_rhel_like:
        node.remoter.run(
            r"for conf in $( rpm -qc $(rpm -qa | grep scylla) | grep -v contains ) "
            r"/etc/systemd/system/{var-lib-scylla,var-lib-systemd-coredump}.mount; "
            r"do if test -e $conf; then sudo cp -v $conf $conf.autobackup; fi; done"
        )
    else:
        node.remoter.run(
            r"for conf in $(cat /var/lib/dpkg/info/scylla-*server.conffiles "
            r"/var/lib/dpkg/info/scylla-*conf.conffiles "
            r"/var/lib/dpkg/info/scylla-*jmx.conffiles | grep -v init ) "
            r"/etc/systemd/system/{var-lib-scylla,var-lib-systemd-coredump}.mount; "
            r"do if test -e $conf; then sudo cp -v $conf $conf.backup; fi; done"
        )


def recover_conf(node):
    if node.distro.is_rhel_like:
        node.remoter.run(
            r"for conf in $( rpm -qc $(rpm -qa | grep scylla) | grep -v contains ) "
            r"/etc/systemd/system/{var-lib-scylla,var-lib-systemd-coredump}.mount; "
            r"do if test -e $conf.autobackup; then sudo cp -v $conf.autobackup $conf; fi; done"
        )
    else:
        node.remoter.run(
            r"for conf in $(cat /var/lib/dpkg/info/scylla-*server.conffiles "
            r"/var/lib/dpkg/info/scylla-*conf.conffiles "
            r"/var/lib/dpkg/info/scylla-*jmx.conffiles | grep -v init ); do "
            r"if test -e $conf.backup; then sudo cp -v $conf.backup $conf; fi; done"
        )


class UpgradeTest(FillDatabaseData, loader_utils.LoaderUtilsMixin):
    """
    Test a Scylla cluster upgrade.
    """

    def setUp(self):
        super().setUp()
        self.do_truncates = self.params.get("enable_truncate_checks_on_node_upgrade") or False
        self.stacks = {}
        for node in self.db_cluster.nodes:
            self.configure_event_filtering(node)

    def configure_event_filtering(self, node):
        self.stacks[node] = contextlib.ExitStack()
        # ignoring those oversized allocation errors, till both ends of the upgrade would have
        # fixes for https://github.com/scylladb/scylladb/issues/24660
        self.stacks[node].enter_context(
            DbEventsFilter(
                node=node,
                db_event=DatabaseLogEvent.OVERSIZED_ALLOCATION,
                line=r"seastar::rpc::client::wait_for_reply",
                extra_time_to_expiration=30,
            )
        )

    orig_ver = None
    new_ver = None

    # `major_release` (eg: 2.1 <-> 2.2, 2017.1 <-> 2018.1)
    # `reinstall` (opensource <-> enterprise, enterprise <-> opensource)
    # `minor_release` (eg: 2.2.1 <-> 2.2.5, 2018.1.0 <-> 2018.1.1)
    upgrade_rollback_mode = None

    # expected format version after upgrade and nodetool upgradesstables called
    # would be recalculated after all the cluster finish upgrade
    expected_sstable_format_version = "mc"

    system_upgrade_timeout = 30 * 60

    def should_do_complex_profile(self) -> bool:
        """
        since tablets doesn't support complex profiles with index, we need to replace them with simple profiles without index
        """
        if is_tablets_feature_enabled(self.db_cluster.nodes[0]):
            if SkipPerIssues("https://github.com/scylladb/scylladb/issues/22677", params=self.params):
                return False
        return True

    @retrying(n=5)
    def _query_from_one_table(self, session, query, table_name) -> list:
        return self.rows_to_list(session.execute(SimpleStatement(query.format(table_name)), timeout=300))

    def read_data_from_truncated_tables(self, session):
        truncate_query = "SELECT COUNT(*) FROM {}"
        tables_name = self.get_tables_name_of_keyspace(session=session, keyspace_name="truncate_ks")
        for table_name in tables_name:
            self.actions_log.info(f"Start read data from {table_name} tables")
            try:
                count = self._query_from_one_table(session=session, query=truncate_query, table_name=table_name)
                self.assertEqual(
                    str(count[0][0]),
                    "0",
                    msg=f"Expected that there is no data in the table truncate_ks.{table_name}, but found {count[0][0]} rows",
                )
                self.actions_log.info(f"Finish read data from {table_name} tables ")
            except Exception as details:  # noqa: BLE001
                InfoEvent(
                    message=f"Failed read data from {table_name} tables. Error: {details!s}. Traceback: {traceback.format_exc()}",
                    severity=Severity.ERROR,
                ).publish()

    def validate_truncated_entries_for_table(self, session, system_truncated=False):
        tables_id = self.get_tables_id_of_keyspace(session=session, keyspace_name="truncate_ks")

        for table_id in tables_id:
            if system_truncated:
                # validate truncation entries in the system.truncated table - expected entry
                truncated_time = self.get_truncated_time_from_system_truncated(session=session, table_id=table_id)
                self.assertTrue(
                    truncated_time, msg="Expected truncated entry in the system.truncated table, but it's not found"
                )

            # validate truncation entries in the system.local table - not expected entry
            truncated_time = self.get_truncated_time_from_system_local(session=session)
            if system_truncated:
                self.assertEqual(
                    truncated_time,
                    [[None]],
                    msg="Not expected truncated entry in the system.local table, but it's found",
                )
            else:
                self.assertTrue(
                    truncated_time, msg="Expected truncated entry in the system.local table, but it's not found"
                )

    @truncate_entries
    def upgrade_node(self, node, upgrade_sstables=True):
        self._upgrade_node(node=node, upgrade_sstables=upgrade_sstables)

    @decorate_with_context(
        [
            ignore_abort_requested_errors,
            ignore_ycsb_connection_refused,
            ignore_raft_topology_cmd_failing,
        ]
    )
    # https://github.com/scylladb/scylla/issues/10447#issuecomment-1194155163
    def _upgrade_node(self, node, upgrade_sstables=True, new_scylla_repo=None, new_version=None):  # noqa: PLR0915
        self.actions_log.info(f"Starting {node.name} node upgrade")
        new_scylla_repo = new_scylla_repo or self.params.get("new_scylla_repo")
        new_version = new_version or self.params.get("new_version")
        upgrade_node_packages = self.params.get("upgrade_node_packages")

        scylla_yaml_updates = {}
        if not self.params.get("disable_raft"):
            scylla_yaml_updates.update({"consistent_cluster_management": True})
            scylla_yaml_updates.update({"force_schema_commit_log": True})

        if self.params.get("enable_tablets_on_upgrade"):
            scylla_yaml_updates.update({"enable_tablets": True})
            scylla_yaml_updates.update({"tablets_mode_for_new_keyspaces": "enabled"})
        if self.params.get("enable_views_with_tablets_on_upgrade"):
            scylla_yaml_updates.update({"experimental_features": ["views-with-tablets"]})

        self.actions_log.info(f"Upgrading {node.name} Node")
        # We assume that if update_db_packages is not empty we install packages from there.
        # In this case we don't use upgrade based on new_scylla_repo(ignored sudo yum update scylla...)
        result = node.remoter.run("scylla --version")
        self.orig_ver = result.stdout.strip()

        if upgrade_node_packages:
            self.actions_log.info("Started upgrade_node_packages")
            # update_scylla_packages
            self.actions_log.info("sending upgrade packages to node")
            with self.actions_log.action_scope("sending upgrade packages to node"):
                node.remoter.send_files(upgrade_node_packages, "/tmp/scylla", verbose=True)
            # node.remoter.run('sudo yum update -y --skip-broken', connect_timeout=900)
            node.remoter.run("sudo yum install python34-PyYAML -y")
            # replace the packages
            node.remoter.run(r"rpm -qa scylla\*")
            # flush all memtables to SSTables
            with self.actions_log.action_scope("stopping node"):
                node.run_nodetool("drain", timeout=15 * 60, coredump_on_timeout=True, long_running=True, retry=0)
                node.run_nodetool("snapshot")
                node.stop_scylla_server()
            with self.actions_log.action_scope("upgrading packages"):
                node.remoter.run("sudo rpm -UvhR --oldpackage /tmp/scylla/*development*", ignore_status=True)
                node.remoter.run("sudo rpm -URvh --replacefiles /tmp/scylla/*.rpm | true")
                node.remoter.run(r"rpm -qa scylla\*")
        elif new_scylla_repo:
            # backup the data
            node.remoter.run("sudo cp /etc/scylla/scylla.yaml /etc/scylla/scylla.yaml-backup")
            if node.distro.is_rhel_like:
                node.remoter.run("sudo cp /etc/yum.repos.d/scylla.repo ~/scylla.repo-backup")
            else:
                node.remoter.run("sudo cp /etc/apt/sources.list.d/scylla.list ~/scylla.list-backup")
            backup_conf(node)
            assert new_scylla_repo.startswith("http")
            node.download_scylla_repo(new_scylla_repo)
            # flush all memtables to SSTables
            with self.actions_log.action_scope("stop node"):
                node.run_nodetool("drain", timeout=15 * 60, coredump_on_timeout=True, long_running=True, retry=0)
                node.run_nodetool("snapshot")
                node.stop_scylla_server(verify_down=False)

            orig_is_enterprise = node.is_product_enterprise
            if node.distro.is_rhel_like:
                result = node.remoter.run("sudo yum search scylla-enterprise-server 2>&1", ignore_status=True)
                new_is_enterprise = "scylla-enterprise-server" in result.stdout
            else:
                result = node.remoter.run("sudo apt-cache search scylla-enterprise-server", ignore_status=True)
                new_is_enterprise = "scylla-enterprise-server" in result.stdout

            scylla_pkg = "scylla-enterprise" if new_is_enterprise else "scylla"
            ver_suffix = rf"\*{new_version}" if new_version else ""
            scylla_pkg_ver = f"{scylla_pkg}{ver_suffix}"

            if orig_is_enterprise and ComparableScyllaVersion(
                self.orig_ver
            ) < "2025.1.0~dev" <= ComparableScyllaVersion(self.params.scylla_version_upgrade_target):
                InfoEvent(
                    message=f"upgrade_node - orig is enterprise and version {self.orig_ver} < 2025.1.0~dev and "
                    f"scylla_version_upgrade_target is {self.params.scylla_version_upgrade_target}"
                ).publish()
                scylla_pkg = "scylla"
                scylla_pkg_ver = f"{scylla_pkg}{ver_suffix}"
                InfoEvent(message="Rollback mode is set to reinstall").publish()
                self.upgrade_rollback_mode = "reinstall"

                if self.params.get("use_preinstalled_scylla"):
                    scylla_pkg_ver += f" {scylla_pkg}-machine-image"
            with self.actions_log.action_scope("updating packages"):
                if node.distro.is_rhel_like:
                    node.remoter.run(rf"sudo yum update {scylla_pkg_ver}\* -y", retry=3, timeout=600)
                else:
                    node.remoter.sudo("apt-get update", retry=3, timeout=600)
                    node.remoter.sudo(
                        f"DEBIAN_FRONTEND=noninteractive apt-get dist-upgrade {scylla_pkg_ver} -y"
                        f' -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold"',
                        retry=3,
                        timeout=600,
                    )

        node.remoter.sudo(
            f"sed -i 's/--num-io-queues/--num-io-groups/' {node.add_install_prefix('/etc/scylla.d/io.conf')}"
        )

        with self.actions_log.action_scope("reload systemd and run scylla sysconfig setup"):
            check_reload_systemd_config(node)
            node.run_scylla_sysconfig_setup()
        if scylla_yaml_updates:
            with self.actions_log.action_scope(f"update_scylla_yaml with: {scylla_yaml_updates}"):
                self._update_scylla_yaml_on_node(node_to_update=node, updates=scylla_yaml_updates)
        node.forget_scylla_version()
        node.drop_raft_property()

        # remove filters on new release, as error should be fixed
        self.stacks[node].close()

        with self.actions_log.action_scope("start_scylla_server"):
            node.start_scylla_server(verify_up_timeout=500)
        self.db_cluster.get_db_nodes_cpu_mode()
        result = node.remoter.run("scylla --version")
        new_ver = result.stdout.strip()
        assert self.orig_ver != new_ver, "scylla-server version didn't change during upgrade"
        self.new_ver = new_ver
        self._update_argus_upgraded_version(node, new_ver)
        if upgrade_sstables:
            self.upgradesstables_if_command_available(node)

        self._verify_node_exporter_after_upgrade(node)

        self.db_cluster.wait_all_nodes_un()
        self.actions_log.info(f"Upgrade {node.name} node completed")

    def _verify_node_exporter_after_upgrade(self, node):
        """Verify scylla-node-exporter service runs the binary shipped by the installed package."""
        with self.actions_log.action_scope("verify_node_exporter"):
            # Check the service is active
            result = node.remoter.run("systemctl is-active scylla-node-exporter", ignore_status=True)
            assert result.ok, f"scylla-node-exporter service is NOT active on {node.name}"
            self.log.info("scylla-node-exporter service is active on %s", node.name)

            # Get the ExecStart binary path from the systemd unit
            result = node.remoter.run(
                r"systemctl show scylla-node-exporter -p ExecStart --value | grep -oP 'path=\K[^;\s]+'",
            )
            service_binary = result.stdout.strip()
            self.log.info("scylla-node-exporter ExecStart binary on %s: %s", node.name, service_binary)
            assert service_binary, f"Could not determine ExecStart binary for scylla-node-exporter on {node.name}"

            # Verify the binary is owned by the scylla-node-exporter package
            if node.distro.is_rhel_like or node.distro.is_sles:
                pkg_result = node.remoter.run(f"rpm -qf {service_binary}", ignore_status=True)
            else:
                pkg_result = node.remoter.run(f"dpkg -S {service_binary}", ignore_status=True)
            assert pkg_result.ok and "scylla-node-exporter" in pkg_result.stdout, (
                f"Binary {service_binary} on {node.name} is NOT owned by scylla-node-exporter package. "
                f"Output: {pkg_result.stdout.strip()}"
            )
            self.log.info(
                "Confirmed %s is owned by scylla-node-exporter package on %s",
                service_binary,
                node.name,
            )

            # Log the binary version for reference
            result = node.remoter.run(f"{service_binary} --version 2>&1 | head -1")
            self.log.info("node_exporter binary version on %s: %s", node.name, result.stdout.strip())

    def upgrade_os(self, nodes):
        def upgrade(node):
            node.upgrade_system()

        def wait_for_package_manager(node):
            if node.distro.is_rhel_like:
                lock_file = "/var/run/yum.pid"
            elif node.distro.is_sles:
                lock_file = "/run/zypp.pid"
            else:
                lock_file = "/var/lib/dpkg/lock-frontend"
            self.log.warning("OS upgrade timed out on %s, waiting for package manager to finish...", node.name)
            node.remoter.run(f"while fuser {lock_file} >/dev/null 2>&1; do sleep 10; done", timeout=10 * 60)
            self.log.info("OS upgrade on %s completed after initial timeout", node.name)

        if not self.params.get("upgrade_node_system"):
            return

        with self.actions_log.action_scope("upgrade OS on nodes"):
            try:
                ParallelObject(objects=nodes, timeout=self.system_upgrade_timeout).run(upgrade)
            except ParallelObjectException as exc:
                timed_out = [r for r in exc.results if isinstance(r.exc, TimeoutError)]
                other_failures = [r for r in exc.results if r.exc and not isinstance(r.exc, TimeoutError)]
                if other_failures:
                    raise
                # package manager may still be running on remote node, wait for it to finish
                for result in timed_out:
                    wait_for_package_manager(result.obj)

    @truncate_entries
    # https://github.com/scylladb/scylla/issues/10447#issuecomment-1194155163
    def rollback_node(self, node, upgrade_sstables=True):
        self._rollback_node(node=node, upgrade_sstables=upgrade_sstables)

    @decorate_with_context(
        [
            ignore_abort_requested_errors,
            ignore_ycsb_connection_refused,
            ignore_raft_topology_cmd_failing,
        ]
    )
    def _rollback_node(self, node, upgrade_sstables=True):
        self.actions_log.info(f"Starting {node.name} node rollback")
        result = node.remoter.run("scylla --version")
        orig_ver = result.stdout.strip()
        # flush all memtables to SSTables
        with self.actions_log.action_scope("stop node"):
            node.run_nodetool("drain", timeout=15 * 60, coredump_on_timeout=True, long_running=True, retry=0)
            # backup the data
            node.run_nodetool("snapshot")
            node.stop_scylla_server(verify_down=False)

        if node.distro.is_rhel_like:
            node.remoter.run("sudo cp ~/scylla.repo-backup /etc/yum.repos.d/scylla.repo")
            node.remoter.run("sudo chown root.root /etc/yum.repos.d/scylla.repo")
            node.remoter.run("sudo chmod 644 /etc/yum.repos.d/scylla.repo")
        else:
            node.remoter.run("sudo cp ~/scylla.list-backup /etc/apt/sources.list.d/scylla.list")
            node.remoter.run("sudo chown root.root /etc/apt/sources.list.d/scylla.list")
            node.remoter.run("sudo chmod 644 /etc/apt/sources.list.d/scylla.list")
        node.update_repo_cache()

        if re.findall(r"\d+.\d+", self.orig_ver)[0] == re.findall(r"\d+.\d+", self.new_ver)[0]:
            self.upgrade_rollback_mode = "minor_release"
        self.actions_log.info("downgrading packages")

        if self.upgrade_rollback_mode == "reinstall" or not node.distro.is_rhel_like:
            scylla_pkg_ver = node.scylla_pkg()

            if self.params.get("use_preinstalled_scylla"):
                scylla_pkg_ver += f" {scylla_pkg_ver}-machine-image"

            if node.distro.is_rhel_like:
                node.remoter.run(r"sudo yum remove scylla\* -y", retry=3)
                node.remoter.run(rf"sudo yum install {scylla_pkg_ver} -y", retry=3, timeout=600)
            else:
                node.remoter.sudo(r"apt-get remove scylla\* -y", retry=3)
                node.remoter.sudo(
                    rf"DEBIAN_FRONTEND=noninteractive apt-get install {scylla_pkg_ver} -y"
                    rf' -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold"',
                    retry=3,
                    timeout=600,
                )
            recover_conf(node)
            node.remoter.run("sudo systemctl daemon-reload")
        elif self.upgrade_rollback_mode == "minor_release":
            node.remoter.run(r"sudo yum downgrade scylla\*%s-\* -y" % self.orig_ver.split("-")[0], retry=3, timeout=600)
        else:
            node.remoter.run(r"sudo yum downgrade scylla\* -y", retry=3, timeout=600)
            recover_conf(node)
            node.remoter.run("sudo systemctl daemon-reload")

        result = node.remoter.run("scylla --version")
        new_ver = result.stdout.strip()
        self.actions_log.info("finished downgrading packages")

        node.remoter.run("sudo cp /etc/scylla/scylla.yaml-backup /etc/scylla/scylla.yaml")

        node.drop_raft_property()

        # enable the filter since we are moving to older version that might have still unfixed issues
        self.stacks[node].close()
        self.configure_event_filtering(node)

        # Current default 300s aren't enough for upgrade test of Debian 9.
        # Related issue: https://github.com/scylladb/scylla-cluster-tests/issues/1726
        node.run_scylla_sysconfig_setup()
        node.start_scylla_server(verify_up_timeout=500)

        assert orig_ver != new_ver, "scylla-server version didn't change during rollback"

        if upgrade_sstables:
            self.upgradesstables_if_command_available(node)

        self.db_cluster.wait_all_nodes_un()
        self.actions_log.info(f"Node {node.name} rollback completed")

    def upgradesstables_if_command_available(self, node, queue=None):
        upgradesstables_available = False
        upgradesstables_supported = node.remoter.run(
            'nodetool help | grep -q upgradesstables && echo "yes" || echo "no"'
        )
        if "yes" in upgradesstables_supported.stdout:
            upgradesstables_available = True
            with self.actions_log.action_scope("Upgrading SSTables"):
                node.run_nodetool(sub_cmd="upgradesstables")
        if queue:
            queue.put(upgradesstables_available)
            queue.task_done()

    def get_highest_supported_sstable_version(self):
        """
        find the highest sstable format version supported in the cluster

        :return:
        """
        enabled_sstable_format_features = []
        for node in self.db_cluster.nodes:
            enabled_sstable_format_features.extend(get_node_supported_sstable_versions(node.system_log))

        if not enabled_sstable_format_features:
            for node in self.db_cluster.nodes:
                with self.db_cluster.cql_connection_patient_exclusive(node) as session:
                    enabled_sstable_format_features.extend(get_node_enabled_sstable_version(session))

        return max(set(enabled_sstable_format_features))

    def wait_for_sstable_upgrade(self, node, queue=None):
        all_tables_upgraded = True

        def wait_for_node_to_finish():
            try:
                result = node.remoter.sudo(
                    r"find /var/lib/scylla/data/system -type f ! -path '*snapshots*' -printf %f\\n"
                )
                all_sstable_files = result.stdout.splitlines()

                sstable_version_regex = re.compile(r"(\w+)-[^-]+-(.+)\.(db|txt|sha1|crc32)")

                sstable_versions = {
                    sstable_version_regex.search(f).group(1)
                    for f in all_sstable_files
                    if sstable_version_regex.search(f)
                }

                assert len(sstable_versions) == 1, f"expected all table format to be the same found {sstable_versions}"
                assert list(sstable_versions)[0] == self.expected_sstable_format_version, (
                    f"expected to format version to be '{self.expected_sstable_format_version}', found '{list(sstable_versions)[0]}'"
                )
            except Exception as ex:  # noqa: BLE001
                self.log.warning(ex)
                return False
            else:
                return True

        try:
            wait.wait_for(
                func=wait_for_node_to_finish,
                step=30,
                timeout=900,
                throw_exc=True,
                text="Waiting until upgardesstables is finished",
            )
        except tenacity.RetryError:
            all_tables_upgraded = False
        finally:
            if queue:
                queue.put(all_tables_upgraded)
                queue.task_done()

    default_params = {"timeout": 650000}

    def test_upgrade_cql_queries(self):
        """
        Run a set of different cql queries against various types/tables before
        and after upgrade of every node to check the consistency of data
        """
        self.upgrade_os(self.db_cluster.nodes)

        self.actions_log.info("Populating DB with tables and data")
        self.fill_db_data()
        self.actions_log.info("Running queries to verify data before upgrade")
        self.verify_db_data()

        self.actions_log.info("Starting cassandra-stress write workload for 10M partitions")
        # YAML: stress_cmd: cassandra-stress write cl=QUORUM n=10000000
        # -schema 'replication(strategy=NetworkTopologyStrategy,replication_factor=3)'
        # -mode cql3 native -rate threads=1000 -pop seq=1..10000000
        self._run_all_stress_cmds([], params={"stress_cmd": self.params.get("stress_cmd")})

        self.actions_log.info("Waiting for cassandra-stress to populate data")
        time.sleep(600)

        self.actions_log.info("Starting cassandra-stress read workload for 60m")
        # YAML: stress_cmd_1: cassandra-stress read cl=QUORUM duration=60m
        # -schema 'replication(strategy=NetworkTopologyStrategy,replication_factor=3)'
        # -mode cql3 native -rate threads=100 -pop seq=1..10000000
        stress_queues = self._run_all_stress_cmds([], params={"stress_cmd": self.params.get("stress_cmd_1")})

        self.actions_log.info("Waiting for cassandra-stress to start before upgrade")
        time.sleep(300)

        nodes_num = len(self.db_cluster.nodes)
        # prepare an array containing the indexes
        indexes = list(range(nodes_num))
        # shuffle it so we will upgrade the nodes in a random order
        random.shuffle(indexes)

        # upgrade all the nodes in random order
        for i in indexes:
            self.db_cluster.node_to_upgrade = self.db_cluster.nodes[i]
            self.upgrade_node(self.db_cluster.node_to_upgrade)
            time.sleep(300)

        self.actions_log.info("Running queries to verify data after upgrade")
        self.verify_db_data()
        for stress_queue in stress_queues:
            self.verify_stress_thread(stress_queue)

    def fill_and_verify_db_data(self, note, pre_fill=False, rewrite_data=True):
        if pre_fill:
            self.actions_log.info("Populating DB with tables and data")
            self.fill_db_data()
        self.actions_log.info(f"Running queries to verify data (stage: {note})")
        self.verify_db_data()
        if rewrite_data:
            self.actions_log.info("Re-populating DB with tables and data")
            self.fill_db_data()

    # Added to cover the issue #5621: upgrade from 3.1 to 3.2 fails on std::logic_error (Column idx_token doesn't exist
    # in base and this view is not backing a secondary index)
    # @staticmethod
    def search_for_idx_token_error_after_upgrade(self, node, step):
        self.log.debug(f"Search for idx_token error. Step {step}")
        idx_token_error = list(
            node.follow_system_log(patterns=["Column idx_token doesn't exist"], start_from_beginning=True)
        )
        if idx_token_error:
            IndexSpecialColumnErrorEvent(
                message=f"Node: {node.name}. Step: {step}. "
                f'Found error: index special column "idx_token" is not recognized'
            ).publish()

    @staticmethod
    def shuffle_nodes_and_alternate_dcs(nodes: List[BaseNode]) -> List[BaseNode]:
        """Shuffles provided nodes based on dc (alternates dc's).
        based on https://stackoverflow.com/a/21482016/3891194"""
        dc_nodes = {}
        for node in nodes:
            dc_nodes.setdefault(node.dc_idx, []).append(node)
        for dc_node_list in dc_nodes.values():
            random.shuffle(dc_node_list)

        return [x for x in chain.from_iterable(zip_longest(*dc_nodes.values())) if x]

    @staticmethod
    def _update_scylla_yaml_on_node(node_to_update: BaseNode, updates: dict):
        with node_to_update.remote_scylla_yaml() as scylla_yaml:
            scylla_yaml.update(updates)

    def test_rolling_upgrade(self):  # noqa: PLR0914, PLR0915
        """
        Upgrade half of nodes in the cluster, and start special read workload
        during the stage. Checksum method is changed to xxhash from Scylla 2.2,
        we want to use this case to verify the read (cl=ALL) workload works
        well, upgrade all nodes to new version in the end.
        """
        self.upgrade_os(self.db_cluster.nodes)

        self.actions_log.info("Preparing test keyspaces and tables")
        # prepare test keyspaces and tables before upgrade to avoid schema change during mixed cluster.
        self.prepare_keyspaces_and_tables()
        self.actions_log.info("Running s-b to create schemas to avoid #11459")
        # Use large_partition_stress_during_upgrade for schema creation if defined
        large_partition_stress_during_upgrade = self.params.get("large_partition_stress_during_upgrade")
        if large_partition_stress_during_upgrade:
            sb_create_schema_cmds = [f"{cmd} -duration=1m" for cmd in large_partition_stress_during_upgrade]
            sb_create_schemas = self._run_all_stress_cmds([], params={"stress_cmd": sb_create_schema_cmds})
            for sb_create_schema in sb_create_schemas:
                self.verify_stress_thread(sb_create_schema)
        self.fill_and_verify_db_data("BEFORE UPGRADE", pre_fill=True)

        # write workload during entire test
        self.actions_log.info("Starting cassandra-stress write workload for entire test")
        write_stress_during_entire_test = self.params.get("write_stress_during_entire_test")
        entire_write_cs_thread_pools = self._run_all_stress_cmds(
            [], params={"stress_cmd": write_stress_during_entire_test}
        )

        # Let to write_stress_during_entire_test complete the schema changes
        self.metric_has_data(
            metric_query='sct_cassandra_stress_write_gauge{type="ops", keyspace="keyspace_entire_test"}'
            'or sct_cql_stress_cassandra_stress_write_gauge{type="ops", keyspace="keyspace_entire_test"}',
            n=10,
        )

        if self.do_truncates:
            # Prepare keyspace and tables for truncate test
            self.fill_db_data_for_truncate_test(insert_rows=NUMBER_OF_ROWS_FOR_TRUNCATE_TEST)

        # generate random order to upgrade
        nodes_num = len(self.db_cluster.nodes)
        # prepare an array containing the indexes
        indexes = list(range(nodes_num))
        # shuffle it so we will upgrade the nodes in a random order
        random.shuffle(indexes)

        if self.should_do_complex_profile():
            keyspace = "keyspace_complex"
            table = "user_with_ck"
        else:
            keyspace = "keyspace_entire_test"
            table = "standard1"

        self.actions_log.info("Running stress workload before upgrade")
        if self.should_do_complex_profile():
            # complex workload: prepare write
            self.actions_log.info("Starting complex c-s workload (5M) to prepare data")
            stress_cmd_complex_prepare = self.params.get("stress_cmd_complex_prepare")
            complex_cs_thread_pools = self._run_all_stress_cmds([], params={"stress_cmd": stress_cmd_complex_prepare})

            # wait for the complex workload to finish
            for pool in complex_cs_thread_pools:
                self.verify_stress_thread(pool)

        self.actions_log.info("Checking paged query before upgrading nodes")
        self.paged_query(keyspace=keyspace)
        self.actions_log.info("Finished checking paged query before upgrading nodes")

        # prepare write workload
        self.actions_log.info("Starting prepare write workload (n=10000000)")
        prepare_write_stress = self.params.get("prepare_write_stress")
        prepare_write_cs_thread_pools = self._run_all_stress_cmds([], params={"stress_cmd": prepare_write_stress})
        self.actions_log.info("Waiting for cassandra-stress to start before upgrade")
        self.metric_has_data(
            metric_query='sct_cassandra_stress_write_gauge{type="ops", keyspace="keyspace1"}'
            'or sct_cql_stress_cassandra_stress_write_gauge{type="ops", keyspace="keyspace1"}',
            n=5,
        )

        gemini_thread = None
        if self.params.get("run_gemini_in_rolling_upgrade"):
            self.actions_log.info("Starting gemini during upgrade")
            gemini_cmd = self.params.get("gemini_cmd")
            if self.enable_cdc_for_tables:
                gemini_cmd += " --table-options \"cdc={'enabled': true}\""
            gemini_thread = self.run_gemini(gemini_cmd)
            self.metric_has_data(metric_query="sum(increase(gemini_cql_requests[1m]))", n=10)
        else:
            self.actions_log.info("Gemini workload is disabled for this rolling upgrade test")

        with ignore_upgrade_schema_errors():
            step = "Step1 - Upgrade First Node "
            self.actions_log.info(step)
            # upgrade first node
            self.db_cluster.node_to_upgrade = self.db_cluster.nodes[indexes[0]]
            self.upgrade_node(self.db_cluster.node_to_upgrade)
            self.db_cluster.node_to_upgrade.check_node_health()

            # wait for the prepare write workload to finish
            for pool in prepare_write_cs_thread_pools:
                self.verify_stress_thread(pool)

            # read workload (cl=QUORUM)
            self.actions_log.info("Starting read workload (cl=QUORUM n=10000000)")
            stress_cmd_read_cl_quorum = self.params.get("stress_cmd_read_cl_quorum")
            read_stress_queues = self._run_all_stress_cmds([], params={"stress_cmd": stress_cmd_read_cl_quorum})
            # wait for the read workload to finish
            for queue in read_stress_queues:
                self.verify_stress_thread(queue)
            self.actions_log.info("Completed first node upgrade")
            self.search_for_idx_token_error_after_upgrade(
                node=self.db_cluster.node_to_upgrade, step=step + " - after upgraded one node"
            )

            # read workload
            self.actions_log.info("Starting read workload for 10m")
            stress_cmd_read_10m = self.params.get("stress_cmd_read_10m")
            read_10m_cs_thread_pools = self._run_all_stress_cmds([], params={"stress_cmd": stress_cmd_read_10m})

            self.actions_log.info("Running stress-bench large partitions workload during upgrade")
            large_partition_stress_during_upgrade = self.params.get("large_partition_stress_during_upgrade")
            if large_partition_stress_during_upgrade:
                self._run_all_stress_cmds([], params={"stress_cmd": large_partition_stress_during_upgrade})

            self.actions_log.info("Waiting 60s for workloads to start before upgrade")
            time.sleep(60)

            step = "Step2 - Upgrade Second Node "
            self.actions_log.info(step)
            # upgrade second node
            self.db_cluster.node_to_upgrade = self.db_cluster.nodes[indexes[1]]
            self.upgrade_node(self.db_cluster.node_to_upgrade)
            self.db_cluster.node_to_upgrade.check_node_health()

            # wait for the 10m read workload to finish
            for pool in read_10m_cs_thread_pools:
                self.verify_stress_thread(pool)
            self.fill_and_verify_db_data("after upgraded two nodes")
            self.search_for_idx_token_error_after_upgrade(
                node=self.db_cluster.node_to_upgrade, step=step + " - after upgraded two nodes"
            )

            # read workload (60m)
            self.actions_log.info("Starting read workload for 60m")
            stress_cmd_read_60m = self.params.get("stress_cmd_read_60m")
            read_60m_cs_thread_pools = self._run_all_stress_cmds([], params={"stress_cmd": stress_cmd_read_60m})
            self.actions_log.info("Waiting 60s for cassandra-stress to start before rollback")
            time.sleep(60)

            self.actions_log.info("Step3 - Rollback Second Node")
            # rollback second node
            self.rollback_node(self.db_cluster.nodes[indexes[1]])
            self.db_cluster.nodes[indexes[1]].check_node_health()

        step = "Step4 - Verify data during mixed cluster mode "
        self.actions_log.info(step)
        self.fill_and_verify_db_data("after rollback the second node")

        self.actions_log.info("Repair the first upgraded Node")
        self.db_cluster.nodes[indexes[0]].run_nodetool(sub_cmd="repair", timeout=7200, coredump_on_timeout=True)
        self.search_for_idx_token_error_after_upgrade(node=self.db_cluster.node_to_upgrade, step=step)

        with ignore_upgrade_schema_errors():
            step = "Step5 - Upgrade rest of the Nodes "
            self.actions_log.info(step)
            for i in indexes[1:]:
                self.db_cluster.node_to_upgrade = self.db_cluster.nodes[i]
                self.upgrade_node(self.db_cluster.node_to_upgrade)
                self.db_cluster.node_to_upgrade.check_node_health()
                self.fill_and_verify_db_data("after upgraded %s" % self.db_cluster.node_to_upgrade.name)
                self.search_for_idx_token_error_after_upgrade(node=self.db_cluster.node_to_upgrade, step=step)
        self.actions_log.info("Step5.1 - run raft topology upgrade procedure")
        self.run_raft_topology_upgrade_procedure()
        InfoEvent(message="Step5.2 - check limited voters feature").publish()
        self.validate_limited_voters_feature_enabled()

        self.actions_log.info("Step6 - Verify stress results after upgrade ")
        self.actions_log.info("Waiting for stress threads to complete after upgrade")
        # wait for the 60m read workload to finish
        for pool in read_60m_cs_thread_pools:
            self.verify_stress_thread(pool)

        for pool in entire_write_cs_thread_pools:
            self.verify_stress_thread(pool)

        self.actions_log.info("Step7 - Upgrade sstables to latest supported version ")
        # figure out what is the last supported sstable version
        self.expected_sstable_format_version = self.get_highest_supported_sstable_version()

        # run 'nodetool upgradesstables' on all nodes and check/wait for all file to be upgraded
        upgradesstables = self.db_cluster.run_func_parallel(func=self.upgradesstables_if_command_available)

        # only check sstable format version if all nodes had 'nodetool upgradesstables' available
        if all(upgradesstables):
            self.actions_log.info("Upgrading sstables if new version is available")
            tables_upgraded = self.db_cluster.run_func_parallel(func=self.wait_for_sstable_upgrade)
            assert all(tables_upgraded), f"Failed to upgrade the sstable format {tables_upgraded}"

        # Verify sstabledump / scylla sstable dump-data
        self.actions_log.info("Starting sstabledump to verify correctness of sstables")
        first_node = self.db_cluster.nodes[0]

        dump_cmd = get_sstable_data_dump_command(first_node, keyspace, table)
        first_node.remoter.run(
            f"for i in `sudo find /var/lib/scylla/data/{keyspace}/ -type f |grep -v manifest.json |"
            "grep -v snapshots |grep -v '\\.tmp$' |head -n 1`; do echo $i; "
            f"sudo {dump_cmd} $i 1>/tmp/sstabledump.output || "
            "exit 1; done",
            verbose=True,
        )

        self.actions_log.info("Step8 - Run stress and verify after upgrading entire cluster")
        self.actions_log.info("Starting verification stresses after cluster upgrade")
        stress_after_cluster_upgrade = self.params.get("stress_after_cluster_upgrade")
        self._run_all_stress_cmds([], params={"stress_cmd": stress_after_cluster_upgrade})
        verify_stress_after_cluster_upgrade = self.params.get("verify_stress_after_cluster_upgrade")
        verify_stress_cs_thread_pools = self._run_all_stress_cmds(
            [], params={"stress_cmd": verify_stress_after_cluster_upgrade}
        )
        for pool in verify_stress_cs_thread_pools:
            self.verify_stress_thread(pool)

        if self.should_do_complex_profile():
            # complex workload: verify data by simple read cl=ALL
            self.actions_log.info("Starting c-s complex workload to verify data by simple read")
            stress_cmd_complex_verify_read = self.params.get("stress_cmd_complex_verify_read")
            complex_cs_thread_pools = self._run_all_stress_cmds(
                [], params={"stress_cmd": stress_cmd_complex_verify_read}
            )
            # wait for the read complex workload to finish
            for pool in complex_cs_thread_pools:
                self.verify_stress_thread(pool)

        self.actions_log.info("Will check paged query after upgrading all nodes")
        self.paged_query(keyspace=keyspace)
        self.actions_log.info("Done checking paged query after upgrading nodes")

        # After adjusted the workloads, there is a entire write workload, and it uses a fixed duration for catching
        # the data lose.
        # But the execute time of workloads are not exact, so let only use basic prepare write & read verify for
        # complex workloads,and comment two complex workloads.
        #
        # TODO: retest commented workloads and decide to enable or delete them.
        # if self.should_do_complex_profile():
        #   # complex workload: verify data by multiple ops
        #   self.log.info('Starting c-s complex workload to verify data by multiple ops')
        #   stress_cmd_complex_verify_more = self.params.get('stress_cmd_complex_verify_more')
        #   stress_cmd_complex_verify_more = self.replace_complex_profile_if_needed(stress_cmd_complex_verify_more)
        #   complex_cs_thread_pool = self.run_stress_thread(stress_cmd=stress_cmd_complex_verify_more)

        #   # wait for the complex workload to finish
        #   self.verify_stress_thread(complex_cs_thread_pool)

        #   # complex workload: verify data by delete 1/10 data
        #   self.log.info('Starting c-s complex workload to verify data by delete')
        #   stress_cmd_complex_verify_delete = self.params.get('stress_cmd_complex_verify_delete')
        #   stress_cmd_complex_verify_delete = self.replace_complex_profile_if_needed(stress_cmd_complex_verify_delete)
        #   complex_cs_thread_pool = self.run_stress_thread(stress_cmd=stress_cmd_complex_verify_delete)
        #   # wait for the complex workload to finish
        #   self.verify_stress_thread(complex_cs_thread_pool)

        # During the test we filter and ignore some specific errors, but we want to allow only certain amount of them
        step = "Step9 - Search for errors that we filter during the test "
        self.actions_log.info(step)
        self.actions_log.info("Checking how many failed_to_load_schem errors happened during the test")
        error_factor = 3
        schema_load_error_num = self.count_log_errors(search_pattern="Failed to load schema version", step=step)
        # Warning example:
        # workload prioritization - update_service_levels_from_distributed_data: an error occurred while retrieving
        # configuration (exceptions::read_failure_exception (Operation failed for system_distributed.service_levels
        # - received 0 responses and 1 failures from 1 CL=ONE.))
        workload_prioritization_error_num = self.count_log_errors(
            search_pattern="workload prioritization.*read_failure_exception",
            step=step,
            search_for_idx_token_error=False,
        )

        self.actions_log.info(
            "schema_load_error_num: %s; workload_prioritization_error_num: %s"
            % (schema_load_error_num, workload_prioritization_error_num)
        )

        # Issue #https://github.com/scylladb/scylla-enterprise/issues/1391
        # By Eliran's comment: For 'Failed to load schema version' error which is expected and non offensive is
        # to count the 'workload prioritization' warning and subtract that amount from the amount of overall errors.
        load_error_num = schema_load_error_num - workload_prioritization_error_num
        assert load_error_num <= error_factor * 8 * len(self.db_cluster.nodes), (
            "Only allowing shards_num * %d schema load errors per host during the "
            "entire test, actual: %d" % (error_factor, schema_load_error_num)
        )

        if gemini_thread:
            self.actions_log.info("Step10 - Verify that gemini did not failed during upgrade")
            self.verify_gemini_results(queue=gemini_thread)
        else:
            self.actions_log.info("Step10 - Skipping Gemini verification as Gemini was not run during this test")

        self.actions_log.info("all nodes were upgraded, and last workaround is verified.")

    def run_raft_topology_upgrade_procedure(self):
        # wait features is enabled on nodes after upgrade
        feature_state_per_node = []
        for node in self.db_cluster.nodes:
            result = wait_for(
                func=self.db_cluster.is_features_enabled_on_node,
                timeout=60,
                step=1,
                text=f"Check feature enabled on node {node.name}",
                throw_exc=False,
                feature_list=[CONSISTENT_TOPOLOGY_CHANGES_FEATURE],
                node=node,
            )
            feature_state_per_node.append(result)
        if not all(feature_state_per_node):
            self.actions_log.info(
                "Step5.1 - Consistent topology changes is not supported. Cluster stay with gossip topology mode"
            )
            return
        raft_upgrade = RaftUpgradeProcedure(self.db_cluster.nodes[0])
        raft_upgrade.start_upgrade_procedure()
        for node in self.db_cluster.nodes:
            RaftUpgradeProcedure(node).wait_upgrade_procedure_done()
        self.actions_log.info("Step5.1 - raft topology upgrade procedure done")

    def _start_and_wait_for_node_upgrade(self, node: BaseNode, step: int) -> None:
        self.actions_log.info(f"Step {step} - Upgrade {node.name} node in {node.dc_idx} dc")
        self.upgrade_node(node, upgrade_sstables=self.params.get("upgrade_sstables"))
        node.check_node_health()

    def _start_and_wait_for_node_rollback(self, node: BaseNode, step: int) -> None:
        self.actions_log.info(f"Step {step} - Rollback {node.name} node in {node.dc_idx} dc")
        self.rollback_node(node, upgrade_sstables=self.params.get("upgrade_sstables"))
        node.check_node_health()

    def _run_stress_workload(
        self,
        workload_name: str,
        wait_for_finish: bool = False,
        round_robin: bool = False,
        reset_loader_cycle: bool = False,
    ) -> [CassandraStressThread]:
        """Runs workload from param name specified in test-case yaml"""
        if reset_loader_cycle:
            # NOTE: it will make stress commands be scheduled onto loaders starting from the first one
            #       not keeping state from previous test stages.
            #       Useful for running DC-specific stress commands from proper loaders in proper DCs.
            self.loaders._loader_cycle = None
        self.actions_log.info(f"Starting {workload_name}")
        stress_commands = self.params.get(workload_name)
        workload_thread_pools = []
        if isinstance(stress_commands, str):
            stress_commands = [stress_commands]

        for stress_command in stress_commands:
            workload_thread_pools.append(self.run_stress_thread(stress_cmd=stress_command, round_robin=round_robin))

        if self.params.get("alternator_port"):
            self.pre_create_alternator_tables()
        if wait_for_finish is True:
            self.actions_log.info(f"Waiting for {workload_name} to finish")
            for thread_pool in workload_thread_pools:
                self.verify_stress_thread(thread_pool)
        else:
            self.actions_log.info("Sleeping for 60s to let the stress command(s) start before the next steps...")
            time.sleep(60)
        return workload_thread_pools

    def _add_sla_credentials_to_stress_commands(self, workloads_with_sla: list):
        if not workloads_with_sla:
            self.log.debug("No workloads supplied")
            return

        stress_during_entire_upgrade = self.params.get(workloads_with_sla[0])
        if not [cmd for cmd in stress_during_entire_upgrade if "<sla credentials " in cmd]:
            self.log.debug("No need to set SLA credentials for stress command: %s", stress_during_entire_upgrade)
            return

        roles = []
        service_level_shares = self.params.get("service_level_shares")
        # OSS Scylla version does not support "shares" option for Service Level
        if service_level_shares and not self.db_cluster.nodes[0].is_enterprise:
            service_level_shares = [None for _ in service_level_shares]

        with self.db_cluster.cql_connection_patient(
            node=self.db_cluster.nodes[0], user=loader_utils.DEFAULT_USER, password=loader_utils.DEFAULT_USER_PASSWORD
        ) as session:
            for index, shares in enumerate(service_level_shares):
                roles.append(create_sla_auth(session=session, shares=shares, index=str(index)))

        self.add_sla_credentials_to_stress_cmds(
            workload_names=workloads_with_sla,
            roles=roles,
            params=self.params,
            parent_class_name=self.__class__.__name__,
        )

    def test_generic_cluster_upgrade(self):
        """
        Upgrade specified number of nodes in the cluster, and start special read workload
        during the stage. Checksum method is changed to xxhash from Scylla 2.2,
        we want to use this case to verify the read (cl=ALL) workload works
        well, upgrade all nodes to new version in the end.

        For multi-dc upgrades, alternates upgraded nodes between dc's.
        """
        self.upgrade_os(self.db_cluster.nodes)

        if self.do_truncates:
            # Prepare keyspace and tables for truncate test
            self.fill_db_data_for_truncate_test(insert_rows=NUMBER_OF_ROWS_FOR_TRUNCATE_TEST)

        self._add_sla_credentials_to_stress_commands(
            workloads_with_sla=["stress_during_entire_upgrade", "stress_after_cluster_upgrade"]
        )
        step = itertools_count(start=1)
        self._run_stress_workload("stress_before_upgrade", wait_for_finish=True)
        stress_thread_pools = self._run_stress_workload("stress_during_entire_upgrade", wait_for_finish=False)

        # generate random order to upgrade
        nodes_to_upgrade = self.shuffle_nodes_and_alternate_dcs(list(self.db_cluster.nodes))

        # Upgrade all nodes that should be rollback later
        upgraded_nodes = []
        for node_to_upgrade in nodes_to_upgrade[: self.params.get("num_nodes_to_rollback")]:
            self._start_and_wait_for_node_upgrade(node_to_upgrade, step=next(step))
            upgraded_nodes.append(node_to_upgrade)

        # Rollback all nodes that where upgraded (not necessarily in the same order)
        random.shuffle(upgraded_nodes)
        self.actions_log.info("Rollback nodes")
        for node in upgraded_nodes:
            self._start_and_wait_for_node_rollback(node, step=next(step))

        # Upgrade all nodes
        self.actions_log.info("Upgrade nodes")
        for node_to_upgrade in nodes_to_upgrade:
            self._start_and_wait_for_node_upgrade(node_to_upgrade, step=next(step))
        self.actions_log.info("All nodes were upgraded successfully")

        self.actions_log.info("Run raft topology upgrade procedure")
        self.run_raft_topology_upgrade_procedure()
        self.validate_limited_voters_feature_enabled()

        self.actions_log.info("Waiting for stress_during_entire_upgrade to finish")
        for stress_thread_pool in stress_thread_pools:
            self.verify_stress_thread(stress_thread_pool)

        self._run_stress_workload("stress_after_cluster_upgrade", wait_for_finish=True)

    def test_cluster_upgrade_latency_regression(self):
        """Check latency regression after a ScyllaDB cluster upgrade using latte stress commands.

        Number of 'before' and 'after' commands must match. Their latency values will be compaired.

        - Write initial latte data (prepare_write_cmd)
        - Wait for end of compactions
        - Read latte data generating report file (stress_before_upgrade)
        - Write latency results from the 'stress_before_upgrade' to Argus
        - Run a read latte stress (stress_during_entire_upgrade) not waiting for it's end
        - Upgrade the DB cluster
        * self.run_raft_topology_upgrade_procedure()
        - Wait for the end of the stress command (stress_during_entire_upgrade)
        - Wait for end of compactions
        - Read latte data (stress_after_cluster_upgrade) generating report file
        - Write latency results from the 'stress_after_cluster_upgrade' to Argus
        """
        self.upgrade_os(self.db_cluster.nodes)

        InfoEvent(message="Step1 - Populate DB data").publish()
        if self.do_truncates:
            # Prepare keyspace and tables for truncate test
            self.fill_db_data_for_truncate_test(insert_rows=NUMBER_OF_ROWS_FOR_TRUNCATE_TEST)
        self.run_prepare_write_cmd()

        InfoEvent(message="Step2 - Run 'read' command before upgrade").publish()
        step = itertools_count(start=1)
        stress_before_upgrade_thread_pools = self._run_stress_workload(
            "stress_before_upgrade",
            wait_for_finish=False,
            round_robin=True,
            reset_loader_cycle=True,
        )
        stress_before_upgrade_results = []
        for stress_before_upgrade_thread_pool in stress_before_upgrade_thread_pools:
            stress_before_upgrade_results.append(self.get_stress_results(stress_before_upgrade_thread_pool))
        self.log.info("Stress results before upgrade: %s", stress_before_upgrade_results)

        result_table = argus_results.LatteStressLatencyComparison()
        # NOTE: write 'before' results to Argus
        for i in range(len(stress_before_upgrade_results)):
            row = f"#{i + 1}"
            result_table.add_result(
                column="before_ops",
                row=row,
                status=argus_results.Status.UNSET,
                value=int(stress_before_upgrade_results[i][0]["op rate"]),
            )
            result_table.add_result(
                column="before_mean",
                row=row,
                status=argus_results.Status.UNSET,
                value=float(stress_before_upgrade_results[i][0]["latency mean"]),
            )
            result_table.add_result(
                column="before_p99",
                row=row,
                status=argus_results.Status.UNSET,
                value=float(stress_before_upgrade_results[i][0]["latency 99th percentile"]),
            )
        argus_results.submit_results_to_argus(argus_client=self.test_config.argus_client(), result_table=result_table)

        stress_during_entire_upgrade_thread_pools = self._run_stress_workload(
            "stress_during_entire_upgrade",
            wait_for_finish=False,
            round_robin=True,
            reset_loader_cycle=True,
        )

        InfoEvent(message="Step3 - Upgrade cluster to '%s' version" % self.params.get("new_version")).publish()

        InfoEvent(message="Upgrade part of nodes nodes before roll-back").publish()
        nodes_to_upgrade = self.shuffle_nodes_and_alternate_dcs(list(self.db_cluster.nodes))
        upgraded_nodes = []
        for node_to_upgrade in nodes_to_upgrade[: self.params.get("num_nodes_to_rollback")]:
            self._start_and_wait_for_node_upgrade(node_to_upgrade, step=next(step))
            upgraded_nodes.append(node_to_upgrade)

        # Rollback all nodes that where upgraded (not necessarily in the same order)
        random.shuffle(upgraded_nodes)
        InfoEvent(message="Roll-back following nodes: %s" % ", ".join(node.name for node in upgraded_nodes)).publish()
        for node in upgraded_nodes:
            self._start_and_wait_for_node_rollback(node, step=next(step))

        InfoEvent(message="Upgrade all nodes").publish()
        for node_to_upgrade in nodes_to_upgrade:
            self._start_and_wait_for_node_upgrade(node_to_upgrade, step=next(step))
        InfoEvent(message="All nodes were upgraded successfully").publish()

        InfoEvent(message="Step4 - Run raft topology upgrade procedure").publish()
        self.run_raft_topology_upgrade_procedure()
        self.validate_limited_voters_feature_enabled()

        InfoEvent(message="Step5 - Wait for stress_during_entire_upgrade to finish").publish()
        for stress_during_entire_upgrade_thread_pool in stress_during_entire_upgrade_thread_pools:
            self.verify_stress_thread(stress_during_entire_upgrade_thread_pool)
        self.wait_no_compactions_running(n=240, sleep_time=30)

        InfoEvent(message="Step6 - run 'stress_after_cluster_upgrade' stress command(s)").publish()
        time.sleep(60)
        stress_after_upgrade_thread_pools = self._run_stress_workload(
            "stress_after_cluster_upgrade",
            wait_for_finish=False,
            round_robin=True,
            reset_loader_cycle=True,
        )
        stress_after_upgrade_results = []
        for stress_after_upgrade_thread_pool in stress_after_upgrade_thread_pools:
            stress_after_upgrade_results.append(self.get_stress_results(stress_after_upgrade_thread_pool))
        self.log.info("Stress results after upgrade: %s", stress_after_upgrade_results)

        # NOTE: write 'after' results to Argus
        for i in range(len(stress_after_upgrade_results)):
            row = f"#{i + 1}"
            result_table.add_result(
                column="after_p99",
                row=row,
                status=argus_results.Status.UNSET,
                value=float(stress_after_upgrade_results[i][0]["latency 99th percentile"]),
            )
            result_table.add_result(
                column="after_mean",
                row=row,
                status=argus_results.Status.UNSET,
                value=float(stress_after_upgrade_results[i][0]["latency mean"]),
            )
            result_table.add_result(
                column="after_ops",
                row=row,
                status=argus_results.Status.UNSET,
                value=int(stress_after_upgrade_results[i][0]["op rate"]),
            )
        argus_results.submit_results_to_argus(argus_client=self.test_config.argus_client(), result_table=result_table)

        assert len(stress_before_upgrade_results) > 0
        for stress_before_upgrade_result in stress_before_upgrade_results:
            assert len(stress_before_upgrade_result) > 0
        assert len(stress_before_upgrade_results) == len(stress_after_upgrade_results)
        for stress_after_upgrade_result in stress_after_upgrade_results:
            assert len(stress_after_upgrade_result) > 0
        for i in range(len(stress_before_upgrade_results)):
            for j in range(len(stress_before_upgrade_results[i])):
                assert "latency 99th percentile" in stress_before_upgrade_results[i][j]
                current_latency_before = float(stress_before_upgrade_results[i][j]["latency 99th percentile"])
                assert current_latency_before > 0
                assert "latency 99th percentile" in stress_after_upgrade_results[i][j]
                current_latency_after = float(stress_after_upgrade_results[i][j]["latency 99th percentile"])
                assert current_latency_after > 0

    def test_kubernetes_scylla_upgrade(self):
        """
        Run a set of different cql queries against various types/tables before
        and after upgrade of every node to check the consistency of data
        """
        self.k8s_cluster.check_scylla_cluster_sa_annotations()
        InfoEvent(message="Step1 - Populate DB with many types of tables and data").publish()
        target_upgrade_version = self.params.get("new_version")
        self.prepare_keyspaces_and_tables()
        InfoEvent(message="Step2 - Populate some data before upgrading cluster").publish()
        self.fill_and_verify_db_data("", pre_fill=True)
        InfoEvent(message="Step3 - Starting c-s write workload").publish()
        for pool in self._run_all_stress_cmds([], params={"stress_cmd": self.params.get("stress_cmd_w")}):
            self.verify_stress_thread(pool)
        InfoEvent(message="Step4 - Starting c-s read workload").publish()
        for pool in self._run_all_stress_cmds([], params={"stress_cmd": self.params.get("stress_cmd_r")}):
            self.verify_stress_thread(pool)
        InfoEvent(message="Step5 - Upgrade cluster to %s" % target_upgrade_version).publish()
        self.db_cluster.upgrade_scylla_cluster(target_upgrade_version)
        InfoEvent(message="Step6 - Wait till cluster got upgraded").publish()
        self.wait_till_scylla_is_upgraded_on_all_nodes(target_upgrade_version)
        self.k8s_cluster.check_scylla_cluster_sa_annotations()

        InfoEvent(message="Step7 - Upgrade sstables").publish()
        if self.params.get("upgrade_sstables"):
            self.expected_sstable_format_version = self.get_highest_supported_sstable_version()
            upgradesstables = self.db_cluster.run_func_parallel(func=self.upgradesstables_if_command_available)
            # only check sstable format version if all nodes had 'nodetool upgradesstables' available
            if all(upgradesstables):
                InfoEvent(message="Waiting until jmx is up across the board").publish()
                self.wait_till_jmx_on_all_nodes()
                InfoEvent(message="Upgrading sstables if new version is available").publish()
                tables_upgraded = self.db_cluster.run_func_parallel(func=self.wait_for_sstable_upgrade)
                assert all(tables_upgraded), f"Failed to upgrade the sstable format {tables_upgraded}"
        else:
            InfoEvent(message="Upgrade of sstables is disabled").publish()

        InfoEvent(message="Step8 - Verify data after upgrade").publish()
        self.fill_and_verify_db_data(note="after all nodes upgraded")
        InfoEvent(message="Step9 - Starting c-s read workload").publish()
        for pool in self._run_all_stress_cmds([], params={"stress_cmd": self.params.get("stress_cmd_r")}):
            self.verify_stress_thread(pool)
        InfoEvent(message="Step10 - Starting c-s write workload").publish()
        for pool in self._run_all_stress_cmds([], params={"stress_cmd": self.params.get("stress_cmd_w")}):
            self.verify_stress_thread(pool)
        InfoEvent(message="Step11 - Starting c-s read workload").publish()
        for pool in self._run_all_stress_cmds([], params={"stress_cmd": self.params.get("stress_cmd_r")}):
            self.verify_stress_thread(pool)
        InfoEvent(message="Step12 - Search for errors in scylla log").publish()
        for node in self.db_cluster.nodes:
            self.search_for_idx_token_error_after_upgrade(node=node, step=f"{node!s} after upgrade")
        InfoEvent(message="Step13 - Checking how many failed_to_load_scheme errors happened during the test").publish()
        error_factor = 3
        schema_load_error_num = self.count_log_errors(
            search_pattern="Failed to load schema version", step="AFTER UPGRADE"
        )
        # Warning example:
        # workload prioritization - update_service_levels_from_distributed_data: an error occurred while retrieving
        # configuration (exceptions::read_failure_exception (Operation failed for system_distributed.service_levels
        # - received 0 responses and 1 failures from 1 CL=ONE.))
        workload_prioritization_error_num = self.count_log_errors(
            search_pattern="workload prioritization.*read_failure_exception",
            step="AFTER UPGRADE",
            search_for_idx_token_error=False,
        )
        InfoEvent(
            message="schema_load_error_num: %s; workload_prioritization_error_num: %s"
            % (schema_load_error_num, workload_prioritization_error_num)
        ).publish()

        # Issue #https://github.com/scylladb/scylla-enterprise/issues/1391
        # By Eliran's comment: For 'Failed to load schema version' error which is expected and non offensive is
        # to count the 'workload prioritization' warning and subtract that amount from the amount of overall errors.
        load_error_num = schema_load_error_num - workload_prioritization_error_num
        assert load_error_num <= error_factor * 8 * len(self.db_cluster.nodes), (
            "Only allowing shards_num * %d schema load errors per host during the "
            "entire test, actual: %d" % (error_factor, schema_load_error_num)
        )

    def _get_current_operator_image_tag(self):
        return (
            self.k8s_cluster.kubectl(
                "get deployment scylla-operator -o custom-columns=:..image --no-headers",
                namespace=self.k8s_cluster._scylla_operator_namespace,
            )
            .stdout.strip()
            .split(":")[-1]
        )

    def test_kubernetes_operator_upgrade(self):
        self.k8s_cluster.check_scylla_cluster_sa_annotations()

        InfoEvent(message="Step1 - Populate DB with data").publish()
        self.prepare_keyspaces_and_tables()
        self.fill_and_verify_db_data("", pre_fill=True)

        InfoEvent(message="Step2 - Run c-s write workload").publish()
        for pool in self._run_all_stress_cmds([], params={"stress_cmd": self.params.get("stress_cmd_w")}):
            self.verify_stress_thread(pool)

        InfoEvent(message="Step3 - Run c-s read workload").publish()
        for pool in self._run_all_stress_cmds([], params={"stress_cmd": self.params.get("stress_cmd_r")}):
            self.verify_stress_thread(pool)

        InfoEvent(message="Step4 - Upgrade scylla-operator").publish()
        old_scylla_pods_uids = {}
        for scylla_pod in self.db_cluster.nodes:
            old_scylla_pods_uids[scylla_pod.name] = scylla_pod.k8s_pod_uid
        progressing_last_transition_time_before_upgrade = self.k8s_cluster.kubectl(
            f"get --no-headers scyllacluster {self.params.get('k8s_scylla_cluster_name')}"
            " -o jsonpath='{.status.conditions[?(@.type==\"Progressing\")].lastTransitionTime}'",
            namespace=self.db_cluster.namespace,
        ).stdout.strip()
        base_docker_image_tag = self._get_current_operator_image_tag()
        upgrade_docker_image = self.params.get("k8s_scylla_operator_upgrade_docker_image") or ""
        self.k8s_cluster.upgrade_scylla_operator(
            self.params.get("k8s_scylla_operator_upgrade_helm_repo")
            or self.params.get("k8s_scylla_operator_helm_repo"),
            self.params.get("k8s_scylla_operator_upgrade_chart_version") or "latest",
            upgrade_docker_image,
        )

        InfoEvent(message="Step5 - Validate scylla-operator version after upgrade").publish()
        actual_docker_image_tag = self._get_current_operator_image_tag()
        self.assertNotEqual(base_docker_image_tag, actual_docker_image_tag)
        expected_docker_image_tag = upgrade_docker_image.split(":")[-1]
        if not expected_docker_image_tag:
            operator_chart_info = self.k8s_cluster.helm(f"ls -n {self.k8s_cluster._scylla_operator_namespace} -o json")
            expected_docker_image_tag = json.loads(operator_chart_info)[0]["app_version"]
        self.assertEqual(expected_docker_image_tag, actual_docker_image_tag)

        InfoEvent(message="Step6 - Wait for the update of Scylla cluster").publish()
        for scylla_pod in self.db_cluster.nodes[::-1]:
            scylla_pod.wait_till_k8s_pod_get_uid(
                ignore_uid=old_scylla_pods_uids[scylla_pod.name], throw_exc=True, timeout=900
            )
        for scylla_pod in self.db_cluster.nodes[::-1]:
            scylla_pod.wait_for_pod_readiness()

        for condition in ("Available=True", "Progressing=False", "Degraded=False"):
            self.k8s_cluster.kubectl_wait(
                f"scyllacluster {self.params.get('k8s_scylla_cluster_name')} --for=condition={condition}",
                namespace=self.db_cluster.namespace,
                timeout=600,
            )
        scylla_cluster_conditions = json.loads(
            self.k8s_cluster.kubectl(
                f"get scyllacluster {self.params.get('k8s_scylla_cluster_name')} -o json",
                namespace=self.db_cluster.namespace,
            ).stdout.strip()
        )["status"]["conditions"]
        for condition in scylla_cluster_conditions:
            if condition["type"] == "Progressing":
                assert condition["lastTransitionTime"] > progressing_last_transition_time_before_upgrade
            assert condition["reason"] != "Error", "'ScyllaCluster' failed to be updated: %s" % condition

        pods_with_wrong_image_tags = []
        pods = json.loads(
            self.k8s_cluster.kubectl("get pods -o json", namespace=self.db_cluster.namespace).stdout.strip()
        )["items"]
        for pod in pods:
            # NOTE: in pods we have 'containers' and 'initContainers'. So, walk over the both types
            #       and check that the 'scylla-operator' images are updated.
            for container_type in ("c", "initC"):
                for container in pod.get("status", {}).get(f"{container_type}ontainerStatuses", []):
                    image = container["image"].split("/")[-1]
                    image_name, image_tag = image.split(":")
                    if image_name == "scylla-operator" and image_tag != expected_docker_image_tag:
                        pods_with_wrong_image_tags.append(
                            {
                                "namespace": self.db_cluster.namespace,
                                "pod_name": pod["metadata"]["name"],
                                "container_name": container["name"],
                                "image": image,
                            }
                        )
        assert not pods_with_wrong_image_tags, (
            f"Found pods that have unexpected scylla-operator image tags.\n"
            f"Expected is '{expected_docker_image_tag}'.\n"
            f"Pods: {json.dumps(pods_with_wrong_image_tags, indent=2)}"
        )

        self.k8s_cluster.check_scylla_cluster_sa_annotations()

        InfoEvent(message="Step7 - Add new member to the Scylla cluster").publish()
        peer_db_node = self.db_cluster.nodes[0]
        new_nodes = self.db_cluster.add_nodes(
            count=1, dc_idx=peer_db_node.dc_idx, rack=peer_db_node.rack, enable_auto_bootstrap=True
        )
        self.db_cluster.wait_for_init(node_list=new_nodes, timeout=40 * 60)
        self.db_cluster.wait_sts_rollout_restart(pods_to_wait=1)
        self.db_cluster.wait_for_nodes_up_and_normal(nodes=new_nodes)
        self.monitors.reconfigure_scylla_monitoring()
        self.k8s_cluster.check_scylla_cluster_sa_annotations()

        InfoEvent(message="Step8 - Verify data in the Scylla cluster").publish()
        self.fill_and_verify_db_data(note="after operator upgrade and scylla member addition")

        InfoEvent(message="Step9 - Run c-s read workload").publish()
        for pool in self._run_all_stress_cmds([], params={"stress_cmd": self.params.get("stress_cmd_r")}):
            self.verify_stress_thread(pool)

    def test_kubernetes_platform_upgrade(self):
        self.k8s_cluster.check_scylla_cluster_sa_annotations()

        InfoEvent(message="Step1 - Populate DB with data").publish()
        self.run_prepare_write_cmd()
        self.prepare_keyspaces_and_tables()
        self.fill_and_verify_db_data("", pre_fill=True)

        InfoEvent(message="Step2 - Start 'stress_during_entire_upgrade'").publish()
        stress_during_entire_upgrade_thread_pools = self._run_stress_workload(
            "stress_during_entire_upgrade", wait_for_finish=False
        )

        InfoEvent(message="Step3 - Upgrade kubernetes platform").publish()
        upgrade_version, scylla_pool = self.k8s_cluster.upgrade_kubernetes_platform(
            pod_objects=self.db_cluster.nodes, use_additional_scylla_nodepool=True
        )

        InfoEvent(message="Step4 - Validate K8S versions after the upgrade").publish()
        data_plane_versions = (
            self.k8s_cluster.kubectl(
                f"get nodes --no-headers -l '{self.k8s_cluster.POOL_LABEL_NAME} "
                f"in ({self.k8s_cluster.AUXILIARY_POOL_NAME},{scylla_pool.name})' "
                "-o custom-columns=:.status.nodeInfo.kubeletVersion"
            )
            .stdout.strip()
            .split()
        )
        # Output example:
        #   GKE                EKS
        #   v1.19.13-gke.700   v1.21.2-eks-c1718fb
        #   v1.19.13-gke.700   v1.21.2-eks-c1718fb
        assert all(ver.startswith(f"v{upgrade_version}.") for ver in data_plane_versions), (
            f"Got unexpected K8S data plane version(s): {data_plane_versions}"
        )

        control_plane_versions = self.k8s_cluster.kubectl("version").stdout.splitlines()
        # Output example:
        # Client Version: v1.20.4
        # Server Version: v1.19.13-gke.700
        # WARNING: foo
        assert any(ver.startswith(f"Server Version: v{upgrade_version}.") for ver in control_plane_versions), (
            f"Got unexpected K8S control plane version: {control_plane_versions}"
        )

        InfoEvent(message="Step5 - Check Scylla-operator pods").publish()
        self.k8s_cluster.kubectl_wait(
            "--all --for=condition=Ready pod", namespace=self.k8s_cluster._scylla_operator_namespace, timeout=600
        )

        InfoEvent(message="Step6 - Check Scylla pods").publish()
        self.k8s_cluster.kubectl_wait(
            "--all --for=condition=Ready pod", namespace=self.db_cluster.namespace, timeout=600
        )
        self.k8s_cluster.check_scylla_cluster_sa_annotations()

        InfoEvent(message="Step7 - Add new member to the Scylla cluster").publish()
        peer_db_node = self.db_cluster.nodes[0]
        new_nodes = self.db_cluster.add_nodes(
            count=1, dc_idx=peer_db_node.dc_idx, rack=peer_db_node.rack, enable_auto_bootstrap=True
        )
        self.db_cluster.wait_for_init(node_list=new_nodes, timeout=40 * 60)
        self.db_cluster.wait_sts_rollout_restart(pods_to_wait=1)
        self.db_cluster.wait_for_nodes_up_and_normal(nodes=new_nodes)
        self.monitors.reconfigure_scylla_monitoring()
        self.k8s_cluster.check_scylla_cluster_sa_annotations()

        InfoEvent(message="Step8 - Wait for the end of 'stress_during_entire_upgrade'").publish()
        for stress_during_entire_upgrade_thread_pool in stress_during_entire_upgrade_thread_pools:
            self.verify_stress_thread(stress_during_entire_upgrade_thread_pool)

        InfoEvent(message="Step9 - Run 'fill_and_verify_db_data'").publish()
        self.fill_and_verify_db_data(note="after operator upgrade and scylla member addition")

        InfoEvent(message="Step10 - Run 'stress_cmd_r' to check all the preloaded data").publish()
        read_stress_threads = self._run_all_stress_cmds(
            [],
            params={
                "stress_cmd": self.params.get("stress_cmd_r"),
                "round_robin": self.params.get("round_robin"),
            },
        )
        for read_stress_thread in read_stress_threads:
            self.verify_stress_thread(read_stress_thread)

    def wait_till_scylla_is_upgraded_on_all_nodes(self, target_version: str) -> None:
        def _is_cluster_upgraded() -> bool:
            for node in self.db_cluster.nodes:
                node.forget_scylla_version()
                if node.scylla_version.replace("~", "-") != target_version.replace("~", "-") or not node.db_up:
                    return False
            return True

        wait.wait_for(
            func=_is_cluster_upgraded,
            step=30,
            text="Waiting until all nodes in the cluster are upgraded",
            timeout=900,
            throw_exc=True,
        )

    def wait_till_jmx_on_all_nodes(self):
        for node in self.db_cluster.nodes:
            node.wait_jmx_up(timeout=300)

    def count_log_errors(self, search_pattern, step, search_for_idx_token_error=True):
        schema_load_error_num = 0
        for node in self.db_cluster.nodes:
            errors = list(node.follow_system_log(patterns=[search_pattern], start_from_beginning=True))
            schema_load_error_num += len(errors)
            if search_for_idx_token_error:
                self.search_for_idx_token_error_after_upgrade(node=node, step=step)
        return schema_load_error_num

    @cache
    def _update_argus_upgraded_version(self, node, new_version):
        try:
            ver = re.search(r"(?P<version>[\w.~]+)-0\.?(?P<date>[0-9]{8,8})\.(?P<commit_id>\w+)", new_version)
            if ver:
                package = Package(
                    name="scylla-server-upgraded",
                    date=ver.group("date"),
                    version=ver.group("version"),
                    revision_id=ver.group("commit_id"),
                    build_id=node.get_scylla_build_id() or "#NO_BUILDID",
                )
                self.log.info("Saving upgraded Scylla version...")
                self.test_config.argus_client().update_scylla_version(version=ver.group("version"))
                self.test_config.argus_client().submit_packages([package])
            else:
                self.log.warning("Couldn't extract version from %s", new_version)
        except Exception as exc:
            self.log.exception("Failed to save upgraded Scylla version in Argus", exc_info=exc)

    def validate_limited_voters_feature_enabled(self):
        InfoEvent(message="Check that limited voters feature enabled after upgrade for all nodes")
        max_voter_number = 5
        feature_state_per_node = []
        for node in self.db_cluster.nodes:
            result = wait_for(
                func=self.db_cluster.is_features_enabled_on_node,
                timeout=60,
                step=1,
                text=f"Check feature enabled on node {node.name}",
                throw_exc=False,
                feature_list=["GROUP0_LIMITED_VOTERS"],
                node=node,
            )
            feature_state_per_node.append(result)
        if all(not enabled for enabled in feature_state_per_node):
            InfoEvent(message="Limited voters feature is not supported or enabled on all nodes").publish()
            return
        elif all(feature_state_per_node):
            InfoEvent(message="Limited voters feature enabled on all nodes").publish()
            num_of_nodes = len(self.db_cluster.nodes)
            num_of_non_voters = len(self.db_cluster.nodes[0].raft.get_group0_non_voters())
            if num_of_nodes <= max_voter_number and num_of_non_voters > 0:
                InfoEvent(
                    message=f"Number of voters: {num_of_nodes - num_of_non_voters} are less than expected: {max_voter_number}",
                    severity=Severity.ERROR,
                ).publish()
            elif num_of_nodes > max_voter_number and (num_of_nodes - num_of_non_voters) > max_voter_number:
                InfoEvent(
                    message=f"Number of voters: {num_of_nodes - num_of_non_voters} is higher than expected: {max_voter_number}",
                    severity=Severity.ERROR,
                ).publish()
            else:
                InfoEvent(message="Feature limited voters is enabled correctly").publish()
            return
        else:
            nodes_without_feature = [
                node.name
                for _, node in filter(lambda state: not state[0], zip(feature_state_per_node, self.db_cluster.nodes))
            ]
            raise Group0LimitedVotersFeatureNotEnableOnNodes(
                f"Limited voters feature not enabled on all nodes: {','.join(nodes_without_feature)}"
            )


class UpgradeCustomTest(UpgradeTest):
    def test_custom_profile_rolling_upgrade(self):
        """
        Run a load of a custom profile.
        Upgrade half of nodes in the cluster, and run rollback.
        Upgrade all nodes to new version in the end.
        """
        cs_user_profiles = self.prepare_data_before_upgrade()
        self._custom_profile_rolling_upgrade(cs_user_profiles=cs_user_profiles)

    def test_custom_profile_sequential_rolling_upgrade(self):
        """
        Run a load of a custom profile.
        Perform sequential rolling upgrade with upgrade path:
        step 1. 2021.1 -> 2022.1
        step 2. 2022.1 -> 2023.1
        In each step: upgrade half of nodes in the cluster, and run rollback.
        Upgrade all nodes to new version in the end.
        """
        cs_user_profiles = self.prepare_data_before_upgrade()
        # Temporary solution. If we want to have sequential upgrade test, we will need to change the test parameters and build upgrade path
        new_version = "2022.1"
        new_scylla_repo = "https://s3.amazonaws.com/downloads.scylladb.com/deb/ubuntu/scylla-2022.1.listscylla.list"
        InfoEvent(message=f"Starting rolling upgrade test to {new_version}").publish()
        self._custom_profile_rolling_upgrade(
            cs_user_profiles=cs_user_profiles, new_version=new_version, new_scylla_repo=new_scylla_repo
        )
        InfoEvent(message=f"Finished rolling upgrade test to {new_version}").publish()

        new_version = "2023.1"
        new_scylla_repo = "https://s3.amazonaws.com/downloads.scylladb.com/deb/ubuntu/scylla-2023.1.list"
        InfoEvent(message=f"Starting rolling upgrade test to {new_version}").publish()
        self._custom_profile_rolling_upgrade(
            cs_user_profiles=cs_user_profiles, new_version=new_version, new_scylla_repo=new_scylla_repo
        )
        InfoEvent(message=f"Finished rolling upgrade test to {new_version}").publish()

    def prepare_data_before_upgrade(self):
        InfoEvent(message="Running a prepare load for the initial custom data").publish()
        if not (prepare_cs_user_profiles := self.params.get("prepare_cs_user_profiles")):
            raise SyntaxError("Parameter 'prepare_cs_user_profiles' is not supplied")

        user_profiles, duration_per_cs_profile = self.parse_cs_user_profiles_param(prepare_cs_user_profiles)
        stress_before_upgrade = self.run_cs_user_profiles(
            cs_profiles=user_profiles, duration_per_cs_profile=duration_per_cs_profile
        )
        for queue in stress_before_upgrade:
            self.verify_stress_thread(queue)

        # write workload during entire test
        if not (cs_user_profiles := self.params.get("cs_user_profiles")):
            raise SyntaxError("Parameter 'cs_user_profiles' is not supplied")

        return cs_user_profiles

    def _custom_profile_rolling_upgrade(self, cs_user_profiles, new_scylla_repo=None, new_version=None):
        self.upgrade_os(self.db_cluster.nodes)

        InfoEvent(message="Starting write workload during entire test").publish()
        user_profiles, duration_per_cs_profile = self.parse_cs_user_profiles_param(cs_user_profiles)
        entire_write_thread_pool = self.run_cs_user_profiles(
            cs_profiles=user_profiles, duration_per_cs_profile=duration_per_cs_profile
        )

        # Let wait for start writing data
        for stress_cmd in user_profiles:
            _, profile = get_profile_content(stress_cmd)
            keyspace_name = profile.get("keyspace")
            self.log.debug("keyspace_name: %s", keyspace_name)
            if not keyspace_name:
                continue

            try:
                self.metric_has_data(
                    metric_query='sct_cassandra_stress_user_gauge{type="ops", keyspace="%s"}' % keyspace_name, n=10
                )
            except Exception as err:  # noqa: BLE001
                InfoEvent(
                    f"Get metrix data for keyspace {keyspace_name} failed with error: {err}", severity=Severity.ERROR
                ).publish()

        # generate random order to upgrade
        nodes_num = len(self.db_cluster.nodes)
        # prepare an array containing the indexes
        indexes = list(range(nodes_num))
        # shuffle it so we will upgrade the nodes in a random order
        random.shuffle(indexes)

        with ignore_upgrade_schema_errors():
            step = "Step1 - Upgrade First Node "
            InfoEvent(message=step).publish()
            # upgrade first node
            self.db_cluster.node_to_upgrade = self.db_cluster.nodes[indexes[0]]
            # Call "_upgrade_node" to prevent running truncate test
            self._upgrade_node(
                self.db_cluster.node_to_upgrade, new_scylla_repo=new_scylla_repo, new_version=new_version
            )
            self.db_cluster.node_to_upgrade.check_node_health()

            InfoEvent(message="after upgraded one node").publish()
            self.search_for_idx_token_error_after_upgrade(
                node=self.db_cluster.node_to_upgrade, step=step + " - after upgraded one node"
            )

            step = "Step2 - Upgrade Second Node "
            InfoEvent(message=step).publish()
            # upgrade second node
            self.db_cluster.node_to_upgrade = self.db_cluster.nodes[indexes[1]]
            # Call "_upgrade_node" to prevent running truncate test
            self._upgrade_node(
                self.db_cluster.node_to_upgrade, new_scylla_repo=new_scylla_repo, new_version=new_version
            )
            self.db_cluster.node_to_upgrade.check_node_health()

            self.search_for_idx_token_error_after_upgrade(
                node=self.db_cluster.node_to_upgrade, step=step + " - after upgraded two nodes"
            )

            InfoEvent(message="Step3 - Rollback Second Node ").publish()
            # rollback second node, use "_rollback_node" to prevent running truncate test
            self._rollback_node(self.db_cluster.nodes[indexes[1]])
            self.db_cluster.nodes[indexes[1]].check_node_health()

        step = "Step4 - Verify data during mixed cluster mode "
        InfoEvent(message=step).publish()
        InfoEvent(message="Repair the first upgraded Node").publish()
        self.db_cluster.nodes[indexes[0]].run_nodetool(sub_cmd="repair")
        self.search_for_idx_token_error_after_upgrade(node=self.db_cluster.node_to_upgrade, step=step)

        with ignore_upgrade_schema_errors():
            step = "Step5 - Upgrade rest of the Nodes "
            InfoEvent(message=step).publish()
            for i in indexes[1:]:
                self.db_cluster.node_to_upgrade = self.db_cluster.nodes[i]
                # Call "_upgrade_node" to prevent running truncate test
                self._upgrade_node(
                    self.db_cluster.node_to_upgrade, new_scylla_repo=new_scylla_repo, new_version=new_version
                )
                self.db_cluster.node_to_upgrade.check_node_health()
                self.search_for_idx_token_error_after_upgrade(node=self.db_cluster.node_to_upgrade, step=step)

        InfoEvent(message="Step6 - Verify stress results after upgrade ").publish()
        InfoEvent(message="Waiting for stress threads to complete after upgrade").publish()

        for queue in entire_write_thread_pool:
            self.verify_stress_thread(queue)

        InfoEvent(message="Step7 - Upgrade sstables to latest supported version ").publish()
        # figure out what is the last supported sstable version
        self.expected_sstable_format_version = self.get_highest_supported_sstable_version()

        # run 'nodetool upgradesstables' on all nodes and check/wait for all file to be upgraded
        upgradesstables = self.db_cluster.run_func_parallel(func=self.upgradesstables_if_command_available)

        # only check sstable format version if all nodes had 'nodetool upgradesstables' available
        if all(upgradesstables):
            InfoEvent(message="Upgrading sstables if new version is available").publish()
            tables_upgraded = self.db_cluster.run_func_parallel(func=self.wait_for_sstable_upgrade)
            assert all(tables_upgraded), f"Failed to upgrade the sstable format {tables_upgraded}"

        InfoEvent(message="all nodes were upgraded, and last workaround is verified.").publish()
