Редактировать postgres_service.py
This commit is contained in:
Markov Andrey
@@ -5,6 +5,7 @@
|
||||
управление доступом к личным документам.
|
||||
|
||||
Добавлена поддержка флага indexed для транзакционности индексации.
|
||||
ВНИМАНИЕ: Создание таблиц теперь выполняется через Alembic, а не в коде.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -33,12 +34,17 @@ class PostgresService:
|
||||
self.pool: Optional[Pool] = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Подключение и инициализация таблиц
|
||||
# Подключение и проверка таблиц
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def connect(self) -> 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,
|
||||
@@ -50,13 +56,42 @@ class PostgresService:
|
||||
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:
|
||||
"""
|
||||
Проверяет наличие основных таблиц. Если какая-то таблица отсутствует,
|
||||
выводит предупреждение и советует запустить миграции Alembic.
|
||||
"""
|
||||
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:
|
||||
"""
|
||||
Создаёт таблицы (если не существуют) с учётом всех необходимых колонок и индексов.
|
||||
Использует схему из schema_postgresql.sql (адаптированную для PostgreSQL).
|
||||
Создаёт таблицы (если не существуют) — используется только для тестов.
|
||||
В продакшене используйте Alembic.
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
# 1. Таблица documents
|
||||
@@ -74,12 +109,10 @@ class PostgresService:
|
||||
indexed BOOLEAN DEFAULT FALSE
|
||||
)
|
||||
""")
|
||||
# Индексы
|
||||
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)")
|
||||
# Полнотекстовый индекс (GIN) – используем 'russian', при необходимости заменить на 'simple'
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS ft_source_name
|
||||
ON documents USING GIN (to_tsvector('russian', source_name))
|
||||
@@ -95,7 +128,7 @@ class PostgresService:
|
||||
""")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_access_user ON document_access(user_jid)")
|
||||
|
||||
# 3. Тип ENUM для history.role (если не существует)
|
||||
# 3. Тип ENUM для history.role
|
||||
await conn.execute("""
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE history_role AS ENUM ('user', 'assistant');
|
||||
@@ -103,7 +136,6 @@ class PostgresService:
|
||||
WHEN duplicate_object THEN NULL;
|
||||
END $$;
|
||||
""")
|
||||
# Таблица history
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS history (
|
||||
id SERIAL PRIMARY KEY,
|
||||
@@ -125,6 +157,19 @@ class PostgresService:
|
||||
)
|
||||
""")
|
||||
|
||||
# 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:
|
||||
@@ -143,12 +188,9 @@ class PostgresService:
|
||||
metadata: dict = None,
|
||||
room_jid: str = None,
|
||||
file_hash: str = None,
|
||||
indexed: bool = False # новый параметр
|
||||
indexed: bool = False
|
||||
) -> int:
|
||||
"""
|
||||
Добавляет запись о документе. Возвращает doc_id.
|
||||
:param indexed: начальное состояние индексации (по умолчанию False)
|
||||
"""
|
||||
"""Добавляет запись о документе. Возвращает 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(
|
||||
@@ -176,9 +218,7 @@ class PostgresService:
|
||||
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(
|
||||
"""
|
||||
@@ -217,10 +257,7 @@ class PostgresService:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
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)",
|
||||
@@ -258,7 +295,7 @@ class PostgresService:
|
||||
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)
|
||||
|
||||
@@ -268,7 +305,7 @@ class PostgresService:
|
||||
await conn.execute("DELETE FROM history WHERE room_jid = $1", room_jid)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Полнотекстовый поиск (fallback, если Qdrant недоступен)
|
||||
# Полнотекстовый поиск (fallback)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def fulltext_search(self, query: str, user_jid: str, room_jid: str = None) -> List[Dict]:
|
||||
@@ -311,7 +348,7 @@ class PostgresService:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
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",
|
||||
@@ -349,7 +386,10 @@ class PostgresService:
|
||||
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)
|
||||
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]:
|
||||
@@ -361,7 +401,10 @@ class PostgresService:
|
||||
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)
|
||||
result = await conn.execute(
|
||||
"DELETE FROM room_templates WHERE room_jid = $1 AND name = $2",
|
||||
room_jid, name
|
||||
)
|
||||
return result == "DELETE 1"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@@ -409,7 +452,7 @@ class PostgresService:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def clear_user_kb(self, user_jid: str) -> List[int]:
|
||||
"""Удаляет все личные документы пользователя (не глобальные и не комнатные). Возвращает список удалённых doc_id."""
|
||||
"""Удаляет все личные документы пользователя. Возвращает список удалённых doc_id."""
|
||||
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",
|
||||
@@ -440,7 +483,7 @@ class PostgresService:
|
||||
return doc_ids
|
||||
|
||||
async def delete_documents_by_ids(self, doc_ids: List[int]) -> None:
|
||||
"""Удаляет документы по списку id (используется в global_remove)."""
|
||||
"""Удаляет документы по списку id."""
|
||||
if not doc_ids:
|
||||
return
|
||||
async with self.pool.acquire() as conn:
|
||||
|
||||
Reference in New Issue
Block a user