更新python后端

bbit_ai/app/config/pgDb.py

@@ -1,22 +1,69 @@
-
-import psycopg
-from langchain_postgres import PostgresChatMessageHistory
-from psycopg_pool import ConnectionPool
+import logging
+import time
 from contextlib import contextmanager
+from typing import Optional
+import psycopg
+from psycopg_pool import ConnectionPool
 
-# conn = psycopg.connect("postgresql://postgres:123456@10.10.10.9/ktor2")
-database_name = "ai_chat_history"
-pool = ConnectionPool("postgresql://postgres:123456@10.10.10.9/ktor2")
+logger = logging.getLogger("PGPool")
+logger.setLevel(logging.INFO)
 
-@contextmanager
-def getConn():
-    with pool.connection() as temp:
-        temp.autocommit = True   # 如果你想所有连接默认 autocommit
-        yield temp               # 把 conn 暴露给外部使用
+class PGPool:
+    """
+    PostgreSQL 连接池封装
+    """
 
+    def __init__(
+        self,
+        uri: str,
+        min_size: int = 1,
+        max_size: int = 20,
+        max_idle: int = 30,
+        max_lifetime: int = 300,
+        timeout: int = 10,
+        check: bool = False,
+    ):
+        """
+        :param uri: PostgreSQL 连接 URI
+        """
+        self.uri = uri
+        self.pool = ConnectionPool(
+            self.uri,
+            min_size=min_size,
+            max_size=max_size,
+            max_idle=max_idle,
+            max_lifetime=max_lifetime,
+            timeout=timeout,
+            check=check,
+        )
 
-def init():
-    with getConn() as connection:
-        PostgresChatMessageHistory.create_tables(connection, database_name)
-
-init()
+    @contextmanager
+    def getConn(self, retries: int = 2, delay: float = 1.0):
+        """
+        获取数据库连接，带重试机制，自动健康检查。
+        使用方式：
+            with pg_pool.get_conn() as conn:
+                with conn.cursor() as cur:
+                    cur.execute(...)
+        """
+        attempt = 0
+        while attempt <= retries:
+            try:
+                with self.pool.connection() as conn:
+                    conn.autocommit = True
+                    yield conn
+                    return
+            except psycopg.OperationalError as e:
+                logger.warning(f"数据库连接异常: {e}. 尝试重试 ({attempt+1}/{retries})")
+                self.pool.check()  # 丢掉坏连接，重新建
+                attempt += 1
+                time.sleep(delay)
+            except Exception as e:
+                logger.error(f"SQL执行异常: {e}")
+                raise
+        raise psycopg.OperationalError("无法获取数据库连接，多次重试失败")
+pg_pool = PGPool(
+    uri="postgresql://postgres:123456@10.10.10.9/ktor2",
+    min_size=1,
+    max_size=20,
+)