From fa4c46173ff8fb317490faea0b17ac5cc358716a Mon Sep 17 00:00:00 2001
From: poturaevpetr <poturaevpetr@gmail.com>
Date: Thu, 25 Dec 2025 21:57:30 +0500
Subject: [PATCH] add write audio to db

---
 autoLoader/database/__init__.py | 28 ++++++++++++----
 autoLoader/loader/loader.py     | 58 +++++++++++++++++++++++++--------
 requirements.txt                |  3 +-
 3 files changed, 68 insertions(+), 21 deletions(-)

diff --git a/autoLoader/database/__init__.py b/autoLoader/database/__init__.py
index 7ce0921..218f9be 100644
--- a/autoLoader/database/__init__.py
+++ b/autoLoader/database/__init__.py
@@ -1,7 +1,8 @@
 from sqlalchemy import create_engine
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import sessionmaker, Session
 from autoLoader.config import DATABASE_URL
+from contextlib import contextmanager
 
 # Создание engine
 engine = create_engine(
@@ -12,11 +13,11 @@ engine = create_engine(
 # SessionLocal
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
-# Base
-Base = declarative_base()
+# Base (используем тот же Base, что и в apiApp.database)
+from apiApp.database import Base
 
 
-# Зависимость для получения сессии БД
+# Зависимость для получения сессии БД (для FastAPI)
 def get_db():
     db = SessionLocal()
     try:
@@ -25,7 +26,22 @@ def get_db():
         db.close()
 
 
-# from apiApp.database.Operator import Operator
+# Контекстный менеджер для использования в loader
+@contextmanager
+def get_db_session():
+    """Контекстный менеджер для работы с БД в loader"""
+    db = SessionLocal()
+    try:
+        yield db
+        db.commit()
+    except Exception as e:
+        db.rollback()
+        raise e
+    finally:
+        db.close()
+
+
+# Импортируем модели из apiApp.database
 from apiApp.database.Audio import Audio
 from apiApp.database.AiConclusion import AiConclusion
-from apiApp.database.ConclusionVersion import ConclusionVersion
\ No newline at end of file
+from apiApp.database.ConclusionVersion import ConclusionVersion
diff --git a/autoLoader/loader/loader.py b/autoLoader/loader/loader.py
index f250484..743eb3a 100644
--- a/autoLoader/loader/loader.py
+++ b/autoLoader/loader/loader.py
@@ -1,9 +1,10 @@
 from autoLoader.config import FILESAPTH
 from autoLoader.loader import ConnectorSFTP
 import datetime, os
-from autoLoader.database import *
+from autoLoader.database import Audio, get_db_session
 
 local_path = os.path.join(os.getcwd(), FILESAPTH)
+
 class Loader():
     def __init__(self):
         self.call_types = ['in']
@@ -15,7 +16,6 @@ class Loader():
         return False
 
     def load(self):
-
         date_now = datetime.datetime.now()# - datetime.timedelta(days=1)
         remote_path = f"/{date_now.strftime('%Y/%m/%d')}"
 
@@ -29,30 +29,60 @@ class Loader():
             print("Не удалось подключиться к SFTP. Завершение работы.")
             exit(1)
 
-        db = get_db()
-        
         try:
             listdir = sftp_client.listdir()
             os.makedirs(local_path, exist_ok = True)
-            diskdir = os.listdir(local_path)
-    
+
             for file in listdir:
                 if self.filter_call(filename=file):
                     remote_file = f"{file}".lstrip('/')
                     filepath = os.path.join(local_path, file)
-                    try:
-                        sftp_client.get(remote_file, filepath)  # Скачиваем файл
-                        audio = Audio()
-                        audio.index_date = datetime.datetime.now()
-                        audio.filename = file
-                        db.session.add(audio)
-                        db.session.commit()
 
+                    # Проверяем, существует ли файл локально
+                    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_path = filepath
+                            audio.file_size = file_size
+
+                            db.add(audio)
+                            # commit произойдёт автоматически при выходе из контекста
+
+                        print(f"✅ Файл {file} сохранён в БД")
+
                     except Exception as e:
-                        print(f"Ошибка при скачивании файла {remote_file}: {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()
+
diff --git a/requirements.txt b/requirements.txt
index 6c293fa..49b4215 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,4 +4,5 @@ sqlalchemy==2.0.35
 pydantic==2.9.2
 python-multipart==0.0.12
 aiofiles==24.1.0
-psycopg2-binary
\ No newline at end of file
+psycopg2-binary
+paramiko
\ No newline at end of file