From b11cd99db330bbc13b5c409298629778e0caace0 Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 22:32:42 +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 | 101 ++++++++++++++++++++++--------- 1 file changed, 72 insertions(+), 29 deletions(-) diff --git a/rag/services/postgres_service.py b/rag/services/postgres_service.py index f356d34..cd55b94 100644 --- a/rag/services/postgres_service.py +++ b/rag/services/postgres_service.py @@ -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 ) - await self._init_tables() + 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: