metadata_manager.py - master - redcondevil/Test

metadata_manager.py
221 строка8 KB
redcondevil
Test
25 ноя 2025, 18:11
25 ноя 2025, 18:11923d4d9
import json
import os
import sqlite3
from pathlib import Path
from typing import Dict, Optional, List

class MetadataManager:
    def __init__(self, metadata_file='metadata.db'):
        self.metadata_file = metadata_file
        self.conn = None
        self.init_database()
    
    def init_database(self):
        """Инициализация SQLite базы данных"""
        self.conn = sqlite3.connect(self.metadata_file, check_same_thread=False)
        self.conn.execute('PRAGMA journal_mode=WAL')
        cursor = self.conn.cursor()
        
        # Создаем таблицу если она не существует
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS metadata (
                file_path TEXT PRIMARY KEY,
                tags TEXT,
                rating INTEGER DEFAULT 0,
                comment TEXT,
                source_url TEXT,
                last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # Проверяем существование колонки source_url и добавляем если нужно
        cursor.execute("PRAGMA table_info(metadata)")
        columns = [column[1] for column in cursor.fetchall()]
        
        if 'source_url' not in columns:
            print("🔄 Adding source_url column to metadata table...")
            cursor.execute('ALTER TABLE metadata ADD COLUMN source_url TEXT')
        
        cursor.execute('''
            CREATE INDEX IF NOT EXISTS idx_file_path 
            ON metadata(file_path)
        ''')
        
        self.conn.commit()
    
    def get_media_metadata(self, file_path: str) -> Dict:
        """Получение метаданных для медиафайла из базы данных"""
        cursor = self.conn.cursor()
        cursor.execute(
            'SELECT tags, rating, comment, source_url FROM metadata WHERE file_path = ?',
            (file_path,)
        )
        result = cursor.fetchone()
        
        if result:
            tags, rating, comment, source_url = result
            return {
                'tags': tags or self.extract_tags_from_path(file_path),
                'rating': rating or 0,
                'comment': comment or '',
                'source_url': source_url or ''
            }
        else:
            return {
                'tags': self.extract_tags_from_path(file_path),
                'rating': 0,
                'comment': '',
                'source_url': ''
            }
    
    def update_media_metadata(self, file_path: str, tags: str = None, 
                            rating: int = None, comment: str = None, source_url: str = None):
        """Обновление метаданных медиафайла"""
        cursor = self.conn.cursor()
        
        # Получаем текущие данные для частичного обновления
        current_data = self.get_media_metadata(file_path)
        
        # Используем UPSERT для атомарного обновления
        cursor.execute('''
            INSERT OR REPLACE INTO metadata 
            (file_path, tags, rating, comment, source_url, last_updated)
            VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
        ''', (
            file_path,
            tags if tags is not None else current_data['tags'],
            rating if rating is not None else current_data['rating'],
            comment if comment is not None else current_data['comment'],
            source_url if source_url is not None else current_data.get('source_url', '')
        ))
        
        self.conn.commit()
    
    def extract_tags_from_path(self, path: str) -> str:
        """Извлечение тегов из пути файла"""
        directory = os.path.dirname(path)
        tags = []
        
        # Исключаем системные папки
        excluded_folders = {'Users', 'home', 'Pictures', 'Images', 'photos', 
                          'Videos', 'videos', 'Downloads', 'Desktop', 'Documents'}
        
        for folder in directory.split(os.sep):
            folder = folder.strip()
            if (folder and folder not in excluded_folders and 
                not folder.startswith('.') and
                not folder.startswith('~')):
                if folder not in tags:
                    tags.append(folder)
        
        return ', '.join(tags)
    
    def batch_update_metadata(self, updates_dict: Dict):
        """Пакетное обновление метаданных"""
        cursor = self.conn.cursor()
        
        try:
            cursor.execute('BEGIN TRANSACTION')
            
            for file_path, metadata in updates_dict.items():
                cursor.execute('''
                    INSERT OR REPLACE INTO metadata 
                    (file_path, tags, rating, comment, source_url, last_updated)
                    VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
                ''', (
                    file_path,
                    metadata.get('tags', ''),
                    metadata.get('rating', 0),
                    metadata.get('comment', ''),
                    metadata.get('source_url', '')
                ))
            
            self.conn.commit()
            
        except Exception as e:
            self.conn.rollback()
            raise e
    
    def search_metadata(self, query: str, min_rating: int = 0, has_url: bool = False) -> Dict:
        """Расширенный поиск по метаданным"""
        cursor = self.conn.cursor()
        
        # Базовый запрос
        base_query = '''
            SELECT file_path, tags, rating, comment, source_url 
            FROM metadata 
            WHERE 1=1
        '''
        params = []
        
        # Поиск по тегам и комментариям
        if query:
            search_pattern = f'%{query}%'
            base_query += ' AND (tags LIKE ? OR comment LIKE ?)'
            params.extend([search_pattern, search_pattern])
        
        # Фильтр по рейтингу
        if min_rating > 0:
            base_query += ' AND rating >= ?'
            params.append(min_rating)
        
        # Фильтр по наличию URL
        if has_url:
            base_query += ' AND source_url != "" AND source_url IS NOT NULL'
        
        cursor.execute(base_query, params)
        
        results = {}
        for row in cursor.fetchall():
            file_path, tags, rating, comment, source_url = row
            
            # Проверяем существование файла
            if os.path.exists(file_path):
                results[file_path] = {
                    'tags': tags,
                    'rating': rating,
                    'comment': comment,
                    'source_url': source_url
                }
        
        return results

    def get_all_tagged_files(self) -> Dict:
        """Получение всех файлов с метаданными"""
        cursor = self.conn.cursor()
        cursor.execute('''
            SELECT file_path, tags, rating, comment, source_url 
            FROM metadata 
            WHERE file_path IS NOT NULL
        ''')
        
        results = {}
        for row in cursor.fetchall():
            file_path, tags, rating, comment, source_url = row
            if os.path.exists(file_path):
                results[file_path] = {
                    'tags': tags,
                    'rating': rating,
                    'comment': comment,
                    'source_url': source_url
                }
        
        return results

    def get_unique_tags(self) -> List[str]:
        """Получение уникальных тегов из базы данных"""
        cursor = self.conn.cursor()
        cursor.execute('SELECT tags FROM metadata WHERE tags != ""')
        
        all_tags = set()
        for (tags,) in cursor.fetchall():
            if tags:
                tag_list = [tag.strip() for tag in tags.split(',')]
                all_tags.update(tag_list)
        
        return sorted(list(all_tags))
    
    def __del__(self):
        """Закрытие соединения с базой данных"""
        if hasattr(self, 'conn') and self.conn:
            self.conn.close()