From deeefdf95d186d727e599067f8ff57de361463a1 Mon Sep 17 00:00:00 2001
From: Alexander Burmak <Alex-Burmak@yandex-team.ru>
Date: Fri, 16 Feb 2024 14:23:21 +0300
Subject: [PATCH] Moved timeout for alter queries to config (#103)

---
 ch_tools/chadmin/internal/partition.py | 52 +++++---------------------
 ch_tools/common/config.py              |  3 +-
 2 files changed, 12 insertions(+), 43 deletions(-)

diff --git a/ch_tools/chadmin/internal/partition.py b/ch_tools/chadmin/internal/partition.py
index a2138801..f19d02d4 100644
--- a/ch_tools/chadmin/internal/partition.py
+++ b/ch_tools/chadmin/internal/partition.py
@@ -1,21 +1,12 @@
 from ch_tools.chadmin.internal.utils import execute_query
 
-DEFAULT_TIMEOUT = 600
-
 
 def attach_partition(ctx, database, table, partition_id, dry_run=False):
     """
     Attach the specified table partition.
     """
     query = f"ALTER TABLE `{database}`.`{table}` ATTACH PARTITION ID '{partition_id}'"
-    execute_query(
-        ctx,
-        query,
-        timeout=DEFAULT_TIMEOUT,
-        format_=None,
-        echo=True,
-        dry_run=dry_run,
-    )
+    _execute_query(ctx, query, dry_run)
 
 
 def detach_partition(ctx, database, table, partition_id, dry_run=False):
@@ -23,14 +14,7 @@ def detach_partition(ctx, database, table, partition_id, dry_run=False):
     Detach the specified table partition.
     """
     query = f"ALTER TABLE `{database}`.`{table}` DETACH PARTITION ID '{partition_id}'"
-    execute_query(
-        ctx,
-        query,
-        timeout=DEFAULT_TIMEOUT,
-        format_=None,
-        echo=True,
-        dry_run=dry_run,
-    )
+    _execute_query(ctx, query, dry_run)
 
 
 def drop_partition(ctx, database, table, partition_id, dry_run=False):
@@ -38,14 +22,7 @@ def drop_partition(ctx, database, table, partition_id, dry_run=False):
     Drop the specified table partition.
     """
     query = f"ALTER TABLE `{database}`.`{table}` DROP PARTITION ID '{partition_id}'"
-    execute_query(
-        ctx,
-        query,
-        timeout=DEFAULT_TIMEOUT,
-        format_=None,
-        echo=True,
-        dry_run=dry_run,
-    )
+    _execute_query(ctx, query, dry_run)
 
 
 def optimize_partition(ctx, database, table, partition_id, dry_run=False):
@@ -53,14 +30,7 @@ def optimize_partition(ctx, database, table, partition_id, dry_run=False):
     Optimize the specified table partition.
     """
     query = f"OPTIMIZE TABLE `{database}`.`{table}` PARTITION ID '{partition_id}'"
-    execute_query(
-        ctx,
-        query,
-        timeout=DEFAULT_TIMEOUT,
-        format_=None,
-        echo=True,
-        dry_run=dry_run,
-    )
+    _execute_query(ctx, query, dry_run)
 
 
 def materialize_ttl_in_partition(ctx, database, table, partition_id, dry_run=False):
@@ -68,11 +38,9 @@ def materialize_ttl_in_partition(ctx, database, table, partition_id, dry_run=Fal
     Materialize TTL for the specified table partition.
     """
     query = f"ALTER TABLE `{database}`.`{table}` MATERIALIZE TTL IN PARTITION ID '{partition_id}'"
-    execute_query(
-        ctx,
-        query,
-        timeout=DEFAULT_TIMEOUT,
-        format_=None,
-        echo=True,
-        dry_run=dry_run,
-    )
+    _execute_query(ctx, query, dry_run)
+
+
+def _execute_query(ctx, query, dry_run):
+    timeout = ctx.obj["config"]["clickhouse"]["alter_table_timeout"]
+    execute_query(ctx, query, timeout=timeout, format_=None, echo=True, dry_run=dry_run)
diff --git a/ch_tools/common/config.py b/ch_tools/common/config.py
index 424cdcdc..ec34a5e8 100644
--- a/ch_tools/common/config.py
+++ b/ch_tools/common/config.py
@@ -14,11 +14,12 @@
         "port": 8443,
         "user": None,
         "password": None,
-        "timeout": 60,
         "settings": {},
         "monitoring_user": None,
         "monitoring_password": None,
         "distributed_ddl_path": "/clickhouse/task_queue/ddl",
+        "timeout": 60,
+        "alter_table_timeout": 600,
     },
     "object_storage": {
         "clean": {