From 478c79f241f8141bd5284c701d7b2056b91054b9 Mon Sep 17 00:00:00 2001 From: MikhailBurdukov Date: Wed, 29 May 2024 14:47:29 +0000 Subject: [PATCH] fix --- ch_backup/logic/access.py | 64 +++++++++++++++++++++++++++------------ ch_backup/util.py | 33 -------------------- 2 files changed, 45 insertions(+), 52 deletions(-) diff --git a/ch_backup/logic/access.py b/ch_backup/logic/access.py index 9e0e9d36..0165887a 100644 --- a/ch_backup/logic/access.py +++ b/ch_backup/logic/access.py @@ -5,6 +5,8 @@ import os import re import shutil +from contextlib import contextmanager +from tempfile import TemporaryDirectory, mkdtemp from typing import Any, Dict, List, Sequence, Union from kazoo.client import KazooClient @@ -14,15 +16,30 @@ from ch_backup.backup.metadata import BackupStorageFormat from ch_backup.backup_context import BackupContext from ch_backup.logic.backup_manager import BackupManager -from ch_backup.util import ( - chown_dir_contents, - copy_directory_content, - temporary_directory, -) +from ch_backup.util import chown_dir_contents, copy_directory_content CH_MARK_FILE = "need_rebuild_lists.mark" +@contextmanager +def access_control_temp_directory( + directory_path: str, backup_name: str, keep_on_exception: bool = False +) -> Any: + """ + Class to automatically create and remove temporary directory. + + If keep flag is set, then in case of exception, do not remove the folder. + """ + _prefix = backup_name + "_tmp_" + if keep_on_exception: + tmpdir = mkdtemp(prefix=_prefix, dir=directory_path) + yield tmpdir + shutil.rmtree(tmpdir) + else: + with TemporaryDirectory(prefix=_prefix, dir=directory_path) as tmpdir: + yield tmpdir + + class AccessBackup(BackupManager): """ Access backup class @@ -34,16 +51,17 @@ def backup(self, context: BackupContext) -> None: """ clickhouse_access_path = context.ch_ctl_conf["access_control_path"] - backup_tmp_path = os.path.join( - context.ch_ctl_conf["tmp_path"], context.backup_meta.name - ) user = context.ch_ctl_conf["user"] group = context.ch_ctl_conf["group"] - os.makedirs(clickhouse_access_path, exist_ok=True) - shutil.chown(clickhouse_access_path, user, group) + self._prepare_folder(clickhouse_access_path, user, group) + self._prepare_folder(context.ch_ctl_conf["tmp_path"], user, group, False) - with temporary_directory(backup_tmp_path, user, group): + with access_control_temp_directory( + context.ch_ctl_conf["tmp_path"], + context.backup_meta.name, + keep_on_exception=True, + ) as backup_tmp_path: objects = context.ch_ctl.get_access_control_objects() context.backup_meta.set_access_control(objects) access_control = context.backup_meta.access_control @@ -81,19 +99,18 @@ def restore(self, context: BackupContext) -> None: has_replicated_access = self._has_replicated_access(context) clickhouse_access_path = context.ch_ctl_conf["access_control_path"] - restore_tmp_path = os.path.join( - context.ch_ctl_conf["tmp_path"], context.backup_meta.name - ) user = context.ch_ctl_conf["user"] group = context.ch_ctl_conf["group"] - if os.path.exists(clickhouse_access_path): - shutil.rmtree(clickhouse_access_path) + self._prepare_folder(clickhouse_access_path, user, group, True) + self._prepare_folder(context.ch_ctl_conf["tmp_path"], user, group, False) - os.makedirs(clickhouse_access_path) - shutil.chown(clickhouse_access_path, user, group) + with access_control_temp_directory( + context.ch_ctl_conf["tmp_path"], + context.backup_meta.name, + keep_on_exception=True, + ) as restore_tmp_path: - with temporary_directory(restore_tmp_path, user, group): self._download_access_control_list(context, restore_tmp_path, acl_ids) if has_replicated_access: @@ -172,6 +189,15 @@ def fix_admin_user(self, context: BackupContext, dry_run: bool = True) -> None: except FileNotFoundError: logging.debug(f"File {file_path} not found.") + def _prepare_folder( + self, path: str, user: str, group: str, cleanup: bool = False + ) -> None: + if cleanup and os.path.exists(path): + shutil.rmtree(path) + + os.makedirs(path, exist_ok=True) + shutil.chown(path, user, group) + def _download_access_control_list( self, context: BackupContext, restore_tmp_path: str, acl_ids: List[str] ) -> None: diff --git a/ch_backup/util.py b/ch_backup/util.py index b187191c..b778b1e3 100644 --- a/ch_backup/util.py +++ b/ch_backup/util.py @@ -485,36 +485,3 @@ def copy_directory_content(from_path_dir: str, to_path_dir: str) -> None: subpath_to = os.path.join(to_path_dir, subpath) if not os.path.exists(subpath_to): shutil.copy(subpath_from, to_path_dir) - - -# tempfile.TemporaryDirectory : https://docs.python.org/3.11/library/tempfile.html#tempfile.TemporaryDirectory -# It is not possible to keep directory content only when exception occurs. -class temporary_directory: - """ - Class to automatically create and remove temporary directory. - - In case of exception, do not remove the folder. - """ - - def __init__(self, path: str, user: str, group: str): - self._path = path - self._user = user - self._group = group - - def __enter__(self) -> None: - if os.path.exists(self._path): - shutil.rmtree(self._path) - - os.makedirs(self._path) - shutil.chown(self._path, self._user, self._group) - - def __exit__(self, exc_type, exc_value, traceback): - if exc_type is not None: - logging.warning( - "Dont remove tmp dir {} due to exception. {}: {}", - self._path, - exc_type.__name__, - exc_value, - ) - return - shutil.rmtree(self._path)