496 lines
23 KiB
Python
496 lines
23 KiB
Python
# rag/services/postgres_service.py
|
||
"""
|
||
Сервис для работы с PostgreSQL (через asyncpg).
|
||
Обеспечивает пул соединений, CRUD операции для документов, истории диалогов, комнат MUC,
|
||
управление доступом к личным документам.
|
||
|
||
Добавлена поддержка флага indexed для транзакционности индексации.
|
||
Добавлена колонка file_text для хранения исходного текста документа (для переиндексации).
|
||
"""
|
||
|
||
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, create_if_missing: bool = False) -> None:
|
||
"""
|
||
Создаёт пул соединений и проверяет наличие таблиц.
|
||
Таблицы должны быть созданы через Alembic.
|
||
|
||
Аргументы:
|
||
create_if_missing: если True, создаёт таблицы автоматически (только для тестов).
|
||
По умолчанию False.
|
||
"""
|
||
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
|
||
)
|
||
if create_if_missing:
|
||
await self._init_tables()
|
||
logger.warning("Таблицы созданы автоматически (режим тестирования). В продакшене используйте Alembic.")
|
||
else:
|
||
await self._check_tables()
|
||
logger.info(f"Подключение к PostgreSQL БД {self.db_name} установлено (пул)")
|
||
|
||
async def _check_tables(self) -> None:
|
||
"""Проверяет наличие основных таблиц."""
|
||
required_tables = ['documents', 'document_access', 'history', 'rooms', 'room_templates']
|
||
missing = []
|
||
async with self.pool.acquire() as conn:
|
||
for table in required_tables:
|
||
row = await conn.fetchval(
|
||
"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)",
|
||
table
|
||
)
|
||
if not row:
|
||
missing.append(table)
|
||
if missing:
|
||
logger.warning(
|
||
f"Отсутствуют таблицы: {', '.join(missing)}. "
|
||
"Запустите миграции Alembic: `alembic upgrade head`"
|
||
)
|
||
else:
|
||
logger.debug("Все необходимые таблицы существуют")
|
||
|
||
async def _init_tables(self) -> None:
|
||
"""Создаёт таблицы (только для тестов)."""
|
||
async with self.pool.acquire() as conn:
|
||
# 1. Таблица documents с колонкой file_text
|
||
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,
|
||
indexed BOOLEAN DEFAULT FALSE,
|
||
file_text TEXT 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)")
|
||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_indexed ON documents(indexed)")
|
||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_room_source ON documents(room_jid, source_name)")
|
||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_owner_source ON documents(owner_jid, source_name)")
|
||
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 $$;
|
||
""")
|
||
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
|
||
)
|
||
""")
|
||
|
||
# 5. Таблица room_templates
|
||
await conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS room_templates (
|
||
room_jid VARCHAR(255) NOT NULL,
|
||
name VARCHAR(255) NOT NULL,
|
||
file_path TEXT NOT NULL,
|
||
file_hash VARCHAR(64),
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
PRIMARY KEY (room_jid, name)
|
||
)
|
||
""")
|
||
logger.info("Таблицы созданы (или уже существовали)")
|
||
|
||
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,
|
||
indexed: bool = False,
|
||
file_text: Optional[str] = None
|
||
) -> int:
|
||
"""
|
||
Добавляет запись о документе с возможностью сохранить исходный текст.
|
||
Возвращает doc_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, indexed, file_text
|
||
)
|
||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
|
||
RETURNING id
|
||
""",
|
||
source_name, owner_jid, is_global, collection_name,
|
||
metadata_json, room_jid, file_hash, indexed, file_text
|
||
)
|
||
return row['id']
|
||
|
||
async def update_document_indexed(self, doc_id: int, indexed: bool) -> None:
|
||
"""Обновляет флаг indexed для документа."""
|
||
async with self.pool.acquire() as conn:
|
||
await conn.execute(
|
||
"UPDATE documents SET indexed = $1 WHERE id = $2",
|
||
indexed, doc_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:
|
||
"""Сохраняет сообщение в историю."""
|
||
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, offset: int = 0) -> List[Dict]:
|
||
"""
|
||
Возвращает историю диалога с поддержкой пагинации.
|
||
|
||
Аргументы:
|
||
user_jid: JID пользователя.
|
||
room_jid: JID комнаты (None для личного чата).
|
||
limit: максимальное количество сообщений.
|
||
offset: смещение (сколько пропустить).
|
||
|
||
Возвращает:
|
||
List[Dict]: список сообщений с ключами role, content, timestamp.
|
||
"""
|
||
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 OFFSET $3
|
||
""",
|
||
room_jid, limit, offset
|
||
)
|
||
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 OFFSET $3
|
||
""",
|
||
user_jid, limit, offset
|
||
)
|
||
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)
|
||
# ------------------------------------------------------------------
|
||
|
||
async def fulltext_search(self, query: str, user_jid: str, room_jid: str = None) -> List[Dict]:
|
||
"""
|
||
Поиск по source_name с использованием полнотекстового индекса PostgreSQL.
|
||
Возвращает список документов (id, source_name, metadata) только с indexed = TRUE.
|
||
"""
|
||
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 indexed = TRUE
|
||
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 indexed = TRUE
|
||
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:
|
||
"""Удаляет шаблон."""
|
||
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"
|
||
|
||
# ------------------------------------------------------------------
|
||
# Получение списков документов для пользователей (только indexed = TRUE)
|
||
# ------------------------------------------------------------------
|
||
|
||
async def get_available_sources(self, user_jid: str, room_jid: str = None) -> List[str]:
|
||
"""Возвращает список названий документов (source_name), доступных пользователю."""
|
||
async with self.pool.acquire() as conn:
|
||
if room_jid:
|
||
rows = await conn.fetch(
|
||
"SELECT source_name FROM documents WHERE room_jid = $1 AND indexed = TRUE ORDER BY source_name",
|
||
room_jid
|
||
)
|
||
else:
|
||
rows = await conn.fetch(
|
||
"SELECT source_name FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL AND indexed = TRUE ORDER BY source_name",
|
||
user_jid
|
||
)
|
||
return [row['source_name'] for row in rows]
|
||
|
||
async def get_sources_with_type(self, user_jid: str, room_jid: str = None) -> List[Dict[str, Any]]:
|
||
"""Возвращает список документов с указанием, является ли документ глобальным."""
|
||
async with self.pool.acquire() as conn:
|
||
if room_jid:
|
||
rows = await conn.fetch(
|
||
"SELECT source_name, is_global FROM documents WHERE room_jid = $1 AND indexed = TRUE ORDER BY source_name",
|
||
room_jid
|
||
)
|
||
else:
|
||
rows = await conn.fetch(
|
||
"SELECT source_name, is_global FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL AND indexed = TRUE ORDER BY source_name",
|
||
user_jid
|
||
)
|
||
return [{"source_name": row['source_name'], "is_global": row['is_global']} for row in rows]
|
||
|
||
# ------------------------------------------------------------------
|
||
# Очистка баз знаний (удаление всех документов, независимо от indexed)
|
||
# ------------------------------------------------------------------
|
||
|
||
async def clear_user_kb(self, user_jid: str) -> List[int]:
|
||
"""Удаляет все личные документы пользователя."""
|
||
async with self.pool.acquire() as conn:
|
||
rows = await conn.fetch(
|
||
"SELECT id FROM documents WHERE owner_jid = $1 AND is_global = FALSE AND room_jid IS NULL",
|
||
user_jid
|
||
)
|
||
doc_ids = [row['id'] for row in rows]
|
||
if doc_ids:
|
||
await conn.execute("DELETE FROM document_access WHERE doc_id = ANY($1)", doc_ids)
|
||
await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
|
||
return doc_ids
|
||
|
||
async def clear_room_kb(self, room_jid: str) -> List[int]:
|
||
"""Удаляет все документы указанной комнаты."""
|
||
async with self.pool.acquire() as conn:
|
||
rows = await conn.fetch("SELECT id FROM documents WHERE room_jid = $1", room_jid)
|
||
doc_ids = [row['id'] for row in rows]
|
||
if doc_ids:
|
||
await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
|
||
return doc_ids
|
||
|
||
async def clear_global_kb(self) -> List[int]:
|
||
"""Удаляет все глобальные документы."""
|
||
async with self.pool.acquire() as conn:
|
||
rows = await conn.fetch("SELECT id FROM documents WHERE is_global = TRUE")
|
||
doc_ids = [row['id'] for row in rows]
|
||
if doc_ids:
|
||
await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
|
||
return doc_ids
|
||
|
||
async def delete_documents_by_ids(self, doc_ids: List[int]) -> None:
|
||
"""Удаляет документы по списку id."""
|
||
if not doc_ids:
|
||
return
|
||
async with self.pool.acquire() as conn:
|
||
await conn.execute("DELETE FROM document_access WHERE doc_id = ANY($1)", doc_ids)
|
||
await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids) |