Backup s3 metadata simultaneously with local data (#192)

* Upload cloud storage metadata parallel with local data

* Fix bug

* Fix lint

* Check if cloud metadata is old

* Fix comment

* Fix lint

* Delete todo

* Delte todo 2

* Fix python 3.10 pylint

Author: kirillgarbar

ch_backup/backup/layout.py

@@ -197,18 +197,28 @@ def upload_data_part(
             raise StorageError(msg) from e
 
     def upload_cloud_storage_metadata(
-        self, backup_meta: BackupMetadata, disk: Disk, delete_after_upload: bool = False
+        self,
+        backup_meta: BackupMetadata,
+        disk: Disk,
+        table: Table,
+        delete_after_upload: bool = False,
     ) -> bool:
         """
         Upload specified disk metadata files from given directory path as a tarball.
         Returns: whether backed up disk had data.
         """
+        assert table.path_on_disk, f"Table {table} doesn't store data on disk"
+
         backup_name = backup_meta.get_sanitized_name()
         compression = backup_meta.cloud_storage.compressed
         remote_path = _disk_metadata_path(
-            self.get_backup_path(backup_name), disk.name, compression
+            self.get_backup_path(backup_name),
+            table.database,
+            table.name,
+            disk.name,
+            compression,
         )
-        shadow_path = os.path.join(disk.path, "shadow", backup_name)
+        shadow_path = _table_shadow_path(disk.path, backup_name, table.path_on_disk)
         exclude_file_names = ["frozen_metadata.txt"]
         if dir_is_empty(shadow_path, exclude_file_names):
             return False
@@ -219,6 +229,7 @@ def upload_cloud_storage_metadata(
             self._storage_loader.upload_files_tarball_scan(
                 dir_path=shadow_path,
                 remote_path=remote_path,
+                tar_base_dir=table.path_on_disk,
                 exclude_file_names=exclude_file_names,
                 is_async=True,
                 encryption=backup_meta.cloud_storage.encrypted,
@@ -532,6 +543,28 @@ def check_data_part(self, backup_path: str, part: PartMetadata) -> bool:
             )
             return False
 
+    def _get_cloud_storage_metadata_remote_paths(
+        self,
+        backup_name: str,
+        source_disk_name: str,
+        compression: bool,
+    ) -> Sequence[str]:
+        # Check if metadata is stored as 'disks/s3.tar.gz' for backwards compatibility
+        old_style_remote_path = _disk_metadata_path(
+            self.get_backup_path(backup_name), None, None, source_disk_name, compression
+        )
+        if self._storage_loader.path_exists(old_style_remote_path):
+            return [old_style_remote_path]
+        return self._storage_loader.list_dir(
+            str(
+                os.path.join(
+                    self.get_backup_path(backup_name), "disks", source_disk_name
+                )
+            ),
+            recursive=True,
+            absolute=True,
+        )
+
     def download_cloud_storage_metadata(
         self, backup_meta: BackupMetadata, disk: Disk, source_disk_name: str
     ) -> None:
@@ -542,22 +575,23 @@ def download_cloud_storage_metadata(
         compression = backup_meta.cloud_storage.compressed
         disk_path = os.path.join(disk.path, "shadow", backup_name)
         os.makedirs(disk_path, exist_ok=True)
-        remote_path = _disk_metadata_path(
-            self.get_backup_path(backup_name), source_disk_name, compression
+        metadata_remote_paths = self._get_cloud_storage_metadata_remote_paths(
+            backup_name, source_disk_name, compression
         )
 
-        logging.debug(f'Downloading "{disk_path}" files from "{remote_path}"')
-        try:
-            self._storage_loader.download_files(
-                remote_path=remote_path,
-                local_path=disk_path,
-                is_async=True,
-                encryption=backup_meta.cloud_storage.encrypted,
-                compression=compression,
-            )
-        except Exception as e:
-            msg = f'Failed to download tarball file "{remote_path}"'
-            raise StorageError(msg) from e
+        for remote_path in metadata_remote_paths:
+            logging.debug(f'Downloading "{disk_path}" files from "{remote_path}"')
+            try:
+                self._storage_loader.download_files(
+                    remote_path=remote_path,
+                    local_path=disk_path,
+                    is_async=True,
+                    encryption=backup_meta.cloud_storage.encrypted,
+                    compression=compression,
+                )
+            except Exception as e:
+                msg = f'Failed to download tarball file "{remote_path}"'
+                raise StorageError(msg) from e
 
     def delete_backup(self, backup_name: str) -> None:
         """
@@ -715,17 +749,39 @@ def _part_path(
 
 
 def _disk_metadata_path(
-    backup_path: str, disk_name: str, compressed: bool = False
+    backup_path: str,
+    db_name: Optional[str],
+    table_name: Optional[str],
+    disk_name: str,
+    compressed: bool = False,
 ) -> str:
     """
     Returns path to store tarball with cloud storage shadow metadata.
     """
     extension = ".tar"
     if compressed:
         extension += COMPRESSED_EXTENSION
+    if table_name or db_name:
+        assert table_name and db_name
+        return os.path.join(
+            backup_path,
+            "disks",
+            disk_name,
+            _quote(db_name),
+            f"{_quote(table_name)}{extension}",
+        )
     return os.path.join(backup_path, "disks", f"{disk_name}{extension}")
 
 
+def _table_shadow_path(
+    disk_path: str, backup_name: str, table_path_on_disk: str
+) -> str:
+    """
+    Returns path to frozen table data on given disk.
+    """
+    return os.path.join(disk_path, "shadow", backup_name, table_path_on_disk)
+
+
 def _quote(value: str) -> str:
     return quote(value, safe="").translate(
         {

ch_backup/backup/metadata/backup_metadata.py

@@ -254,13 +254,15 @@ def add_database(self, db: Database) -> None:
             "tables": {},
         }
 
-    def get_tables(self, db_name: str) -> Sequence[TableMetadata]:
+    def get_tables(self, db_name: Optional[str] = None) -> Sequence[TableMetadata]:
         """
         Get tables for the specified database.
         """
         result = []
-        for table_name, raw_metadata in self._databases[db_name]["tables"].items():
-            result.append(TableMetadata.load(db_name, table_name, raw_metadata))
+        databases = [db_name] if db_name else self._databases.keys()
+        for db in databases:
+            for table_name, raw_metadata in self._databases[db]["tables"].items():
+                result.append(TableMetadata.load(db, table_name, raw_metadata))
 
         return result
 

ch_backup/backup/metadata/cloud_storage_metadata.py

@@ -38,7 +38,8 @@ def add_disk(self, disk_name: str) -> None:
         """
         Add disk name in backed up disks list.
         """
-        self._disks.append(disk_name)
+        if disk_name not in self._disks:
+            self._disks.append(disk_name)
 
     @property
     def encrypted(self) -> bool:

ch_backup/clickhouse/disks.py

@@ -184,7 +184,7 @@ def _create_temporary_disk(
 
         self._ch_ctl.reload_config()
         source_disk = self._ch_ctl.get_disk(tmp_disk_name)
-        logging.debug(f'Restoring Cloud Storage "shadow" data  of disk "{disk_name}"')
+        logging.debug(f'Restoring Cloud Storage "shadow" data of disk "{disk_name}"')
         self._backup_layout.download_cloud_storage_metadata(
             backup_meta, source_disk, disk_name
         )

ch_backup/clickhouse/models.py

@@ -2,6 +2,7 @@
 ClickHouse resource models.
 """
 
+import os
 import re
 from types import SimpleNamespace
 from typing import List, Optional, Tuple
@@ -75,6 +76,11 @@ def __init__(
         self.paths_with_disks = self._map_paths_to_disks(disks, data_paths)
         self.metadata_path = metadata_path
 
+        self.path_on_disk = None
+        if self.paths_with_disks:
+            path, disk = self.paths_with_disks[0]
+            self.path_on_disk = os.path.relpath(path, disk.path)
+
     def _map_paths_to_disks(
         self, disks: List[Disk], data_paths: List[str]
     ) -> List[Tuple[str, Disk]]:

ch_backup/logic/table.py

@@ -61,6 +61,13 @@ def backup(
 
         backup_name = context.backup_meta.get_sanitized_name()
 
+        if context.cloud_conf.get("cloud_storage", {}).get("encryption", True):
+            logging.debug('Cloud Storage "shadow" backup will be encrypted')
+            context.backup_meta.cloud_storage.encrypt()
+        if context.cloud_conf.get("cloud_storage", {}).get("compression", True):
+            logging.debug('Cloud Storage "shadow" backup will be compressed')
+            context.backup_meta.cloud_storage.compress()
+
         for db in databases:
             self._backup(
                 context,
@@ -71,8 +78,6 @@ def backup(
                 freeze_threads,
             )
 
-        self._backup_cloud_storage_metadata(context)
-
     def _collect_local_metadata_mtime(
         self, context: BackupContext, db: Database, tables: Sequence[str]
     ) -> Dict[str, TableMetadataMtime]:
@@ -153,6 +158,7 @@ def _backup(
                             mtimes,
                             create_statement,
                         )
+                        self._backup_cloud_storage_metadata(context, table)
 
             context.backup_layout.wait()
             context.ch_ctl.remove_freezed_data()
@@ -223,28 +229,23 @@ def _load_create_statement_from_disk(table: Table) -> Optional[bytes]:
             return None
 
     @staticmethod
-    def _backup_cloud_storage_metadata(context: BackupContext) -> None:
+    def _backup_cloud_storage_metadata(context: BackupContext, table: Table) -> None:
         """
         Backup cloud storage metadata files.
         """
-        if context.cloud_conf.get("cloud_storage", {}).get("encryption", True):
-            logging.debug('Cloud Storage "shadow" backup will be encrypted')
-            context.backup_meta.cloud_storage.encrypt()
-        if context.cloud_conf.get("cloud_storage", {}).get("compression", True):
-            logging.debug('Cloud Storage "shadow" backup will be compressed')
-            context.backup_meta.cloud_storage.compress()
-
-        logging.debug('Backing up Cloud Storage disks "shadow" directory')
-        disks = context.ch_ctl.get_disks()
-        for disk in disks.values():
+        logging.debug(
+            'Backing up Cloud Storage disks "shadow" directory of "{}"."{}"',
+            table.database,
+            table.name,
+        )
+        for _, disk in table.paths_with_disks:
             if disk.type == "s3" and not disk.cache_path:
                 if not context.backup_layout.upload_cloud_storage_metadata(
-                    context.backup_meta, disk
+                    context.backup_meta, disk, table
                 ):
                     logging.debug(f'No data frozen on disk "{disk.name}", skipping')
                     continue
                 context.backup_meta.cloud_storage.add_disk(disk.name)
-        logging.debug("Cloud Storage disks has been backed up ")
 
     # pylint: disable=too-many-arguments,too-many-locals,too-many-positional-arguments
     def restore(

ch_backup/storage/async_pipeline/pipeline_builder.py

@@ -112,7 +112,10 @@ def build_decompress_stage(self) -> "PipelineBuilder":
         return self
 
     def build_read_files_tarball_scan_stage(
-        self, dir_path: Path, exclude_file_names: Optional[List[str]] = None
+        self,
+        dir_path: Path,
+        tar_base_dir: Optional[str] = None,
+        exclude_file_names: Optional[List[str]] = None,
     ) -> "PipelineBuilder":
         """
         Build reading files to tarball stage.
@@ -122,7 +125,9 @@ def build_read_files_tarball_scan_stage(
 
         self.append(
             thread_input(
-                ReadFilesTarballScanStage(stage_config, dir_path, exclude_file_names),
+                ReadFilesTarballScanStage(
+                    stage_config, dir_path, tar_base_dir, exclude_file_names
+                ),
                 maxsize=queue_size,
             )
         )

ch_backup/storage/async_pipeline/pipeline_executor.py

@@ -90,6 +90,7 @@ def upload_files_tarball_scan(
         encryption: bool,
         delete: bool,
         compression: bool,
+        tar_base_dir: Optional[str] = None,
         exclude_file_names: Optional[List[str]] = None,
         callback: Optional[Callable] = None,
     ) -> None:
@@ -107,6 +108,7 @@ def upload_files_tarball_scan(
             encryption,
             delete_after=delete,
             compression=compression,
+            tar_base_dir=tar_base_dir,
             exclude_file_names=exclude_file_names,
         )
         self._exec_pipeline(job_id, pipeline, is_async, callback)

ch_backup/storage/async_pipeline/pipelines.py

@@ -67,6 +67,7 @@ def upload_files_tarball_scan_pipeline(
     encrypt: bool,
     delete_after: bool,
     compression: bool,
+    tar_base_dir: Optional[str] = None,
     exclude_file_names: Optional[List[str]] = None,
 ) -> None:
     """
@@ -80,7 +81,9 @@ def upload_files_tarball_scan_pipeline(
     estimated_size = calc_tarball_size_scan(
         base_path, estimated_size, exclude_file_names
     )
-    builder.build_read_files_tarball_scan_stage(base_path, exclude_file_names)
+    builder.build_read_files_tarball_scan_stage(
+        base_path, tar_base_dir, exclude_file_names
+    )
     if compression:
         builder.build_compress_stage()
     if encrypt:

ch_backup/storage/async_pipeline/stages/filesystem/read_files_tarball_stage.py

@@ -22,19 +22,26 @@ def __init__(
         self,
         config: dict,
         base_path: Path,
+        tar_base_dir: Optional[str] = None,
     ) -> None:
         self._chunk_size = config["chunk_size"]
         self._base_path = base_path
         self._file_source: Iterable[Any] = []
+        self._tar_base_dir: Optional[str] = tar_base_dir
 
     def __call__(self) -> Iterator[bytes]:
         """
         Read files and yield them as TAR stream.
         """
         for file_relative_path in self._file_source:
             file_path = self._base_path / file_relative_path
+            file_path_in_tar = (
+                Path(self._tar_base_dir) / file_relative_path
+                if self._tar_base_dir
+                else file_relative_path
+            )
 
-            yield self.make_tar_header(str(file_relative_path), file_path)
+            yield self.make_tar_header(str(file_path_in_tar), file_path)
             yield from self.read_file_content(file_path)
 
     def read_file_content(self, file_path: Path) -> Iterator[bytes]:
@@ -70,9 +77,10 @@ def __init__(
         self,
         config: dict,
         base_path: Path,
+        tar_base_dir: Optional[str] = None,
         exclude_file_names: Optional[List[str]] = None,
     ) -> None:
-        super().__init__(config, base_path)
+        super().__init__(config, base_path, tar_base_dir)
         self._file_source = scan_dir_files(self._base_path, exclude_file_names)
 
 
@@ -82,7 +90,11 @@ class ReadFilesTarballStage(ReadFilesTarballStageBase):
     """
 
     def __init__(
-        self, config: dict, base_path: Path, file_relative_paths: List[Path]
+        self,
+        config: dict,
+        base_path: Path,
+        file_relative_paths: List[Path],
+        tar_base_dir: Optional[str] = None,
     ) -> None:
-        super().__init__(config, base_path)
+        super().__init__(config, base_path, tar_base_dir)
         self._file_source = file_relative_paths