Freeze table by partitions (#219)

* Freeze table by partitions

* ExecPool for generic purposes

* Thread-safe freeze

* Correct timeout for freeze by partitions

* Fix lint

* ExecPool classes without explicit Thread/ProcessPool in constructors

* Remove ClickhouseClientMultithreading

* Small fixes

Author: kirillgarbar

ch_backup/backup/layout.py

@@ -88,7 +88,6 @@ def upload_table_create_statement(
         backup_meta: BackupMetadata,
         db: Database,
         table: Table,
-        create_statement: bytes,
     ) -> None:
         """
         Upload table create statement.
@@ -105,7 +104,7 @@ def upload_table_create_statement(
                 table.name,
             )
             self._storage_loader.upload_data(
-                create_statement,
+                table.create_statement,
                 remote_path,
                 is_async=True,
                 encryption=backup_meta.encrypted,

ch_backup/ch_backup.py

@@ -196,9 +196,7 @@ def backup(
                         databases,
                         db_tables,
                         schema_only=sources.schema_only,
-                        freeze_threads=self._config["multiprocessing"][
-                            "freeze_threads"
-                        ],
+                        multiprocessing_config=self._config["multiprocessing"],
                     )
 
                 # Upload operations are async. Should wait until they are all finished.

ch_backup/clickhouse/client.py

@@ -2,7 +2,8 @@
 ClickHouse client.
 """
 
-from typing import Any
+from contextlib import contextmanager
+from typing import Any, Iterator, Optional
 
 import requests
 
@@ -34,6 +35,8 @@ class ClickhouseClient:
     """
 
     def __init__(self, config: dict, settings: dict = None) -> None:
+        self._config = config
+        self._settings = settings
         host = config["host"]
         protocol = config["protocol"]
         port = config["port"] or (8123 if protocol == "http" else 8443)
@@ -58,6 +61,7 @@ def query(
         settings: dict = None,
         timeout: float = None,
         should_retry: bool = True,
+        new_session: bool = False,
     ) -> Any:
         """
         Execute query.
@@ -68,13 +72,18 @@ def query(
             if timeout is None:
                 timeout = self.timeout
 
-            response = self._session.post(
-                self._url,
-                params=settings,
-                json=post_data,
-                timeout=(self.connect_timeout, timeout),
-                data=query.encode("utf-8"),
-            )
+            # https://github.com/psf/requests/issues/2766
+            # requests.Session object is not guaranteed to be thread-safe.
+            # When using ClickhouseClient with multithreading, "new_session"
+            # should be True, so a separate Session is used for each query.
+            with self._get_session(new_session) as session:
+                response = session.post(
+                    self._url,
+                    params=settings,
+                    json=post_data,
+                    timeout=(self.connect_timeout, timeout),
+                    data=query.encode("utf-8"),
+                )
 
             response.raise_for_status()
         except requests.exceptions.HTTPError as e:
@@ -87,6 +96,21 @@ def query(
         except ValueError:
             return str.strip(response.text)
 
+    @contextmanager
+    def _get_session(
+        self, new_session: Optional[bool] = False
+    ) -> Iterator[requests.Session]:
+        session = (
+            self._create_session(self._config, self._settings)
+            if new_session
+            else self._session
+        )
+        try:
+            yield session
+        finally:
+            if new_session:
+                session.close()
+
     @staticmethod
     def _create_session(config, settings):
         session = requests.Session()

ch_backup/clickhouse/control.py

@@ -21,6 +21,7 @@
 from ch_backup.clickhouse.client import ClickhouseClient
 from ch_backup.clickhouse.models import Database, Disk, FrozenPart, Table
 from ch_backup.exceptions import ClickhouseBackupError
+from ch_backup.storage.async_pipeline.base_pipeline.exec_pool import ThreadExecPool
 from ch_backup.util import (
     chown_dir_contents,
     chown_file,
@@ -167,6 +168,13 @@
 """
 )
 
+FREEZE_PARTITION_SQL = strip_query(
+    """
+    ALTER TABLE `{db_name}`.`{table_name}`
+    FREEZE PARTITION ID '{partition}' WITH NAME '{backup_name}'
+"""
+)
+
 SYSTEM_UNFREEZE_SQL = strip_query(
     """
     SYSTEM UNFREEZE WITH NAME '{backup_name}'
@@ -435,6 +443,15 @@
 """
 )
 
+GET_PARTITIONS = strip_query(
+    """
+    SELECT DISTINCT partition_id
+    FROM system.parts
+    WHERE database = '{db_name}' AND table = '{table_name}' AND active = 1
+    FORMAT JSONCompact
+"""
+)
+
 
 # pylint: disable=too-many-public-methods
 class ClickhouseCTL:
@@ -535,18 +552,64 @@ def attach_table(self, table: Union[TableMetadata, Table]) -> None:
 
         self._ch_client.query(query_sql)
 
-    def freeze_table(self, backup_name: str, table: Table) -> None:
+    def freeze_table(
+        self,
+        backup_name: str,
+        table: Table,
+        threads: int,
+    ) -> None:
         """
         Make snapshot of the specified table.
         """
-        query_sql = FREEZE_TABLE_SQL.format(
+        # Table has no partitions or created with deprecated syntax.
+        # FREEZE PARTITION ID with deprecated syntax throws segmentation fault in CH.
+        freeze_by_partitions = threads > 0 and "PARTITION BY" in table.create_statement
+        if freeze_by_partitions:
+            with ThreadExecPool(max(1, threads)) as pool:
+                if freeze_by_partitions:
+                    partitions_to_freeze = self.list_partitions(table)
+                    for partition in partitions_to_freeze:
+                        query_sql = FREEZE_PARTITION_SQL.format(
+                            db_name=escape(table.database),
+                            table_name=escape(table.name),
+                            backup_name=backup_name,
+                            partition=partition,
+                        )
+                        pool.submit(
+                            f'Freeze partition "{partition}"',
+                            self._ch_client.query,
+                            query_sql,
+                            timeout=self._freeze_timeout,
+                            should_retry=False,
+                            new_session=True,
+                        )
+                pool.wait_all(
+                    keep_going=False,
+                    timeout=self._freeze_timeout,
+                )
+        else:
+            query_sql = FREEZE_TABLE_SQL.format(
+                db_name=escape(table.database),
+                table_name=escape(table.name),
+                backup_name=backup_name,
+            )
+            self._ch_client.query(
+                query_sql,
+                timeout=self._freeze_timeout,
+                should_retry=False,
+                new_session=True,
+            )
+
+    def list_partitions(self, table: Table) -> List[str]:
+        """
+        Get list of active partitions for table.
+        """
+        query_sql = GET_PARTITIONS.format(
             db_name=escape(table.database),
             table_name=escape(table.name),
-            backup_name=backup_name,
-        )
-        self._ch_client.query(
-            query_sql, timeout=self._freeze_timeout, should_retry=False
         )
+        result = self._ch_client.query(query_sql, should_retry=False)["data"]
+        return [partition[0] for partition in result]
 
     def system_unfreeze(self, backup_name: str) -> None:
         """

ch_backup/config.py

@@ -160,6 +160,8 @@ def _as_seconds(t: str) -> int:
         "cloud_storage_restore_workers": 4,
         # The number of threads for parallel freeze of tables
         "freeze_threads": 4,
+        # The number of threads for parallel freeze of partitions. If set to 0, will freeze table in one query.
+        "freeze_partition_threads": 16,
         # The number of threads for parallel drop replica
         "drop_replica_threads": 8,
     },

ch_backup/logic/table.py

@@ -4,7 +4,6 @@
 
 import os
 from collections import deque
-from concurrent.futures import Future, ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
 from functools import partial
 from itertools import chain
@@ -28,6 +27,9 @@
 from ch_backup.exceptions import ClickhouseBackupError
 from ch_backup.logic.backup_manager import BackupManager
 from ch_backup.logic.upload_part_observer import UploadPartObserver
+from ch_backup.storage.async_pipeline.base_pipeline.exec_pool import (
+    ThreadExecPool,
+)
 from ch_backup.util import compare_schema
 
 
@@ -54,7 +56,7 @@ def backup(
         databases: Sequence[Database],
         db_tables: Dict[str, list],
         schema_only: bool,
-        freeze_threads: int,
+        multiprocessing_config: Dict,
     ) -> None:
         """
         Backup tables metadata, MergeTree data and Cloud storage metadata.
@@ -76,7 +78,7 @@ def backup(
                 db_tables[db.name],
                 backup_name,
                 schema_only,
-                freeze_threads,
+                multiprocessing_config,
             )
 
     def _collect_local_metadata_change_times(
@@ -110,7 +112,7 @@ def _backup(
         tables: Sequence[str],
         backup_name: str,
         schema_only: bool,
-        freeze_threads: int,
+        multiprocessing_config: Dict,
     ) -> None:
         """
         Backup single database tables.
@@ -133,33 +135,32 @@ def _backup(
             # Create shadow/increment.txt if not exists manually to avoid
             # race condition with parallel freeze
             context.ch_ctl.create_shadow_increment()
-            futures: List[Future] = []
-            with ThreadPoolExecutor(max_workers=freeze_threads) as pool:
+            with ThreadExecPool(
+                multiprocessing_config.get("freeze_threads", 1)
+            ) as pool:
                 for table in tables_:
-                    future = pool.submit(
+                    pool.submit(
+                        f'Freeze table "{table.database}"."{table.name}"',
                         TableBackup._freeze_table,
                         context,
                         db,
                         table,
                         backup_name,
                         schema_only,
+                        multiprocessing_config.get("freeze_partition_threads", 0),
                     )
-                    futures.append(future)
 
-                for future in as_completed(futures):
-                    table_and_create_statement = future.result()
-                    if table_and_create_statement is not None:
-                        table, create_statement = table_and_create_statement
+                for freezed_table in pool.as_completed(keep_going=False):
+                    if freezed_table is not None:
                         self._backup_freezed_table(
                             context,
                             db,
-                            table,
+                            freezed_table,
                             backup_name,
                             schema_only,
                             change_times,
-                            create_statement,
                         )
-                        self._backup_cloud_storage_metadata(context, table)
+                        self._backup_cloud_storage_metadata(context, freezed_table)
 
             context.backup_layout.wait()
             context.ch_ctl.remove_freezed_data()
@@ -174,13 +175,14 @@ def _freeze_table(
         table: Table,
         backup_name: str,
         schema_only: bool,
-    ) -> Optional[Tuple[Table, bytes]]:
+        freeze_partition_threads: int,
+    ) -> Optional[Table]:
         """
         Freeze table and return it's create statement
         """
         logging.debug('Trying to freeze "{}"."{}"', table.database, table.name)
-
         create_statement = TableBackup._load_create_statement_from_disk(table)
+        table.create_statement = create_statement or ""
         if not create_statement:
             logging.warning(
                 'Skipping table backup for "{}"."{}". Local metadata is empty or absent',
@@ -192,7 +194,9 @@ def _freeze_table(
         # Freeze only MergeTree tables
         if not schema_only and table.is_merge_tree():
             try:
-                context.ch_ctl.freeze_table(backup_name, table)
+                context.ch_ctl.freeze_table(
+                    backup_name, table, freeze_partition_threads
+                )
             except ClickhouseError:
                 if context.ch_ctl.does_table_exist(table.database, table.name):
                     logging.error(
@@ -209,10 +213,10 @@ def _freeze_table(
                 )
                 return None
 
-        return (table, create_statement)
+        return table
 
     @staticmethod
-    def _load_create_statement_from_disk(table: Table) -> Optional[bytes]:
+    def _load_create_statement_from_disk(table: Table) -> Optional[str]:
         """
         Load a create statement of the table from a metadata file on the disk.
         """
@@ -224,7 +228,7 @@ def _load_create_statement_from_disk(table: Table) -> Optional[bytes]:
             )
             return None
         try:
-            return Path(table.metadata_path).read_bytes()
+            return Path(table.metadata_path).read_text("utf-8")
         except OSError as e:
             logging.debug(
                 'Cannot load a create statement of the table "{}"."{}": {}',
@@ -378,7 +382,6 @@ def _backup_freezed_table(
         backup_name: str,
         schema_only: bool,
         change_times: Dict[str, TableMetadataChangeTime],
-        create_statement: bytes,
     ) -> None:
         # Check if table metadata was updated
         new_change_time = self._get_change_time(table.metadata_path)
@@ -400,7 +403,7 @@ def _backup_freezed_table(
         )
         # Backup table metadata
         context.backup_layout.upload_table_create_statement(
-            context.backup_meta, db, table, create_statement
+            context.backup_meta, db, table
         )
         # Backup table data
         if not schema_only: