FileAudioAPI/autoLoader/loader/loader.py

from autoLoader.config import FILESAPTH
from autoLoader.loader.connector import ConnectorSFTP
import datetime, os
from autoLoader.database import Audio, get_db_session
from sqlalchemy import inspect

local_path = os.path.join(os.getcwd(), FILESAPTH)

class Loader():
    def __init__(self):
        self.call_types = ['in']
        self.loaded_files = []  # Список загруженных файлов

    def filter_call(self, filename: str):
        if filename.split("-")[0] in self.call_types:
            return True
        return False

    def check_database(self):
        """Проверяет существование таблиц в БД"""
        from autoLoader.database import engine

        inspector = inspect(engine)
        existing_tables = inspector.get_table_names()

        if 'audio' not in existing_tables:
            print("❌ Таблица 'audio' не существует в базе данных!")
            print("💡 Запустите 'python init_db.py' для создания таблиц")
            return False

        return True

    def send_to_recognition(self, limit: int = None):
        """
        Отправляет загруженные файлы на распознавание

        Args:
            limit: Максимальное количество файлов для отправки
        """
        if not self.loaded_files:
            print("⏭️ Нет файлов для отправки на распознавание")
            return

        from autoLoader.loader import RecognitionChecker

        print(f"\n🎤 Отправка {len(self.loaded_files)} файлов на распознавание...")

        checker = RecognitionChecker()

        # Проверяем доступность API
        if not checker.check_api_availability():
            print("❌ GigaAM API недоступен. Пропускаем распознавание.")
            return

        # Проверяем БД
        if not checker.check_database():
            print("❌ База данных не готова. Пропускаем распознавание.")
            return

        # Фильтруем только что загруженные файлы без заключения
        files_to_send = []
        for audio in self.loaded_files:
            # Проверяем, есть ли у файла заключение
            with get_db_session() as db:
                from autoLoader.database import AiConclusion
                existing_conclusion = db.query(AiConclusion).filter(
                    AiConclusion.audio_id == audio.id
                ).first()

                if existing_conclusion:
                    print(f"⏭️ Файл {audio.filename} уже имеет заключение")
                    continue
                else:
                    # Конвертируем SQLAlchemy объект в словарь
                    files_to_send.append({
                        'id': str(audio.id),
                        'filename': audio.filename,
                        'file_size': audio.file_size,
                        'index_date': audio.index_date.isoformat() if audio.index_date else None
                    })

        if not files_to_send:
            print("✅ Все файлы уже имеют заключения")
            return

        # Отправляем файлы параллельно
        print(f"🚀 Отправка {len(files_to_send)} файлов параллельно...")
        results = checker.send_to_recognition_parallel(files_to_send)

        print(f"✅ Отправлено {results['sent']} из {results['total']} файлов на распознавание")

    def load(self):
        # Проверяем БД перед началом работы
        if not self.check_database():
            exit(1)

        date_now = datetime.datetime.now()# - datetime.timedelta(days=1)
        remote_path = f"/{date_now.strftime('%Y/%m/%d')}"

        connector = ConnectorSFTP()
        connector.connect(remote_path=remote_path)

        sftp_client = connector.sftp
        ssh_client = connector.ssh

        if sftp_client is None or ssh_client is None:
            print("Не удалось подключиться к SFTP. Завершение работы.")
            exit(1)

        try:
            listdir = sftp_client.listdir()
            os.makedirs(local_path, exist_ok = True)

            for file in listdir:
                if self.filter_call(filename=file):
                    remote_file = f"{file}".lstrip('/')
                    filepath = os.path.join(local_path, file)

                    # Проверяем, существует ли файл локально
                    if os.path.exists(filepath):
                        print(f"Файл уже существует локально: {file}")
                        continue

                    try:
                        # Скачиваем файл
                        sftp_client.get(remote_file, filepath)
                        print(f"📥 Скачан файл: {remote_file}")

                        # Получаем размер файла
                        file_size = os.path.getsize(filepath)

                        # Сохраняем в БД через контекстный менеджер
                        with get_db_session() as db:
                            # Проверяем, есть ли уже такой файл в БД
                            existing_audio = db.query(Audio).filter(Audio.filename == file).first()
                            if existing_audio:
                                print(f"⏭️ Файл {file} уже есть в БД, пропускаем")
                                continue

                            # Создаём новую запись
                            audio = Audio()
                            audio.index_date = datetime.datetime.now()
                            audio.filename = file
                            audio.file_size = file_size

                            db.add(audio)
                            db.flush()  # Чтобы получить audio.id

                            # Сохраняем в список для последующей отправки
                            self.loaded_files.append(audio)

                        print(f"✅ Файл {file} сохранён в БД")

                    except Exception as e:
                        print(f"❌ Ошибка при обработке файла {remote_file}: {e}")
                        # Если файл скачался, но ошибка в БД - удаляем файл
                        if os.path.exists(filepath):
                            try:
                                os.remove(filepath)
                                print(f"🗑️ Файл {file} удалён из-за ошибки")
                            except:
                                pass

        finally:
            # Закрываем соединения
            sftp_client.close()
            ssh_client.close()

        # После загрузки всех файлов отправляем на распознавание
        print(f"\n📊 Загружено файлов: {len(self.loaded_files)}")
        self.send_to_recognition()