Add new file
This commit is contained in:
Markov Andrey
356
core/services/postgres_service.py
Normal file
356
core/services/postgres_service.py
Normal file
@@ -0,0 +1,356 @@
|
|||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""
|
||||||
|
Сервис для работы с PostgreSQL (через asyncpg).
|
||||||
|
Обеспечивает пул соединений, CRUD операции для документов, истории диалогов, комнат MUC,
|
||||||
|
управление доступом к личным документам.
|
||||||
|
|
||||||
|
Полностью заменяет MySQLService. Реализованы все методы, необходимые для KBService и других компонентов.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from typing import List, Dict, Any, Optional
|
||||||
|
|
||||||
|
import asyncpg
|
||||||
|
from asyncpg import Pool, Connection
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class PostgresService:
|
||||||
|
"""
|
||||||
|
Асинхронный сервис для работы с PostgreSQL.
|
||||||
|
Использует пул соединений asyncpg.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, host: str, port: int, user: str, password: str, db_name: str):
|
||||||
|
self.host = host
|
||||||
|
self.port = port
|
||||||
|
self.user = user
|
||||||
|
self.password = password
|
||||||
|
self.db_name = db_name
|
||||||
|
self.pool: Optional[Pool] = None
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# Подключение и инициализация таблиц
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
async def connect(self) -> None:
|
||||||
|
"""
|
||||||
|
Создаёт пул соединений и инициализирует таблицы, если их нет.
|
||||||
|
"""
|
||||||
|
self.pool = await asyncpg.create_pool(
|
||||||
|
host=self.host,
|
||||||
|
port=self.port,
|
||||||
|
user=self.user,
|
||||||
|
password=self.password,
|
||||||
|
database=self.db_name,
|
||||||
|
min_size=1,
|
||||||
|
max_size=10,
|
||||||
|
command_timeout=60
|
||||||
|
)
|
||||||
|
await self._init_tables()
|
||||||
|
logger.info(f"Подключение к PostgreSQL БД {self.db_name} установлено (пул)")
|
||||||
|
|
||||||
|
async def _init_tables(self) -> None:
|
||||||
|
"""
|
||||||
|
Создаёт таблицы (если не существуют) с учётом всех необходимых колонок и индексов.
|
||||||
|
Использует схему из schema_postgresql.sql (адаптированную для PostgreSQL).
|
||||||
|
"""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
# 1. Таблица documents
|
||||||
|
await conn.execute("""
|
||||||
|
CREATE TABLE IF NOT EXISTS documents (
|
||||||
|
id SERIAL PRIMARY KEY,
|
||||||
|
source_name VARCHAR(255) NOT NULL,
|
||||||
|
upload_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
owner_jid VARCHAR(255) NOT NULL,
|
||||||
|
is_global BOOLEAN DEFAULT FALSE,
|
||||||
|
collection_name VARCHAR(64) NOT NULL,
|
||||||
|
metadata JSONB,
|
||||||
|
room_jid VARCHAR(255) DEFAULT NULL,
|
||||||
|
file_hash VARCHAR(64) DEFAULT NULL
|
||||||
|
)
|
||||||
|
""")
|
||||||
|
# Индексы
|
||||||
|
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_owner ON documents(owner_jid)")
|
||||||
|
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_global ON documents(is_global)")
|
||||||
|
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_room ON documents(room_jid)")
|
||||||
|
# Полнотекстовый индекс (GIN) – используем 'russian', при необходимости заменить на 'simple'
|
||||||
|
await conn.execute("""
|
||||||
|
CREATE INDEX IF NOT EXISTS ft_source_name
|
||||||
|
ON documents USING GIN (to_tsvector('russian', source_name))
|
||||||
|
""")
|
||||||
|
|
||||||
|
# 2. Таблица document_access
|
||||||
|
await conn.execute("""
|
||||||
|
CREATE TABLE IF NOT EXISTS document_access (
|
||||||
|
doc_id INT,
|
||||||
|
user_jid VARCHAR(255),
|
||||||
|
PRIMARY KEY (doc_id, user_jid)
|
||||||
|
)
|
||||||
|
""")
|
||||||
|
await conn.execute("CREATE INDEX IF NOT EXISTS idx_access_user ON document_access(user_jid)")
|
||||||
|
|
||||||
|
# 3. Тип ENUM для history.role (если не существует)
|
||||||
|
await conn.execute("""
|
||||||
|
DO $$ BEGIN
|
||||||
|
CREATE TYPE history_role AS ENUM ('user', 'assistant');
|
||||||
|
EXCEPTION
|
||||||
|
WHEN duplicate_object THEN NULL;
|
||||||
|
END $$;
|
||||||
|
""")
|
||||||
|
# Таблица history
|
||||||
|
await conn.execute("""
|
||||||
|
CREATE TABLE IF NOT EXISTS history (
|
||||||
|
id SERIAL PRIMARY KEY,
|
||||||
|
user_jid VARCHAR(255) NOT NULL,
|
||||||
|
role history_role NOT NULL,
|
||||||
|
content TEXT,
|
||||||
|
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
room_jid VARCHAR(255) DEFAULT NULL
|
||||||
|
)
|
||||||
|
""")
|
||||||
|
await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_user ON history(user_jid)")
|
||||||
|
await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_room ON history(room_jid)")
|
||||||
|
await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_time ON history(timestamp)")
|
||||||
|
|
||||||
|
# 4. Таблица rooms
|
||||||
|
await conn.execute("""
|
||||||
|
CREATE TABLE IF NOT EXISTS rooms (
|
||||||
|
room_jid VARCHAR(255) PRIMARY KEY
|
||||||
|
)
|
||||||
|
""")
|
||||||
|
|
||||||
|
async def close(self) -> None:
|
||||||
|
"""Закрывает пул соединений."""
|
||||||
|
if self.pool:
|
||||||
|
await self.pool.close()
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# Документы
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
async def add_document(
|
||||||
|
self,
|
||||||
|
source_name: str,
|
||||||
|
owner_jid: str,
|
||||||
|
is_global: bool,
|
||||||
|
collection_name: str,
|
||||||
|
metadata: dict = None,
|
||||||
|
room_jid: str = None,
|
||||||
|
file_hash: str = None
|
||||||
|
) -> int:
|
||||||
|
"""
|
||||||
|
Добавляет запись о документе. Возвращает doc_id.
|
||||||
|
|
||||||
|
:param source_name: название документа (source_name)
|
||||||
|
:param owner_jid: JID владельца
|
||||||
|
:param is_global: глобальный ли документ
|
||||||
|
:param collection_name: имя коллекции Qdrant
|
||||||
|
:param metadata: произвольные метаданные (словарь)
|
||||||
|
:param room_jid: JID комнаты (если документ комнатный)
|
||||||
|
:param file_hash: SHA-256 хеш содержимого (для обновления версий)
|
||||||
|
:return: id новой записи
|
||||||
|
"""
|
||||||
|
metadata_json = json.dumps(metadata, ensure_ascii=False) if metadata else None
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
row = await conn.fetchrow(
|
||||||
|
"""
|
||||||
|
INSERT INTO documents (source_name, owner_jid, is_global, collection_name, metadata, room_jid, file_hash)
|
||||||
|
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||||
|
RETURNING id
|
||||||
|
""",
|
||||||
|
source_name, owner_jid, is_global, collection_name, metadata_json, room_jid, file_hash
|
||||||
|
)
|
||||||
|
return row['id']
|
||||||
|
|
||||||
|
async def get_document(self, doc_id: int) -> Optional[Dict]:
|
||||||
|
"""Возвращает запись документа по id."""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
row = await conn.fetchrow("SELECT * FROM documents WHERE id = $1", doc_id)
|
||||||
|
return dict(row) if row else None
|
||||||
|
|
||||||
|
async def add_access(self, doc_id: int, user_jid: str) -> None:
|
||||||
|
"""
|
||||||
|
Даёт пользователю доступ к личному документу (только для не-глобальных, не-комнатных).
|
||||||
|
"""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
await conn.execute(
|
||||||
|
"""
|
||||||
|
INSERT INTO document_access (doc_id, user_jid)
|
||||||
|
VALUES ($1, $2)
|
||||||
|
ON CONFLICT (doc_id, user_jid) DO NOTHING
|
||||||
|
""",
|
||||||
|
doc_id, user_jid
|
||||||
|
)
|
||||||
|
|
||||||
|
async def can_access(self, doc_id: int, user_jid: str, is_global: bool) -> bool:
|
||||||
|
"""Проверяет, имеет ли пользователь доступ к документу."""
|
||||||
|
if is_global:
|
||||||
|
return True
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
row = await conn.fetchrow(
|
||||||
|
"SELECT 1 FROM document_access WHERE doc_id = $1 AND user_jid = $2",
|
||||||
|
doc_id, user_jid
|
||||||
|
)
|
||||||
|
return row is not None
|
||||||
|
|
||||||
|
async def find_document_by_name_and_room(self, source_name: str, room_jid: str) -> Optional[Dict]:
|
||||||
|
"""
|
||||||
|
Поиск документа в комнате по имени (source_name). Используется для обновления версий.
|
||||||
|
Возвращает словарь с ключами id, file_hash или None.
|
||||||
|
"""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
row = await conn.fetchrow(
|
||||||
|
"SELECT id, file_hash FROM documents WHERE source_name = $1 AND room_jid = $2 LIMIT 1",
|
||||||
|
source_name, room_jid
|
||||||
|
)
|
||||||
|
return dict(row) if row else None
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# История диалогов
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
async def add_history(self, user_jid: str, role: str, content: str, room_jid: str = None) -> None:
|
||||||
|
"""
|
||||||
|
Сохраняет сообщение в историю.
|
||||||
|
:param role: 'user' или 'assistant'
|
||||||
|
"""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
await conn.execute(
|
||||||
|
"INSERT INTO history (user_jid, role, content, room_jid) VALUES ($1, $2, $3, $4)",
|
||||||
|
user_jid, role, content, room_jid
|
||||||
|
)
|
||||||
|
|
||||||
|
async def get_history(self, user_jid: str, room_jid: str = None, limit: int = 100) -> List[Dict]:
|
||||||
|
"""
|
||||||
|
Возвращает последние limit сообщений, отсортированных от старых к новым.
|
||||||
|
Если room_jid указан – история комнаты, иначе – личная история пользователя.
|
||||||
|
"""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
if room_jid:
|
||||||
|
rows = await conn.fetch(
|
||||||
|
"""
|
||||||
|
SELECT role, content, timestamp
|
||||||
|
FROM history
|
||||||
|
WHERE room_jid = $1
|
||||||
|
ORDER BY timestamp ASC
|
||||||
|
LIMIT $2
|
||||||
|
""",
|
||||||
|
room_jid, limit
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
rows = await conn.fetch(
|
||||||
|
"""
|
||||||
|
SELECT role, content, timestamp
|
||||||
|
FROM history
|
||||||
|
WHERE user_jid = $1 AND room_jid IS NULL
|
||||||
|
ORDER BY timestamp ASC
|
||||||
|
LIMIT $2
|
||||||
|
""",
|
||||||
|
user_jid, limit
|
||||||
|
)
|
||||||
|
return [dict(row) for row in rows]
|
||||||
|
|
||||||
|
async def clear_user_history(self, user_jid: str) -> None:
|
||||||
|
"""Удаляет личную историю пользователя (не трогает комнаты)."""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
await conn.execute("DELETE FROM history WHERE user_jid = $1 AND room_jid IS NULL", user_jid)
|
||||||
|
|
||||||
|
async def clear_room_history(self, room_jid: str) -> None:
|
||||||
|
"""Удаляет историю указанной комнаты."""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
await conn.execute("DELETE FROM history WHERE room_jid = $1", room_jid)
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# Полнотекстовый поиск (fallback, если Qdrant недоступен)
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
async def fulltext_search(self, query: str, user_jid: str, room_jid: str = None) -> List[Dict]:
|
||||||
|
"""
|
||||||
|
Поиск по source_name с использованием полнотекстового индекса PostgreSQL.
|
||||||
|
Возвращает список документов (id, source_name, metadata).
|
||||||
|
"""
|
||||||
|
if not query or len(query.strip()) < 3:
|
||||||
|
return []
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
if room_jid:
|
||||||
|
rows = await conn.fetch(
|
||||||
|
"""
|
||||||
|
SELECT id, source_name, metadata
|
||||||
|
FROM documents
|
||||||
|
WHERE room_jid = $1
|
||||||
|
AND to_tsvector('russian', source_name) @@ plainto_tsquery('russian', $2)
|
||||||
|
LIMIT 10
|
||||||
|
""",
|
||||||
|
room_jid, query
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
rows = await conn.fetch(
|
||||||
|
"""
|
||||||
|
SELECT id, source_name, metadata
|
||||||
|
FROM documents
|
||||||
|
WHERE (owner_jid = $1 OR is_global = TRUE)
|
||||||
|
AND room_jid IS NULL
|
||||||
|
AND to_tsvector('russian', source_name) @@ plainto_tsquery('russian', $2)
|
||||||
|
LIMIT 10
|
||||||
|
""",
|
||||||
|
user_jid, query
|
||||||
|
)
|
||||||
|
return [dict(row) for row in rows]
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# Комнаты MUC
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
async def add_room(self, room_jid: str) -> None:
|
||||||
|
"""Добавляет комнату в список (если ещё не добавлена)."""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
await conn.execute(
|
||||||
|
"INSERT INTO rooms (room_jid) VALUES ($1) ON CONFLICT (room_jid) DO NOTHING",
|
||||||
|
room_jid
|
||||||
|
)
|
||||||
|
|
||||||
|
async def get_rooms(self) -> List[str]:
|
||||||
|
"""Возвращает список JID комнат, в которых бот должен состоять."""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
rows = await conn.fetch("SELECT room_jid FROM rooms")
|
||||||
|
return [row['room_jid'] for row in rows]
|
||||||
|
|
||||||
|
async def remove_room(self, room_jid: str) -> None:
|
||||||
|
"""Удаляет комнату из списка."""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
await conn.execute("DELETE FROM rooms WHERE room_jid = $1", room_jid)
|
||||||
|
|
||||||
|
async def save_template(self, room_jid: str, name: str, file_path: str, file_hash: str) -> bool:
|
||||||
|
"""Сохраняет шаблон в БД (если с таким именем уже есть – перезаписывает)."""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
await conn.execute("""
|
||||||
|
INSERT INTO room_templates (room_jid, name, file_path, file_hash)
|
||||||
|
VALUES ($1, $2, $3, $4)
|
||||||
|
ON CONFLICT (room_jid, name) DO UPDATE SET
|
||||||
|
file_path = EXCLUDED.file_path,
|
||||||
|
file_hash = EXCLUDED.file_hash,
|
||||||
|
created_at = CURRENT_TIMESTAMP
|
||||||
|
""", room_jid, name, file_path, file_hash)
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def get_template(self, room_jid: str, name: str) -> Optional[Dict]:
|
||||||
|
"""Возвращает словарь {file_path, file_hash} или None."""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
row = await conn.fetchrow("SELECT file_path, file_hash FROM room_templates WHERE room_jid = $1 AND name = $2", room_jid, name)
|
||||||
|
return dict(row) if row else None
|
||||||
|
|
||||||
|
async def list_templates(self, room_jid: str) -> List[str]:
|
||||||
|
"""Возвращает список названий шаблонов в комнате."""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
rows = await conn.fetch("SELECT name FROM room_templates WHERE room_jid = $1 ORDER BY name", room_jid)
|
||||||
|
return [row['name'] for row in rows]
|
||||||
|
|
||||||
|
async def delete_template(self, room_jid: str, name: str) -> bool:
|
||||||
|
"""Удаляет шаблон. Возвращает True, если удалён."""
|
||||||
|
async with self.pool.acquire() as conn:
|
||||||
|
result = await conn.execute("DELETE FROM room_templates WHERE room_jid = $1 AND name = $2", room_jid, name)
|
||||||
|
return result == "DELETE 1"
|
||||||
Reference in New Issue
Block a user