src/components/storage.py - master - alsoalgo/CTF

src/components/storage.py
167 строк6 KB
alsoalgo
some improvements
14 дек 2025, 22:23
14 дек 2025, 22:2337d2ec0
"""
Система хранения заявок в файлах
"""
import json
import os
from pathlib import Path
from typing import List, Optional, Dict, Any
from datetime import datetime

from src.models.application import Application, ApplicationStatus


class ApplicationStorage:
    """
    Класс для хранения заявок в файловой системе.
    Каждая заявка сохраняется в отдельный JSON файл.
    """
    
    def __init__(self, storage_dir: str = "storage"):
        """
        Инициализация хранилища.
        
        :param storage_dir: Директория для хранения заявок
        """
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(exist_ok=True)
    
    def save_application(self, application: Application) -> bool:
        """
        Сохраняет заявку в файл.
        
        :param application: Объект Application
        :return: True если успешно сохранено
        """
        try:
            # Обновляем время изменения
            application.updated_at = datetime.now()
            
            # Преобразуем в словарь
            data = application.to_dict()
            
            # Сохраняем в JSON файл
            file_path = self.storage_dir / f"{application.application_id}.json"
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            
            return True
        except Exception as e:
            print(f"Error saving application: {e}")
            return False
    
    def load_application(self, application_id: str) -> Optional[Application]:
        """
        Загружает заявку из файла.
        
        :param application_id: ID заявки
        :return: Объект Application или None
        """
        try:
            file_path = self.storage_dir / f"{application_id}.json"
            if not file_path.exists():
                return None
            
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            return Application.from_dict(data)
        except Exception as e:
            print(f"Error loading application: {e}")
            return None
    
    def list_applications(self, status_filter: Optional[ApplicationStatus] = None) -> List[Application]:
        """
        Возвращает список всех заявок.
        
        :param status_filter: Фильтр по статусу (опционально)
        :return: Список заявок
        """
        applications = []
        
        try:
            for file_path in self.storage_dir.glob("*.json"):
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                    
                    app = Application.from_dict(data)
                    
                    # Применяем фильтр по статусу если указан
                    if status_filter is None or app.status == status_filter:
                        applications.append(app)
                except Exception as e:
                    print(f"Error loading application from {file_path}: {e}")
                    continue
            
            # Сортируем по дате обновления (новые первыми)
            applications.sort(key=lambda x: x.updated_at, reverse=True)
            
        except Exception as e:
            print(f"Error listing applications: {e}")
        
        return applications
    
    def delete_application(self, application_id: str) -> bool:
        """
        Удаляет заявку.
        
        :param application_id: ID заявки
        :return: True если успешно удалено
        """
        try:
            file_path = self.storage_dir / f"{application_id}.json"
            if file_path.exists():
                file_path.unlink()
                return True
            return False
        except Exception as e:
            print(f"Error deleting application: {e}")
            return False
    
    def search_applications(
        self, 
        query: str,
        search_fields: Optional[List[str]] = None
    ) -> List[Application]:
        """
        Поиск заявок по текстовому запросу.
        
        :param query: Поисковый запрос
        :param search_fields: Поля для поиска (по умолчанию: supplier, off_taker, product)
        :return: Список найденных заявок
        """
        if search_fields is None:
            search_fields = ['supplier', 'off_taker', 'product', 'application_id']
        
        query_lower = query.lower()
        all_applications = self.list_applications()
        results = []
        
        for app in all_applications:
            for field in search_fields:
                if hasattr(app, field):
                    value = str(getattr(app, field) or '').lower()
                    if query_lower in value:
                        results.append(app)
                        break
        
        return results
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        Возвращает статистику по заявкам.
        
        :return: Словарь со статистикой
        """
        all_applications = self.list_applications()
        
        stats = {
            'total': len(all_applications),
            'by_status': {}
        }
        
        for status in ApplicationStatus:
            count = len([app for app in all_applications if app.status == status])
            stats['by_status'][status.value] = count
        
        return stats