Редактировать postgres_service.py
This commit is contained in:
Markov Andrey
@@ -5,8 +5,7 @@
|
||||
управление доступом к личным документам.
|
||||
|
||||
Добавлена поддержка флага indexed для транзакционности индексации.
|
||||
Добавлены составные индексы для ускорения полнотекстового поиска (fallback).
|
||||
ВНИМАНИЕ: Создание таблиц теперь выполняется через Alembic, а не в коде.
|
||||
Добавлена колонка file_text для хранения исходного текста документа (для переиндексации).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -58,19 +57,14 @@ class PostgresService:
|
||||
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:
|
||||
@@ -90,12 +84,9 @@ class PostgresService:
|
||||
logger.debug("Все необходимые таблицы существуют")
|
||||
|
||||
async def _init_tables(self) -> None:
|
||||
"""
|
||||
Создаёт таблицы (если не существуют) — используется только для тестов.
|
||||
В продакшене используйте Alembic.
|
||||
"""
|
||||
"""Создаёт таблицы (только для тестов)."""
|
||||
async with self.pool.acquire() as conn:
|
||||
# 1. Таблица documents
|
||||
# 1. Таблица documents с колонкой file_text
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id SERIAL PRIMARY KEY,
|
||||
@@ -107,18 +98,16 @@ class PostgresService:
|
||||
metadata JSONB,
|
||||
room_jid VARCHAR(255) DEFAULT NULL,
|
||||
file_hash VARCHAR(64) DEFAULT NULL,
|
||||
indexed BOOLEAN DEFAULT FALSE
|
||||
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)")
|
||||
# Составные индексы для ускорения полнотекстового поиска (fallback)
|
||||
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)")
|
||||
# Полнотекстовый индекс (GIN) – используем 'russian', при необходимости заменить на 'simple'
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS ft_source_name
|
||||
ON documents USING GIN (to_tsvector('russian', source_name))
|
||||
@@ -194,18 +183,26 @@ class PostgresService:
|
||||
metadata: dict = None,
|
||||
room_jid: str = None,
|
||||
file_hash: str = None,
|
||||
indexed: bool = False
|
||||
indexed: bool = False,
|
||||
file_text: Optional[str] = None
|
||||
) -> int:
|
||||
"""Добавляет запись о документе. Возвращает doc_id."""
|
||||
"""
|
||||
Добавляет запись о документе с возможностью сохранить исходный текст.
|
||||
Возвращает 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)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
|
||||
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
|
||||
source_name, owner_jid, is_global, collection_name,
|
||||
metadata_json, room_jid, file_hash, indexed, file_text
|
||||
)
|
||||
return row['id']
|
||||
|
||||
@@ -326,7 +323,6 @@ class PostgresService:
|
||||
"""
|
||||
Поиск по source_name с использованием полнотекстового индекса PostgreSQL.
|
||||
Возвращает список документов (id, source_name, metadata) только с indexed = TRUE.
|
||||
Для ускорения используются составные индексы (room_jid, source_name) и (owner_jid, source_name).
|
||||
"""
|
||||
if not query or len(query.strip()) < 3:
|
||||
return []
|
||||
@@ -386,7 +382,7 @@ class PostgresService:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
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)
|
||||
@@ -414,7 +410,7 @@ class PostgresService:
|
||||
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",
|
||||
@@ -427,10 +423,7 @@ class PostgresService:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def get_available_sources(self, user_jid: str, room_jid: str = None) -> List[str]:
|
||||
"""
|
||||
Возвращает список названий документов (source_name), доступных пользователю.
|
||||
Только те, у которых indexed = TRUE.
|
||||
"""
|
||||
"""Возвращает список названий документов (source_name), доступных пользователю."""
|
||||
async with self.pool.acquire() as conn:
|
||||
if room_jid:
|
||||
rows = await conn.fetch(
|
||||
@@ -445,10 +438,7 @@ class PostgresService:
|
||||
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]]:
|
||||
"""
|
||||
Возвращает список документов с указанием, является ли документ глобальным.
|
||||
Только indexed = TRUE.
|
||||
"""
|
||||
"""Возвращает список документов с указанием, является ли документ глобальным."""
|
||||
async with self.pool.acquire() as conn:
|
||||
if room_jid:
|
||||
rows = await conn.fetch(
|
||||
@@ -467,7 +457,7 @@ class PostgresService:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def clear_user_kb(self, user_jid: str) -> List[int]:
|
||||
"""Удаляет все личные документы пользователя. Возвращает список удалённых 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",
|
||||
@@ -480,7 +470,7 @@ class PostgresService:
|
||||
return doc_ids
|
||||
|
||||
async def clear_room_kb(self, room_jid: str) -> List[int]:
|
||||
"""Удаляет все документы указанной комнаты. Возвращает список doc_id."""
|
||||
"""Удаляет все документы указанной комнаты."""
|
||||
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]
|
||||
@@ -489,7 +479,7 @@ class PostgresService:
|
||||
return doc_ids
|
||||
|
||||
async def clear_global_kb(self) -> List[int]:
|
||||
"""Удаляет все глобальные документы. Возвращает список doc_id."""
|
||||
"""Удаляет все глобальные документы."""
|
||||
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]
|
||||
|
||||
Reference in New Issue
Block a user