fix: 模特管理

python_core/database/model_postgres.py

@@ -5,24 +5,63 @@
 import uuid
 from datetime import datetime
 from typing import List, Optional, Dict, Any
-import psycopg2
-from psycopg2.extras import RealDictCursor
+from contextlib import contextmanager
+from python_core.config import settings
 from python_core.database.types import Model
-from python_core.database.db_postgres import DatabasePostgres
-from python_core.utils.logger import logger
+from python_core.utils.logger import setup_logger
+
+# 尝试导入 psycopg2，如果失败则提供友好的错误信息
+try:
+    import psycopg2
+    import psycopg2.extras
+    PSYCOPG2_AVAILABLE = True
+except ImportError as e:
+    PSYCOPG2_AVAILABLE = False
+    PSYCOPG2_ERROR = str(e)
+
+logger = setup_logger(__name__)
 
 
-class ModelTablePostgres(DatabasePostgres):
+class ModelTablePostgres:
     """模特表 PostgreSQL 实现"""
-    
+
     def __init__(self):
-        super().__init__()
+        if not PSYCOPG2_AVAILABLE:
+            error_msg = f"""
+PostgreSQL support requires psycopg2 package.
+Please install it using: pip install psycopg2-binary
+
+原始错误: {PSYCOPG2_ERROR}
+"""
+            logger.error(error_msg)
+            raise ImportError(error_msg)
+
+        self.db_url = settings.db
+        self.table_name = "models"
+
+        # 初始化模特表
         self._init_model_table()
+
+    @contextmanager
+    def _get_connection(self):
+        """获取数据库连接的上下文管理器"""
+        conn = None
+        try:
+            conn = psycopg2.connect(self.db_url)
+            yield conn
+        except Exception as e:
+            if conn:
+                conn.rollback()
+            logger.error(f"Database connection error: {e}")
+            raise e
+        finally:
+            if conn:
+                conn.close()
     
     def _init_model_table(self):
         """初始化模特表"""
         try:
-            with self.get_connection() as conn:
+            with self._get_connection() as conn:
                 with conn.cursor() as cursor:
                     # 创建模特表
                     cursor.execute("""
@@ -99,8 +138,8 @@ class ModelTablePostgres(DatabasePostgres):
                     is_cloud: bool = False) -> Optional[Model]:
         """创建模特"""
         try:
-            with self.get_connection() as conn:
-                with conn.cursor(cursor_factory=RealDictCursor) as cursor:
+            with self._get_connection() as conn:
+                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                     cursor.execute("""
                         INSERT INTO models (model_number, model_image, user_id, is_cloud)
                         VALUES (%s, %s, %s, %s)
@@ -131,8 +170,8 @@ class ModelTablePostgres(DatabasePostgres):
     def get_model_by_id(self, model_id: str) -> Optional[Model]:
         """根据ID获取模特"""
         try:
-            with self.get_connection() as conn:
-                with conn.cursor(cursor_factory=RealDictCursor) as cursor:
+            with self._get_connection() as conn:
+                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                     cursor.execute("""
                         SELECT * FROM models WHERE id = %s
                     """, (model_id,))
@@ -150,8 +189,8 @@ class ModelTablePostgres(DatabasePostgres):
     def get_model_by_number(self, model_number: str, user_id: str = "default") -> Optional[Model]:
         """根据模特编号获取模特"""
         try:
-            with self.get_connection() as conn:
-                with conn.cursor(cursor_factory=RealDictCursor) as cursor:
+            with self._get_connection() as conn:
+                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                     cursor.execute("""
                         SELECT * FROM models 
                         WHERE model_number = %s AND user_id = %s
@@ -172,8 +211,8 @@ class ModelTablePostgres(DatabasePostgres):
                       offset: int = 0) -> List[Model]:
         """获取所有模特"""
         try:
-            with self.get_connection() as conn:
-                with conn.cursor(cursor_factory=RealDictCursor) as cursor:
+            with self._get_connection() as conn:
+                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                     # 构建查询条件
                     conditions = []
                     params = []
@@ -224,7 +263,7 @@ class ModelTablePostgres(DatabasePostgres):
                 logger.warning("没有有效的更新字段")
                 return False
             
-            with self.get_connection() as conn:
+            with self._get_connection() as conn:
                 with conn.cursor() as cursor:
                     # 构建更新语句
                     set_clauses = []
@@ -266,7 +305,7 @@ class ModelTablePostgres(DatabasePostgres):
     def delete_model(self, model_id: str, hard_delete: bool = False) -> bool:
         """删除模特"""
         try:
-            with self.get_connection() as conn:
+            with self._get_connection() as conn:
                 with conn.cursor() as cursor:
                     if hard_delete:
                         # 硬删除
@@ -298,31 +337,33 @@ class ModelTablePostgres(DatabasePostgres):
                      include_cloud: bool = True, limit: int = 50) -> List[Model]:
         """搜索模特"""
         try:
-            with self.get_connection() as conn:
-                with conn.cursor(cursor_factory=RealDictCursor) as cursor:
-                    # 构建查询条件
-                    conditions = ["is_active = true"]
-                    params = [f"%{query}%"]
-                    
+            with self._get_connection() as conn:
+                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
+                    # 简化搜索逻辑
+                    search_pattern = f"%{query}%"
+
                     if include_cloud:
-                        conditions.append("(user_id = %s OR is_cloud = true)")
-                        params.append(user_id)
+                        cursor.execute("""
+                            SELECT * FROM models
+                            WHERE model_number ILIKE %s
+                            AND is_active = true
+                            AND (user_id = %s OR is_cloud = true)
+                            ORDER BY
+                                CASE WHEN model_number ILIKE %s THEN 1 ELSE 2 END,
+                                created_at DESC
+                            LIMIT %s
+                        """, (search_pattern, user_id, search_pattern, limit))
                     else:
-                        conditions.append("user_id = %s")
-                        params.append(user_id)
-                    
-                    params.append(limit)
-                    
-                    where_clause = " AND ".join(conditions)
-                    
-                    cursor.execute(f"""
-                        SELECT * FROM models 
-                        WHERE model_number ILIKE %s AND {where_clause}
-                        ORDER BY 
-                            CASE WHEN model_number ILIKE %s THEN 1 ELSE 2 END,
-                            created_at DESC
-                        LIMIT %s
-                    """, [f"%{query}%"] + params)
+                        cursor.execute("""
+                            SELECT * FROM models
+                            WHERE model_number ILIKE %s
+                            AND is_active = true
+                            AND user_id = %s
+                            ORDER BY
+                                CASE WHEN model_number ILIKE %s THEN 1 ELSE 2 END,
+                                created_at DESC
+                            LIMIT %s
+                        """, (search_pattern, user_id, search_pattern, limit))
                     
                     rows = cursor.fetchall()
                     return [self._row_to_model(row) for row in rows]
@@ -337,7 +378,7 @@ class ModelTablePostgres(DatabasePostgres):
                        include_inactive: bool = False) -> int:
         """获取模特数量"""
         try:
-            with self.get_connection() as conn:
+            with self._get_connection() as conn:
                 with conn.cursor() as cursor:
                     # 构建查询条件
                     conditions = []
@@ -371,7 +412,7 @@ class ModelTablePostgres(DatabasePostgres):
     def toggle_model_status(self, model_id: str) -> bool:
         """切换模特状态"""
         try:
-            with self.get_connection() as conn:
+            with self._get_connection() as conn:
                 with conn.cursor() as cursor:
                     cursor.execute("""
                         UPDATE models