From a4b62ef13b2a82f9a3498d100e487e8ae9626be5 Mon Sep 17 00:00:00 2001
From: Denis Potapov <denis.potapov@aiven.io>
Date: Thu, 28 Sep 2023 11:30:46 +0200
Subject: [PATCH] Adjust default values for ClickHousePlugin parameters

Calculate the default value for max_concurrent_ parameters based on the
number of Clickhouse nodes.
---
 .../coordinator/plugins/clickhouse/plugin.py  | 16 ++++++
 .../plugins/clickhouse/test_plugin.py         | 56 +++++++++++++++++++
 2 files changed, 72 insertions(+)
 create mode 100644 tests/unit/coordinator/plugins/clickhouse/test_plugin.py

diff --git a/astacus/coordinator/plugins/clickhouse/plugin.py b/astacus/coordinator/plugins/clickhouse/plugin.py
index 80408eee..258a0b02 100644
--- a/astacus/coordinator/plugins/clickhouse/plugin.py
+++ b/astacus/coordinator/plugins/clickhouse/plugin.py
@@ -54,6 +54,7 @@
 )
 from astacus.coordinator.plugins.zookeeper_config import ZooKeeperConfiguration
 from pathlib import Path
+from pydantic import root_validator
 from typing import List, Sequence
 
 
@@ -83,6 +84,21 @@ class ClickHousePlugin(CoordinatorPlugin):
     max_concurrent_sync: int = 100
     use_system_unfreeze: bool = True
 
+    @root_validator(pre=True)
+    def _apply_max_concurrent_query_defaults(cls, values):
+        default_value = 10 * len(values["clickhouse"].nodes) if "clickhouse" in values else 100
+        properties = [
+            "max_concurrent_create_databases",
+            "max_concurrent_drop_databases",
+            "max_concurrent_restart_replica",
+            "max_concurrent_restore_replica",
+            "max_concurrent_attach",
+            "max_concurrent_sync",
+        ]
+        for property in properties:
+            values.setdefault(property, default_value)
+        return values
+
     def get_backup_steps(self, *, context: OperationContext) -> List[Step]:
         zookeeper_client = get_zookeeper_client(self.zookeeper)
         clickhouse_clients = get_clickhouse_clients(self.clickhouse)
diff --git a/tests/unit/coordinator/plugins/clickhouse/test_plugin.py b/tests/unit/coordinator/plugins/clickhouse/test_plugin.py
new file mode 100644
index 00000000..4cd5a9da
--- /dev/null
+++ b/tests/unit/coordinator/plugins/clickhouse/test_plugin.py
@@ -0,0 +1,56 @@
+"""
+Copyright (c) 2023 Aiven Ltd
+See LICENSE for details
+"""
+from astacus.coordinator.plugins.clickhouse.config import ClickHouseConfiguration, ClickHouseNode
+from astacus.coordinator.plugins.clickhouse.plugin import ClickHousePlugin
+
+import pytest
+
+
+@pytest.mark.parametrize(
+    "property",
+    [
+        "max_concurrent_create_databases",
+        "max_concurrent_drop_databases",
+        "max_concurrent_restart_replica",
+        "max_concurrent_restore_replica",
+        "max_concurrent_attach",
+        "max_concurrent_sync",
+    ],
+)
+def test_max_concurrent_queries_params_have_default_values(property: str) -> None:
+    plugin = ClickHousePlugin(
+        clickhouse=ClickHouseConfiguration(
+            nodes=[
+                ClickHouseNode(host="host", port=4242),
+                ClickHouseNode(host="host", port=4242),
+                ClickHouseNode(host="host", port=4242),
+            ]
+        )
+    )
+    assert getattr(plugin, property) == 30
+
+
+def test_max_concurrent_queries_params_accept_explicit_values() -> None:
+    plugin = ClickHousePlugin(
+        max_concurrent_create_databases=1,
+        max_concurrent_drop_databases=2,
+        max_concurrent_restart_replica=3,
+        max_concurrent_restore_replica=4,
+        max_concurrent_attach=5,
+        max_concurrent_sync=6,
+        clickhouse=ClickHouseConfiguration(
+            nodes=[
+                ClickHouseNode(host="host", port=4242),
+                ClickHouseNode(host="host", port=4242),
+                ClickHouseNode(host="host", port=4242),
+            ]
+        ),
+    )
+    assert plugin.max_concurrent_create_databases == 1
+    assert plugin.max_concurrent_drop_databases == 2
+    assert plugin.max_concurrent_restart_replica == 3
+    assert plugin.max_concurrent_restore_replica == 4
+    assert plugin.max_concurrent_attach == 5
+    assert plugin.max_concurrent_sync == 6