Browse Source

add write audio to db

dev
poturaevpetr 2 weeks ago
parent
commit
fa4c46173f
  1. 28
      autoLoader/database/__init__.py
  2. 58
      autoLoader/loader/loader.py
  3. 3
      requirements.txt

28
autoLoader/database/__init__.py

@ -1,7 +1,8 @@
from sqlalchemy import create_engine from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base 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 autoLoader.config import DATABASE_URL
from contextlib import contextmanager
# Создание engine # Создание engine
engine = create_engine( engine = create_engine(
@ -12,11 +13,11 @@ engine = create_engine(
# SessionLocal # SessionLocal
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
# Base # Base (используем тот же Base, что и в apiApp.database)
Base = declarative_base() from apiApp.database import Base
# Зависимость для получения сессии БД # Зависимость для получения сессии БД (для FastAPI)
def get_db(): def get_db():
db = SessionLocal() db = SessionLocal()
try: try:
@ -25,7 +26,22 @@ def get_db():
db.close() 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.Audio import Audio
from apiApp.database.AiConclusion import AiConclusion from apiApp.database.AiConclusion import AiConclusion
from apiApp.database.ConclusionVersion import ConclusionVersion from apiApp.database.ConclusionVersion import ConclusionVersion

58
autoLoader/loader/loader.py

@ -1,9 +1,10 @@
from autoLoader.config import FILESAPTH from autoLoader.config import FILESAPTH
from autoLoader.loader import ConnectorSFTP from autoLoader.loader import ConnectorSFTP
import datetime, os import datetime, os
from autoLoader.database import * from autoLoader.database import Audio, get_db_session
local_path = os.path.join(os.getcwd(), FILESAPTH) local_path = os.path.join(os.getcwd(), FILESAPTH)
class Loader(): class Loader():
def __init__(self): def __init__(self):
self.call_types = ['in'] self.call_types = ['in']
@ -15,7 +16,6 @@ class Loader():
return False return False
def load(self): def load(self):
date_now = datetime.datetime.now()# - datetime.timedelta(days=1) date_now = datetime.datetime.now()# - datetime.timedelta(days=1)
remote_path = f"/{date_now.strftime('%Y/%m/%d')}" remote_path = f"/{date_now.strftime('%Y/%m/%d')}"
@ -29,30 +29,60 @@ class Loader():
print("Не удалось подключиться к SFTP. Завершение работы.") print("Не удалось подключиться к SFTP. Завершение работы.")
exit(1) exit(1)
db = get_db()
try: try:
listdir = sftp_client.listdir() listdir = sftp_client.listdir()
os.makedirs(local_path, exist_ok = True) os.makedirs(local_path, exist_ok = True)
diskdir = os.listdir(local_path)
for file in listdir: for file in listdir:
if self.filter_call(filename=file): if self.filter_call(filename=file):
remote_file = f"{file}".lstrip('/') remote_file = f"{file}".lstrip('/')
filepath = os.path.join(local_path, file) 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}") 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: 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: finally:
# Закрываем соединения # Закрываем соединения
sftp_client.close() sftp_client.close()
ssh_client.close() ssh_client.close()

3
requirements.txt

@ -4,4 +4,5 @@ sqlalchemy==2.0.35
pydantic==2.9.2 pydantic==2.9.2
python-multipart==0.0.12 python-multipart==0.0.12
aiofiles==24.1.0 aiofiles==24.1.0
psycopg2-binary psycopg2-binary
paramiko
Loading…
Cancel
Save