Редактировать postgres_service.py

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."""
-        Добавляет запись о документе. Возвращает doc_id.
-        :param indexed: начальное состояние индексации (по умолчанию False)
-        """
         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: