From e60318cbd6a3b3f4db8b508df957097c65afff92 Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 22:48:34 +0000 Subject: [PATCH] =?UTF-8?q?=D0=A0=D0=B5=D0=B4=D0=B0=D0=BA=D1=82=D0=B8?= =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B0=D1=82=D1=8C=20postgres=5Fservice.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- rag/services/postgres_service.py | 62 ++++++++++++++------------------ 1 file changed, 26 insertions(+), 36 deletions(-) diff --git a/rag/services/postgres_service.py b/rag/services/postgres_service.py index a76335f..098c5c0 100644 --- a/rag/services/postgres_service.py +++ b/rag/services/postgres_service.py @@ -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]