# [DEF:superset_tool.utils.fileio:Module] # # @SEMANTICS: file, io, zip, yaml, temp, archive, utility # @PURPOSE: Предоставляет набор утилит для управления файловыми операциями, включая работу с временными файлами, архивами ZIP, файлами YAML и очистку директорий. # @LAYER: Infra # @RELATION: DEPENDS_ON -> superset_tool.exceptions # @RELATION: DEPENDS_ON -> superset_tool.utils.logger # @RELATION: DEPENDS_ON -> pyyaml # @PUBLIC_API: create_temp_file, remove_empty_directories, read_dashboard_from_disk, calculate_crc32, RetentionPolicy, archive_exports, save_and_unpack_dashboard, update_yamls, create_dashboard_export, sanitize_filename, get_filename_from_headers, consolidate_archive_folders # [SECTION: IMPORTS] import os import re import zipfile from pathlib import Path from typing import Any, Optional, Tuple, Dict, List, Union, LiteralString, Generator from contextlib import contextmanager import tempfile from datetime import date, datetime import glob import shutil import zlib from dataclasses import dataclass import yaml from superset_tool.exceptions import InvalidZipFormatError from superset_tool.utils.logger import SupersetLogger # [/SECTION] # [DEF:create_temp_file:Function] # @PURPOSE: Контекстный менеджер для создания временного файла или директории с гарантированным удалением. # @PRE: suffix должен быть строкой, определяющей тип ресурса. # @POST: Временный ресурс создан и путь к нему возвращен; ресурс удален после выхода из контекста. # @PARAM: content (Optional[bytes]) - Бинарное содержимое для записи во временный файл. # @PARAM: suffix (str) - Суффикс ресурса. Если `.dir`, создается директория. # @PARAM: mode (str) - Режим записи в файл (e.g., 'wb'). # @PARAM: logger (Optional[SupersetLogger]) - Экземпляр логгера. # @YIELDS: Path - Путь к временному ресурсу. # @THROW: IOError - При ошибках создания ресурса. @contextmanager def create_temp_file(content: Optional[bytes] = None, suffix: str = ".zip", mode: str = 'wb', dry_run = False, logger: Optional[SupersetLogger] = None) -> Generator[Path, None, None]: logger = logger or SupersetLogger(name="fileio") with logger.belief_scope("Create temporary resource"): resource_path = None is_dir = suffix.startswith('.dir') try: if is_dir: with tempfile.TemporaryDirectory(suffix=suffix) as temp_dir: resource_path = Path(temp_dir) logger.debug("[create_temp_file][State] Created temporary directory: %s", resource_path) yield resource_path else: fd, temp_path_str = tempfile.mkstemp(suffix=suffix) resource_path = Path(temp_path_str) os.close(fd) if content: resource_path.write_bytes(content) logger.debug("[create_temp_file][State] Created temporary file: %s", resource_path) yield resource_path finally: if resource_path and resource_path.exists() and not dry_run: try: if resource_path.is_dir(): shutil.rmtree(resource_path) logger.debug("[create_temp_file][Cleanup] Removed temporary directory: %s", resource_path) else: resource_path.unlink() logger.debug("[create_temp_file][Cleanup] Removed temporary file: %s", resource_path) except OSError as e: logger.error("[create_temp_file][Failure] Error during cleanup of %s: %s", resource_path, e) # [/DEF:create_temp_file:Function] # [DEF:remove_empty_directories:Function] # @PURPOSE: Рекурсивно удаляет все пустые поддиректории, начиная с указанного пути. # @PRE: root_dir должен быть путем к существующей директории. # @POST: Все пустые поддиректории удалены, возвращено их количество. # @PARAM: root_dir (str) - Путь к корневой директории для очистки. # @PARAM: logger (Optional[SupersetLogger]) - Экземпляр логгера. # @RETURN: int - Количество удаленных директорий. def remove_empty_directories(root_dir: str, logger: Optional[SupersetLogger] = None) -> int: logger = logger or SupersetLogger(name="fileio") with logger.belief_scope(f"Remove empty directories in {root_dir}"): logger.info("[remove_empty_directories][Enter] Starting cleanup of empty directories in %s", root_dir) removed_count = 0 if not os.path.isdir(root_dir): logger.error("[remove_empty_directories][Failure] Directory not found: %s", root_dir) return 0 for current_dir, _, _ in os.walk(root_dir, topdown=False): if not os.listdir(current_dir): try: os.rmdir(current_dir) removed_count += 1 logger.info("[remove_empty_directories][State] Removed empty directory: %s", current_dir) except OSError as e: logger.error("[remove_empty_directories][Failure] Failed to remove %s: %s", current_dir, e) logger.info("[remove_empty_directories][Exit] Removed %d empty directories.", removed_count) return removed_count # [/DEF:remove_empty_directories:Function] # [DEF:read_dashboard_from_disk:Function] # @PURPOSE: Читает бинарное содержимое файла с диска. # @PRE: file_path должен указывать на существующий файл. # @POST: Возвращает байты содержимого и имя файла. # @PARAM: file_path (str) - Путь к файлу. # @PARAM: logger (Optional[SupersetLogger]) - Экземпляр логгера. # @RETURN: Tuple[bytes, str] - Кортеж (содержимое, имя файла). # @THROW: FileNotFoundError - Если файл не найден. def read_dashboard_from_disk(file_path: str, logger: Optional[SupersetLogger] = None) -> Tuple[bytes, str]: logger = logger or SupersetLogger(name="fileio") with logger.belief_scope(f"Read dashboard from {file_path}"): path = Path(file_path) assert path.is_file(), f"Файл дашборда не найден: {file_path}" logger.info("[read_dashboard_from_disk][Enter] Reading file: %s", file_path) content = path.read_bytes() if not content: logger.warning("[read_dashboard_from_disk][Warning] File is empty: %s", file_path) return content, path.name # [/DEF:read_dashboard_from_disk:Function] # [DEF:calculate_crc32:Function] # @PURPOSE: Вычисляет контрольную сумму CRC32 для файла. # @PRE: file_path должен быть объектом Path к существующему файлу. # @POST: Возвращает 8-значную hex-строку CRC32. # @PARAM: file_path (Path) - Путь к файлу. # @RETURN: str - 8-значное шестнадцатеричное представление CRC32. # @THROW: IOError - При ошибках чтения файла. def calculate_crc32(file_path: Path) -> str: logger = SupersetLogger(name="fileio") with logger.belief_scope(f"Calculate CRC32 for {file_path}"): with open(file_path, 'rb') as f: crc32_value = zlib.crc32(f.read()) return f"{crc32_value:08x}" # [/DEF:calculate_crc32:Function] # [SECTION: DATA_CLASSES] # [DEF:RetentionPolicy:DataClass] # @PURPOSE: Определяет политику хранения для архивов (ежедневные, еженедельные, ежемесячные). @dataclass class RetentionPolicy: daily: int = 7 weekly: int = 4 monthly: int = 12 # [/DEF:RetentionPolicy:DataClass] # [/SECTION] # [DEF:archive_exports:Function] # @PURPOSE: Управляет архивом экспортированных файлов, применяя политику хранения и дедупликацию. # @PRE: output_dir должен быть путем к существующей директории. # @POST: Старые или дублирующиеся архивы удалены согласно политике. # @RELATION: CALLS -> apply_retention_policy # @RELATION: CALLS -> calculate_crc32 # @PARAM: output_dir (str) - Директория с архивами. # @PARAM: policy (RetentionPolicy) - Политика хранения. # @PARAM: deduplicate (bool) - Флаг для включения удаления дубликатов по CRC32. # @PARAM: logger (Optional[SupersetLogger]) - Экземпляр логгера. def archive_exports(output_dir: str, policy: RetentionPolicy, deduplicate: bool = False, logger: Optional[SupersetLogger] = None) -> None: logger = logger or SupersetLogger(name="fileio") with logger.belief_scope(f"Archive exports in {output_dir}"): output_path = Path(output_dir) if not output_path.is_dir(): logger.warning("[archive_exports][Skip] Archive directory not found: %s", output_dir) return logger.info("[archive_exports][Enter] Managing archive in %s", output_dir) # 1. Collect all zip files zip_files = list(output_path.glob("*.zip")) if not zip_files: logger.info("[archive_exports][State] No zip files found in %s", output_dir) return # 2. Deduplication if deduplicate: logger.info("[archive_exports][State] Starting deduplication...") checksums = {} files_to_remove = [] # Sort by modification time (newest first) to keep the latest version zip_files.sort(key=lambda f: f.stat().st_mtime, reverse=True) for file_path in zip_files: try: crc = calculate_crc32(file_path) if crc in checksums: files_to_remove.append(file_path) logger.debug("[archive_exports][State] Duplicate found: %s (same as %s)", file_path.name, checksums[crc].name) else: checksums[crc] = file_path except Exception as e: logger.error("[archive_exports][Failure] Failed to calculate CRC32 for %s: %s", file_path, e) for f in files_to_remove: try: f.unlink() zip_files.remove(f) logger.info("[archive_exports][State] Removed duplicate: %s", f.name) except OSError as e: logger.error("[archive_exports][Failure] Failed to remove duplicate %s: %s", f, e) # 3. Retention Policy files_with_dates = [] for file_path in zip_files: # Try to extract date from filename # Pattern: ..._YYYYMMDD_HHMMSS.zip or ..._YYYYMMDD.zip match = re.search(r'_(\d{8})_', file_path.name) file_date = None if match: try: date_str = match.group(1) file_date = datetime.strptime(date_str, "%Y%m%d").date() except ValueError: pass if not file_date: # Fallback to modification time file_date = datetime.fromtimestamp(file_path.stat().st_mtime).date() files_with_dates.append((file_path, file_date)) files_to_keep = apply_retention_policy(files_with_dates, policy, logger) for file_path, _ in files_with_dates: if file_path not in files_to_keep: try: file_path.unlink() logger.info("[archive_exports][State] Removed by retention policy: %s", file_path.name) except OSError as e: logger.error("[archive_exports][Failure] Failed to remove %s: %s", file_path, e) # [/DEF:archive_exports:Function] # [DEF:apply_retention_policy:Function] # @PURPOSE: (Helper) Применяет политику хранения к списку файлов, возвращая те, что нужно сохранить. # @PRE: files_with_dates is a list of (Path, date) tuples. # @POST: Returns a set of files to keep. # @PARAM: files_with_dates (List[Tuple[Path, date]]) - Список файлов с датами. # @PARAM: policy (RetentionPolicy) - Политика хранения. # @PARAM: logger (SupersetLogger) - Логгер. # @RETURN: set - Множество путей к файлам, которые должны быть сохранены. def apply_retention_policy(files_with_dates: List[Tuple[Path, date]], policy: RetentionPolicy, logger: SupersetLogger) -> set: with logger.belief_scope("Apply retention policy"): # Сортируем по дате (от новой к старой) sorted_files = sorted(files_with_dates, key=lambda x: x[1], reverse=True) # Словарь для хранения файлов по категориям daily_files = [] weekly_files = [] monthly_files = [] today = date.today() for file_path, file_date in sorted_files: # Ежедневные if (today - file_date).days < policy.daily: daily_files.append(file_path) # Еженедельные elif (today - file_date).days < policy.weekly * 7: weekly_files.append(file_path) # Ежемесячные elif (today - file_date).days < policy.monthly * 30: monthly_files.append(file_path) # Возвращаем множество файлов, которые нужно сохранить files_to_keep = set() files_to_keep.update(daily_files) files_to_keep.update(weekly_files[:policy.weekly]) files_to_keep.update(monthly_files[:policy.monthly]) logger.debug("[apply_retention_policy][State] Keeping %d files according to retention policy", len(files_to_keep)) return files_to_keep # [/DEF:apply_retention_policy:Function] # [DEF:save_and_unpack_dashboard:Function] # @PURPOSE: Сохраняет бинарное содержимое ZIP-архива на диск и опционально распаковывает его. # @PRE: zip_content должен быть байтами валидного ZIP-архива. # @POST: ZIP-файл сохранен, и если unpack=True, он распакован в output_dir. # @PARAM: zip_content (bytes) - Содержимое ZIP-архива. # @PARAM: output_dir (Union[str, Path]) - Директория для сохранения. # @PARAM: unpack (bool) - Флаг, нужно ли распаковывать архив. # @PARAM: original_filename (Optional[str]) - Исходное имя файла для сохранения. # @PARAM: logger (Optional[SupersetLogger]) - Экземпляр логгера. # @RETURN: Tuple[Path, Optional[Path]] - Путь к ZIP-файлу и, если применимо, путь к директории с распаковкой. # @THROW: InvalidZipFormatError - При ошибке формата ZIP. def save_and_unpack_dashboard(zip_content: bytes, output_dir: Union[str, Path], unpack: bool = False, original_filename: Optional[str] = None, logger: Optional[SupersetLogger] = None) -> Tuple[Path, Optional[Path]]: logger = logger or SupersetLogger(name="fileio") with logger.belief_scope("Save and unpack dashboard"): logger.info("[save_and_unpack_dashboard][Enter] Processing dashboard. Unpack: %s", unpack) try: output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) zip_name = sanitize_filename(original_filename) if original_filename else f"dashboard_export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip" zip_path = output_path / zip_name zip_path.write_bytes(zip_content) logger.info("[save_and_unpack_dashboard][State] Dashboard saved to: %s", zip_path) if unpack: with zipfile.ZipFile(zip_path, 'r') as zip_ref: zip_ref.extractall(output_path) logger.info("[save_and_unpack_dashboard][State] Dashboard unpacked to: %s", output_path) return zip_path, output_path return zip_path, None except zipfile.BadZipFile as e: logger.error("[save_and_unpack_dashboard][Failure] Invalid ZIP archive: %s", e) raise InvalidZipFormatError(f"Invalid ZIP file: {e}") from e # [/DEF:save_and_unpack_dashboard:Function] # [DEF:update_yamls:Function] # @PURPOSE: Обновляет конфигурации в YAML-файлах, заменяя значения или применяя regex. # @PRE: path должен быть существующей директорией. # @POST: Все YAML файлы в директории обновлены согласно переданным параметрам. # @RELATION: CALLS -> _update_yaml_file # @THROW: FileNotFoundError - Если `path` не существует. # @PARAM: db_configs (Optional[List[Dict]]) - Список конфигураций для замены. # @PARAM: path (str) - Путь к директории с YAML файлами. # @PARAM: regexp_pattern (Optional[LiteralString]) - Паттерн для поиска. # @PARAM: replace_string (Optional[LiteralString]) - Строка для замены. # @PARAM: logger (Optional[SupersetLogger]) - Экземпляр логгера. def update_yamls(db_configs: Optional[List[Dict[str, Any]]] = None, path: str = "dashboards", regexp_pattern: Optional[LiteralString] = None, replace_string: Optional[LiteralString] = None, logger: Optional[SupersetLogger] = None) -> None: logger = logger or SupersetLogger(name="fileio") with logger.belief_scope("Update YAML configurations"): logger.info("[update_yamls][Enter] Starting YAML configuration update.") dir_path = Path(path) assert dir_path.is_dir(), f"Путь {path} не существует или не является директорией" configs: List[Dict[str, Any]] = db_configs or [] for file_path in dir_path.rglob("*.yaml"): _update_yaml_file(file_path, configs, regexp_pattern, replace_string, logger) # [/DEF:update_yamls:Function] # [DEF:_update_yaml_file:Function] # @PURPOSE: (Helper) Обновляет один YAML файл. # @PRE: file_path должен быть объектом Path к существующему YAML файлу. # @POST: Файл обновлен согласно переданным конфигурациям или регулярному выражению. # @PARAM: file_path (Path) - Путь к файлу. # @PARAM: db_configs (List[Dict]) - Конфигурации. # @PARAM: regexp_pattern (Optional[str]) - Паттерн. # @PARAM: replace_string (Optional[str]) - Замена. # @PARAM: logger (SupersetLogger) - Логгер. def _update_yaml_file(file_path: Path, db_configs: List[Dict[str, Any]], regexp_pattern: Optional[str], replace_string: Optional[str], logger: SupersetLogger) -> None: with logger.belief_scope(f"Update YAML file: {file_path}"): # Читаем содержимое файла try: with open(file_path, 'r', encoding='utf-8') as f: content = f.read() except Exception as e: logger.error("[_update_yaml_file][Failure] Failed to read %s: %s", file_path, e) return # Если задан pattern и replace_string, применяем замену по регулярному выражению if regexp_pattern and replace_string: try: new_content = re.sub(regexp_pattern, replace_string, content) if new_content != content: with open(file_path, 'w', encoding='utf-8') as f: f.write(new_content) logger.info("[_update_yaml_file][State] Updated %s using regex pattern", file_path) except Exception as e: logger.error("[_update_yaml_file][Failure] Error applying regex to %s: %s", file_path, e) # Если заданы конфигурации, заменяем значения (поддержка old/new) if db_configs: try: # Прямой текстовый заменитель для старых/новых значений, чтобы сохранить структуру файла modified_content = content for cfg in db_configs: # Ожидаем структуру: {'old': {...}, 'new': {...}} old_cfg = cfg.get('old', {}) new_cfg = cfg.get('new', {}) for key, old_val in old_cfg.items(): if key in new_cfg: new_val = new_cfg[key] # Заменяем только точные совпадения старого значения в тексте YAML, используя ключ для контекста if isinstance(old_val, str): # Ищем паттерн: key: "value" или key: value key_pattern = re.escape(key) val_pattern = re.escape(old_val) # Группы: 1=ключ+разделитель, 2=открывающая кавычка (опц), 3=значение, 4=закрывающая кавычка (опц) pattern = rf'({key_pattern}\s*:\s*)(["\']?)({val_pattern})(["\']?)' # [DEF:replacer:Function] # @PURPOSE: Функция замены, сохраняющая кавычки если они были. # @PRE: match должен быть объектом совпадения регулярного выражения. # @POST: Возвращает строку с новым значением, сохраняя префикс и кавычки. def replacer(match): with logger.belief_scope("replacer"): prefix = match.group(1) quote_open = match.group(2) quote_close = match.group(4) return f"{prefix}{quote_open}{new_val}{quote_close}" # [/DEF:replacer:Function] modified_content = re.sub(pattern, replacer, modified_content) logger.info("[_update_yaml_file][State] Replaced '%s' with '%s' for key %s in %s", old_val, new_val, key, file_path) # Записываем обратно изменённый контент без парсинга YAML, сохраняем оригинальное форматирование with open(file_path, 'w', encoding='utf-8') as f: f.write(modified_content) except Exception as e: logger.error("[_update_yaml_file][Failure] Error performing raw replacement in %s: %s", file_path, e) # [/DEF:_update_yaml_file:Function] # [DEF:create_dashboard_export:Function] # @PURPOSE: Создает ZIP-архив из указанных исходных путей. # @PRE: source_paths должен содержать существующие пути. # @POST: ZIP-архив создан по пути zip_path. # @PARAM: zip_path (Union[str, Path]) - Путь для сохранения ZIP архива. # @PARAM: source_paths (List[Union[str, Path]]) - Список исходных путей для архивации. # @PARAM: exclude_extensions (Optional[List[str]]) - Список расширений для исключения. # @PARAM: logger (Optional[SupersetLogger]) - Экземпляр логгера. # @RETURN: bool - `True` при успехе, `False` при ошибке. def create_dashboard_export(zip_path: Union[str, Path], source_paths: List[Union[str, Path]], exclude_extensions: Optional[List[str]] = None, logger: Optional[SupersetLogger] = None) -> bool: logger = logger or SupersetLogger(name="fileio") with logger.belief_scope(f"Create dashboard export: {zip_path}"): logger.info("[create_dashboard_export][Enter] Packing dashboard: %s -> %s", source_paths, zip_path) try: exclude_ext = [ext.lower() for ext in exclude_extensions or []] with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: for src_path_str in source_paths: src_path = Path(src_path_str) assert src_path.exists(), f"Путь не найден: {src_path}" for item in src_path.rglob('*'): if item.is_file() and item.suffix.lower() not in exclude_ext: arcname = item.relative_to(src_path.parent) zipf.write(item, arcname) logger.info("[create_dashboard_export][Exit] Archive created: %s", zip_path) return True except (IOError, zipfile.BadZipFile, AssertionError) as e: logger.error("[create_dashboard_export][Failure] Error: %s", e, exc_info=True) return False # [/DEF:create_dashboard_export:Function] # [DEF:sanitize_filename:Function] # @PURPOSE: Очищает строку от символов, недопустимых в именах файлов. # @PRE: filename должен быть строкой. # @POST: Возвращает строку без спецсимволов. # @PARAM: filename (str) - Исходное имя файла. # @RETURN: str - Очищенная строка. def sanitize_filename(filename: str) -> str: logger = SupersetLogger(name="fileio") with logger.belief_scope(f"Sanitize filename: {filename}"): return re.sub(r'[\\/*?:"<>|]', "_", filename).strip() # [/DEF:sanitize_filename:Function] # [DEF:get_filename_from_headers:Function] # @PURPOSE: Извлекает имя файла из HTTP заголовка 'Content-Disposition'. # @PRE: headers должен быть словарем заголовков. # @POST: Возвращает имя файла или None, если заголовок отсутствует. # @PARAM: headers (dict) - Словарь HTTP заголовков. # @RETURN: Optional[str] - Имя файла or `None`. def get_filename_from_headers(headers: dict) -> Optional[str]: logger = SupersetLogger(name="fileio") with logger.belief_scope("Get filename from headers"): content_disposition = headers.get("Content-Disposition", "") if match := re.search(r'filename="?([^"]+)"?', content_disposition): return match.group(1).strip() return None # [/DEF:get_filename_from_headers:Function] # [DEF:consolidate_archive_folders:Function] # @PURPOSE: Консолидирует директории архивов на основе общего слага в имени. # @PRE: root_directory должен быть объектом Path к существующей директории. # @POST: Директории с одинаковым префиксом объединены в одну. # @THROW: TypeError, ValueError - Если `root_directory` невалиден. # @PARAM: root_directory (Path) - Корневая директория для консолидации. # @PARAM: logger (Optional[SupersetLogger]) - Экземпляр логгера. def consolidate_archive_folders(root_directory: Path, logger: Optional[SupersetLogger] = None) -> None: logger = logger or SupersetLogger(name="fileio") with logger.belief_scope(f"Consolidate archives in {root_directory}"): assert isinstance(root_directory, Path), "root_directory must be a Path object." assert root_directory.is_dir(), "root_directory must be an existing directory." logger.info("[consolidate_archive_folders][Enter] Consolidating archives in %s", root_directory) # Собираем все директории с архивами archive_dirs = [] for item in root_directory.iterdir(): if item.is_dir(): # Проверяем, есть ли в директории ZIP-архивы if any(item.glob("*.zip")): archive_dirs.append(item) # Группируем по слагу (части имени до первого '_') slug_groups = {} for dir_path in archive_dirs: dir_name = dir_path.name slug = dir_name.split('_')[0] if '_' in dir_name else dir_name if slug not in slug_groups: slug_groups[slug] = [] slug_groups[slug].append(dir_path) # Для каждой группы консолидируем for slug, dirs in slug_groups.items(): if len(dirs) <= 1: continue # Создаем целевую директорию target_dir = root_directory / slug target_dir.mkdir(exist_ok=True) logger.info("[consolidate_archive_folders][State] Consolidating %d directories under %s", len(dirs), target_dir) # Перемещаем содержимое for source_dir in dirs: if source_dir == target_dir: continue for item in source_dir.iterdir(): dest_item = target_dir / item.name try: if item.is_dir(): shutil.move(str(item), str(dest_item)) else: shutil.move(str(item), str(dest_item)) except Exception as e: logger.error("[consolidate_archive_folders][Failure] Failed to move %s to %s: %s", item, dest_item, e) # Удаляем исходную директорию try: source_dir.rmdir() logger.info("[consolidate_archive_folders][State] Removed source directory: %s", source_dir) except Exception as e: logger.error("[consolidate_archive_folders][Failure] Failed to remove source directory %s: %s", source_dir, e) # [/DEF:consolidate_archive_folders:Function] # [/DEF:superset_tool.utils.fileio:Module]