fix
This commit is contained in:
@@ -10,7 +10,7 @@ from rich.panel import Panel
|
||||
import json
|
||||
import re
|
||||
|
||||
from python_core.services.resource_category_manager import ResourceCategoryManager, ResourceCategory
|
||||
from python_core.services.resource_category_manager_v2 import resource_category_manager_v2, ResourceCategory
|
||||
from python_core.utils.logger import logger
|
||||
|
||||
console = Console()
|
||||
@@ -36,7 +36,7 @@ def list_categories(
|
||||
console.print(f"📋 [bold blue]资源分类列表[/bold blue]")
|
||||
|
||||
# 创建分类管理器
|
||||
manager = ResourceCategoryManager()
|
||||
manager = resource_category_manager_v2
|
||||
|
||||
# 获取分类列表
|
||||
categories = manager.get_all_categories()
|
||||
@@ -133,7 +133,7 @@ def create_category(
|
||||
raise typer.Exit(1)
|
||||
|
||||
# 创建分类管理器
|
||||
manager = ResourceCategoryManager()
|
||||
manager = resource_category_manager_v2
|
||||
|
||||
# 创建分类
|
||||
result = manager.create_category(title, ai_prompt, color)
|
||||
@@ -175,7 +175,7 @@ def get_category(
|
||||
console.print(f"分类ID: {category_id}")
|
||||
|
||||
# 创建分类管理器
|
||||
manager = ResourceCategoryManager()
|
||||
manager = resource_category_manager_v2
|
||||
|
||||
# 获取分类
|
||||
category = manager.get_category_by_id(category_id)
|
||||
@@ -248,7 +248,7 @@ def update_category(
|
||||
raise typer.Exit(1)
|
||||
|
||||
# 创建分类管理器
|
||||
manager = ResourceCategoryManager()
|
||||
manager = resource_category_manager_v2
|
||||
|
||||
# 更新分类
|
||||
result = manager.update_category(
|
||||
@@ -298,7 +298,7 @@ def delete_category(
|
||||
console.print(f"分类ID: {category_id}")
|
||||
|
||||
# 创建分类管理器
|
||||
manager = ResourceCategoryManager()
|
||||
manager = resource_category_manager_v2
|
||||
|
||||
# 获取分类信息
|
||||
category = manager.get_category_by_id(category_id)
|
||||
@@ -343,7 +343,7 @@ def search_categories(
|
||||
console.print(f"关键词: {keyword}")
|
||||
|
||||
# 创建分类管理器
|
||||
manager = ResourceCategoryManager()
|
||||
manager = resource_category_manager_v2
|
||||
|
||||
# 搜索分类
|
||||
categories = manager.search_categories(keyword)
|
||||
@@ -408,7 +408,7 @@ def show_stats():
|
||||
console.print(f"📊 [bold blue]分类统计信息[/bold blue]")
|
||||
|
||||
# 创建分类管理器
|
||||
manager = ResourceCategoryManager()
|
||||
manager = resource_category_manager_v2
|
||||
|
||||
# 获取所有分类
|
||||
categories = manager.get_all_categories()
|
||||
|
||||
660
python_core/database/resource_category_postgres.py
Normal file
660
python_core/database/resource_category_postgres.py
Normal file
@@ -0,0 +1,660 @@
|
||||
|
||||
# 素材分类表 - PostgreSQL 版本
|
||||
|
||||
import uuid
|
||||
from typing import Dict, List, Any, Optional
|
||||
from datetime import datetime
|
||||
from contextlib import contextmanager
|
||||
from python_core.config import settings
|
||||
from python_core.utils.logger import setup_logger
|
||||
from .types import ResourceCategory
|
||||
|
||||
# 尝试导入 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 ResourceCategoryPostgres:
|
||||
"""
|
||||
素材分类表类 - PostgreSQL 版本
|
||||
基于 PostgreSQL 数据库实现的素材分类管理功能
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# 检查 psycopg2 是否可用
|
||||
if not PSYCOPG2_AVAILABLE:
|
||||
error_msg = f"""
|
||||
PostgreSQL 驱动 psycopg2 未安装。请安装:
|
||||
|
||||
方案1(推荐):
|
||||
pip install psycopg2-binary
|
||||
|
||||
方案2:
|
||||
# Ubuntu/Debian
|
||||
sudo apt-get install postgresql libpq-dev python3-dev
|
||||
pip install psycopg2
|
||||
|
||||
# CentOS/RHEL
|
||||
sudo yum install postgresql postgresql-devel python3-devel
|
||||
pip install psycopg2
|
||||
|
||||
# macOS
|
||||
brew install postgresql
|
||||
pip install psycopg2
|
||||
|
||||
原始错误: {PSYCOPG2_ERROR}
|
||||
"""
|
||||
logger.error(error_msg)
|
||||
raise ImportError(error_msg)
|
||||
|
||||
self.db_url = settings.db
|
||||
self.table_name = "resource_categories"
|
||||
|
||||
# 初始化素材分类表
|
||||
# self._init_resource_category_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_resource_category_table(self):
|
||||
"""初始化素材分类表"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
# 创建素材分类表(如果不存在)
|
||||
create_table_sql = """
|
||||
CREATE TABLE IF NOT EXISTS resource_categories (
|
||||
id VARCHAR(36) PRIMARY KEY,
|
||||
title VARCHAR(255) NOT NULL,
|
||||
ai_prompt TEXT DEFAULT '',
|
||||
color VARCHAR(7) DEFAULT '#FF5733',
|
||||
is_cloud BOOLEAN DEFAULT FALSE,
|
||||
user_id VARCHAR(36) NOT NULL,
|
||||
is_active BOOLEAN DEFAULT TRUE,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
"""
|
||||
cursor.execute(create_table_sql)
|
||||
|
||||
# 创建索引
|
||||
indexes = [
|
||||
"CREATE INDEX IF NOT EXISTS idx_resource_categories_title ON resource_categories(title);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_resource_categories_user_id ON resource_categories(user_id);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_resource_categories_is_cloud ON resource_categories(is_cloud);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_resource_categories_is_active ON resource_categories(is_active);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_resource_categories_created_at ON resource_categories(created_at);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_resource_categories_user_title ON resource_categories(user_id, title);"
|
||||
]
|
||||
|
||||
for index_sql in indexes:
|
||||
cursor.execute(index_sql)
|
||||
|
||||
conn.commit()
|
||||
logger.info("Resource category table initialized")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize resource category table: {e}")
|
||||
raise e
|
||||
|
||||
def create_category(self, category: ResourceCategory) -> str:
|
||||
"""
|
||||
创建素材分类
|
||||
|
||||
Args:
|
||||
category: 素材分类信息
|
||||
|
||||
Returns:
|
||||
分类ID
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
# 检查分类名称是否已存在(同一用户下)
|
||||
cursor.execute(
|
||||
"SELECT id FROM resource_categories WHERE title = %s AND user_id = %s AND is_active = TRUE",
|
||||
(category.title, category.user_id)
|
||||
)
|
||||
if cursor.fetchone():
|
||||
raise ValueError(f"Category title '{category.title}' already exists for this user")
|
||||
|
||||
# 生成分类ID(如果没有提供)
|
||||
category_id = category.id or str(uuid.uuid4())
|
||||
|
||||
# 插入分类记录
|
||||
insert_sql = """
|
||||
INSERT INTO resource_categories (
|
||||
id, title, ai_prompt, color, is_cloud, user_id, is_active, created_at, updated_at
|
||||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
"""
|
||||
|
||||
now = datetime.now()
|
||||
cursor.execute(insert_sql, (
|
||||
category_id,
|
||||
category.title,
|
||||
category.ai_prompt,
|
||||
category.color,
|
||||
category.is_cloud,
|
||||
category.user_id,
|
||||
category.is_active,
|
||||
now,
|
||||
now
|
||||
))
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"Created resource category: {category.title} (ID: {category_id})")
|
||||
return category_id
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create resource category '{category.title}': {e}")
|
||||
raise e
|
||||
|
||||
def get_category_by_id(self, category_id: str) -> Optional[ResourceCategory]:
|
||||
"""
|
||||
根据ID获取素材分类
|
||||
|
||||
Args:
|
||||
category_id: 分类ID
|
||||
|
||||
Returns:
|
||||
分类信息,如果不存在返回None
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
||||
cursor.execute(
|
||||
"SELECT * FROM resource_categories WHERE id = %s AND is_active = TRUE",
|
||||
(category_id,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row:
|
||||
return self._row_to_category(row)
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get resource category by ID '{category_id}': {e}")
|
||||
return None
|
||||
|
||||
def get_category_by_title(self, title: str, user_id: str = None) -> Optional[ResourceCategory]:
|
||||
"""
|
||||
根据标题获取素材分类
|
||||
|
||||
Args:
|
||||
title: 分类标题
|
||||
user_id: 用户ID(可选,用于过滤用户分类)
|
||||
|
||||
Returns:
|
||||
分类信息,如果不存在返回None
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
||||
if user_id:
|
||||
cursor.execute(
|
||||
"SELECT * FROM resource_categories WHERE title = %s AND user_id = %s AND is_active = TRUE LIMIT 1",
|
||||
(title, user_id)
|
||||
)
|
||||
else:
|
||||
cursor.execute(
|
||||
"SELECT * FROM resource_categories WHERE title = %s AND is_active = TRUE LIMIT 1",
|
||||
(title,)
|
||||
)
|
||||
|
||||
row = cursor.fetchone()
|
||||
if row:
|
||||
return self._row_to_category(row)
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get resource category by title '{title}': {e}")
|
||||
return None
|
||||
|
||||
def update_category(self, category_id: str, updates: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
更新素材分类信息
|
||||
|
||||
Args:
|
||||
category_id: 分类ID
|
||||
updates: 要更新的字段
|
||||
|
||||
Returns:
|
||||
更新成功返回True
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
# 检查分类是否存在(不限制 is_active 状态)
|
||||
cursor.execute(
|
||||
"SELECT id FROM resource_categories WHERE id = %s",
|
||||
(category_id,)
|
||||
)
|
||||
if not cursor.fetchone():
|
||||
logger.warning(f"Resource category not found: {category_id}")
|
||||
return False
|
||||
|
||||
# 移除不应该直接更新的字段
|
||||
protected_fields = ["id", "created_at"]
|
||||
filtered_updates = {k: v for k, v in updates.items() if k not in protected_fields}
|
||||
|
||||
if not filtered_updates:
|
||||
logger.warning("No valid fields to update")
|
||||
return False
|
||||
|
||||
# 添加更新时间
|
||||
filtered_updates["updated_at"] = datetime.now()
|
||||
|
||||
# 构建更新SQL
|
||||
set_clauses = []
|
||||
values = []
|
||||
for key, value in filtered_updates.items():
|
||||
set_clauses.append(f"{key} = %s")
|
||||
values.append(value)
|
||||
|
||||
values.append(category_id) # WHERE条件的参数
|
||||
|
||||
update_sql = f"UPDATE resource_categories SET {', '.join(set_clauses)} WHERE id = %s"
|
||||
cursor.execute(update_sql, values)
|
||||
|
||||
conn.commit()
|
||||
|
||||
if cursor.rowcount > 0:
|
||||
logger.info(f"Updated resource category: {category_id}")
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"No rows updated for resource category: {category_id}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update resource category '{category_id}': {e}")
|
||||
return False
|
||||
|
||||
def delete_category(self, category_id: str, hard_delete: bool = True) -> bool:
|
||||
"""
|
||||
删除素材分类
|
||||
|
||||
Args:
|
||||
category_id: 分类ID
|
||||
hard_delete: 是否硬删除(True)或软删除(False)
|
||||
|
||||
Returns:
|
||||
删除成功返回True
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
if hard_delete:
|
||||
# 硬删除:直接从数据库中删除记录
|
||||
cursor.execute("DELETE FROM resource_categories WHERE id = %s", (category_id,))
|
||||
else:
|
||||
# 软删除:设置 is_active = FALSE
|
||||
cursor.execute(
|
||||
"UPDATE resource_categories SET is_active = FALSE, updated_at = %s WHERE id = %s",
|
||||
(datetime.now(), category_id)
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
|
||||
if cursor.rowcount > 0:
|
||||
delete_type = "hard deleted" if hard_delete else "soft deleted"
|
||||
logger.info(f"Resource category {delete_type}: {category_id}")
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"No resource category found to delete: {category_id}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete resource category '{category_id}': {e}")
|
||||
return False
|
||||
|
||||
def get_categories_by_user(self, user_id: str, include_cloud: bool = True, limit: int = 100) -> List[ResourceCategory]:
|
||||
"""
|
||||
获取用户的素材分类列表
|
||||
|
||||
Args:
|
||||
user_id: 用户ID
|
||||
include_cloud: 是否包含云端公共分类
|
||||
limit: 最大返回数量
|
||||
|
||||
Returns:
|
||||
分类列表
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
||||
if include_cloud:
|
||||
sql = """
|
||||
SELECT * FROM resource_categories
|
||||
WHERE (user_id = %s OR is_cloud = TRUE) AND is_active = TRUE
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
cursor.execute(sql, (user_id, limit))
|
||||
else:
|
||||
sql = """
|
||||
SELECT * FROM resource_categories
|
||||
WHERE user_id = %s AND is_active = TRUE
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
cursor.execute(sql, (user_id, limit))
|
||||
|
||||
rows = cursor.fetchall()
|
||||
categories = []
|
||||
|
||||
for row in rows:
|
||||
category = self._row_to_category(row)
|
||||
categories.append(category)
|
||||
|
||||
return categories
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get resource categories for user '{user_id}': {e}")
|
||||
return []
|
||||
|
||||
def search_categories(self, query: str, user_id: str = None, include_cloud: bool = True, limit: int = 50) -> List[ResourceCategory]:
|
||||
"""
|
||||
搜索素材分类
|
||||
|
||||
Args:
|
||||
query: 搜索关键词(匹配标题、AI提示词)
|
||||
user_id: 用户ID(可选,用于过滤用户分类)
|
||||
include_cloud: 是否包含云端公共分类
|
||||
limit: 最大返回数量
|
||||
|
||||
Returns:
|
||||
匹配的分类列表
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
||||
search_pattern = f"%{query}%"
|
||||
|
||||
if user_id and include_cloud:
|
||||
sql = """
|
||||
SELECT * FROM resource_categories
|
||||
WHERE (user_id = %s OR is_cloud = TRUE) AND is_active = TRUE
|
||||
AND (title ILIKE %s OR ai_prompt ILIKE %s)
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
cursor.execute(sql, (user_id, search_pattern, search_pattern, limit))
|
||||
elif user_id:
|
||||
sql = """
|
||||
SELECT * FROM resource_categories
|
||||
WHERE user_id = %s AND is_active = TRUE
|
||||
AND (title ILIKE %s OR ai_prompt ILIKE %s)
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
cursor.execute(sql, (user_id, search_pattern, search_pattern, limit))
|
||||
else:
|
||||
sql = """
|
||||
SELECT * FROM resource_categories
|
||||
WHERE is_active = TRUE
|
||||
AND (title ILIKE %s OR ai_prompt ILIKE %s)
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
cursor.execute(sql, (search_pattern, search_pattern, limit))
|
||||
|
||||
rows = cursor.fetchall()
|
||||
categories = []
|
||||
|
||||
for row in rows:
|
||||
category = self._row_to_category(row)
|
||||
categories.append(category)
|
||||
|
||||
return categories
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to search resource categories with query '{query}': {e}")
|
||||
return []
|
||||
|
||||
def get_cloud_categories(self, limit: int = 100) -> List[ResourceCategory]:
|
||||
"""
|
||||
获取云端公共分类
|
||||
|
||||
Args:
|
||||
limit: 最大返回数量
|
||||
|
||||
Returns:
|
||||
云端分类列表
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
||||
sql = """
|
||||
SELECT * FROM resource_categories
|
||||
WHERE is_cloud = TRUE AND is_active = TRUE
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
cursor.execute(sql, (limit,))
|
||||
|
||||
rows = cursor.fetchall()
|
||||
categories = []
|
||||
|
||||
for row in rows:
|
||||
category = self._row_to_category(row)
|
||||
categories.append(category)
|
||||
|
||||
return categories
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get cloud resource categories: {e}")
|
||||
return []
|
||||
|
||||
def get_category_count(self, user_id: str = None, include_cloud: bool = True) -> int:
|
||||
"""
|
||||
获取分类数量
|
||||
|
||||
Args:
|
||||
user_id: 用户ID(可选)
|
||||
include_cloud: 是否包含云端公共分类
|
||||
|
||||
Returns:
|
||||
分类数量
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
if user_id is None:
|
||||
cursor.execute("SELECT COUNT(*) FROM resource_categories WHERE is_active = TRUE")
|
||||
elif include_cloud:
|
||||
cursor.execute(
|
||||
"SELECT COUNT(*) FROM resource_categories WHERE (user_id = %s OR is_cloud = TRUE) AND is_active = TRUE",
|
||||
(user_id,)
|
||||
)
|
||||
else:
|
||||
cursor.execute(
|
||||
"SELECT COUNT(*) FROM resource_categories WHERE user_id = %s AND is_active = TRUE",
|
||||
(user_id,)
|
||||
)
|
||||
|
||||
result = cursor.fetchone()
|
||||
return result[0] if result else 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get resource category count: {e}")
|
||||
return 0
|
||||
|
||||
def batch_create_categories(self, categories: List[ResourceCategory]) -> Dict[str, Any]:
|
||||
"""
|
||||
批量创建分类
|
||||
|
||||
Args:
|
||||
categories: 分类列表
|
||||
|
||||
Returns:
|
||||
创建结果统计
|
||||
"""
|
||||
try:
|
||||
success_count = 0
|
||||
failed_count = 0
|
||||
failed_categories = []
|
||||
|
||||
for category in categories:
|
||||
try:
|
||||
# 检查是否已存在(根据标题和用户ID)
|
||||
existing = self.get_category_by_title(category.title, category.user_id)
|
||||
if existing:
|
||||
logger.warning(f"Category already exists, skipping: {category.title}")
|
||||
continue
|
||||
|
||||
# 创建分类
|
||||
self.create_category(category)
|
||||
success_count += 1
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create category '{category.title}': {e}")
|
||||
failed_count += 1
|
||||
failed_categories.append({
|
||||
"title": category.title,
|
||||
"error": str(e)
|
||||
})
|
||||
|
||||
result = {
|
||||
"total": len(categories),
|
||||
"success": success_count,
|
||||
"failed": failed_count,
|
||||
"failed_categories": failed_categories
|
||||
}
|
||||
|
||||
logger.info(f"Batch create completed: {success_count} success, {failed_count} failed")
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to batch create categories: {e}")
|
||||
return {
|
||||
"total": len(categories),
|
||||
"success": 0,
|
||||
"failed": len(categories),
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
def get_categories_by_color(self, color: str, user_id: str = None, include_cloud: bool = True, limit: int = 50) -> List[ResourceCategory]:
|
||||
"""
|
||||
根据颜色获取分类
|
||||
|
||||
Args:
|
||||
color: 颜色值(hex格式)
|
||||
user_id: 用户ID(可选)
|
||||
include_cloud: 是否包含云端公共分类
|
||||
limit: 最大返回数量
|
||||
|
||||
Returns:
|
||||
匹配的分类列表
|
||||
"""
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
||||
if user_id and include_cloud:
|
||||
sql = """
|
||||
SELECT * FROM resource_categories
|
||||
WHERE (user_id = %s OR is_cloud = TRUE) AND is_active = TRUE
|
||||
AND color = %s
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
cursor.execute(sql, (user_id, color, limit))
|
||||
elif user_id:
|
||||
sql = """
|
||||
SELECT * FROM resource_categories
|
||||
WHERE user_id = %s AND is_active = TRUE
|
||||
AND color = %s
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
cursor.execute(sql, (user_id, color, limit))
|
||||
else:
|
||||
sql = """
|
||||
SELECT * FROM resource_categories
|
||||
WHERE is_active = TRUE AND color = %s
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
cursor.execute(sql, (color, limit))
|
||||
|
||||
rows = cursor.fetchall()
|
||||
categories = []
|
||||
|
||||
for row in rows:
|
||||
category = self._row_to_category(row)
|
||||
categories.append(category)
|
||||
|
||||
return categories
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get categories by color '{color}': {e}")
|
||||
return []
|
||||
|
||||
def activate_category(self, category_id: str) -> bool:
|
||||
"""
|
||||
激活分类(设置 is_active = TRUE)
|
||||
|
||||
Args:
|
||||
category_id: 分类ID
|
||||
|
||||
Returns:
|
||||
激活成功返回True
|
||||
"""
|
||||
try:
|
||||
return self.update_category(category_id, {"is_active": True})
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to activate category '{category_id}': {e}")
|
||||
return False
|
||||
|
||||
def deactivate_category(self, category_id: str) -> bool:
|
||||
"""
|
||||
停用分类(设置 is_active = FALSE)
|
||||
|
||||
Args:
|
||||
category_id: 分类ID
|
||||
|
||||
Returns:
|
||||
停用成功返回True
|
||||
"""
|
||||
try:
|
||||
return self.update_category(category_id, {"is_active": False})
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to deactivate category '{category_id}': {e}")
|
||||
return False
|
||||
|
||||
# 辅助方法
|
||||
def _row_to_category(self, row: Dict[str, Any]) -> ResourceCategory:
|
||||
"""将数据库行转换为ResourceCategory对象"""
|
||||
return ResourceCategory(
|
||||
id=row['id'],
|
||||
title=row['title'],
|
||||
ai_prompt=row['ai_prompt'],
|
||||
color=row['color'],
|
||||
is_cloud=row['is_cloud'],
|
||||
user_id=row['user_id'],
|
||||
is_active=row['is_active'],
|
||||
created_at=row['created_at'].isoformat() if row['created_at'] else "",
|
||||
updated_at=row['updated_at'].isoformat() if row['updated_at'] else ""
|
||||
)
|
||||
|
||||
|
||||
# 创建全局素材分类表实例
|
||||
resource_category_table = ResourceCategoryPostgres()
|
||||
@@ -57,7 +57,7 @@ PostgreSQL 驱动 psycopg2 未安装。请安装:
|
||||
self.table_name = "templates"
|
||||
|
||||
# 初始化模板表
|
||||
self._init_template_table()
|
||||
# self._init_template_table()
|
||||
|
||||
@contextmanager
|
||||
def _get_connection(self):
|
||||
|
||||
@@ -56,7 +56,7 @@ PostgreSQL 驱动 psycopg2 未安装。请安装:
|
||||
self.table_name = "users"
|
||||
|
||||
# 初始化用户表
|
||||
self._init_user_table()
|
||||
# self._init_user_table()
|
||||
|
||||
@contextmanager
|
||||
def _get_connection(self):
|
||||
|
||||
@@ -12,19 +12,8 @@ from datetime import datetime
|
||||
|
||||
from python_core.config import settings
|
||||
from python_core.utils.logger import logger
|
||||
from python_core.utils.jsonrpc import create_response_handler
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResourceCategory:
|
||||
"""素材分类数据结构"""
|
||||
id: str
|
||||
title: str # 分类标题
|
||||
ai_prompt: str # AI识别提示词
|
||||
color: str # 展示颜色 (hex格式,如 #FF5733)
|
||||
created_at: str
|
||||
updated_at: str
|
||||
is_active: bool = True
|
||||
from python_core.database.types import ResourceCategory
|
||||
|
||||
|
||||
class ResourceCategoryManager:
|
||||
|
||||
350
python_core/services/resource_category_manager_v2.py
Normal file
350
python_core/services/resource_category_manager_v2.py
Normal file
@@ -0,0 +1,350 @@
|
||||
"""
|
||||
素材分类管理服务 V2
|
||||
基于 PostgreSQL 数据库的素材分类管理,提供更好的性能和数据一致性
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from typing import List, Dict, Optional
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime
|
||||
|
||||
from python_core.utils.logger import logger
|
||||
from python_core.database.types import ResourceCategory
|
||||
from python_core.database.resource_category_postgres import resource_category_table
|
||||
|
||||
|
||||
class ResourceCategoryManagerV2:
|
||||
"""素材分类管理器 V2 - 基于 PostgreSQL"""
|
||||
|
||||
def __init__(self, user_id: str = "default"):
|
||||
"""
|
||||
初始化分类管理器
|
||||
|
||||
Args:
|
||||
user_id: 用户ID,用于权限控制
|
||||
"""
|
||||
self.user_id = user_id
|
||||
self.category_table = resource_category_table
|
||||
|
||||
def get_all_categories(self, include_cloud: bool = True) -> List[Dict]:
|
||||
"""
|
||||
获取所有分类
|
||||
|
||||
Args:
|
||||
include_cloud: 是否包含云端公共分类
|
||||
|
||||
Returns:
|
||||
分类列表
|
||||
"""
|
||||
try:
|
||||
categories = self.category_table.get_categories_by_user(
|
||||
self.user_id,
|
||||
include_cloud=include_cloud
|
||||
)
|
||||
return [asdict(category) for category in categories]
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get all categories: {e}")
|
||||
return []
|
||||
|
||||
def get_category_by_id(self, category_id: str) -> Optional[Dict]:
|
||||
"""
|
||||
根据ID获取分类
|
||||
|
||||
Args:
|
||||
category_id: 分类ID
|
||||
|
||||
Returns:
|
||||
分类信息,如果不存在返回None
|
||||
"""
|
||||
try:
|
||||
category = self.category_table.get_category_by_id(category_id)
|
||||
if category and category.is_active:
|
||||
return asdict(category)
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get category by ID {category_id}: {e}")
|
||||
return None
|
||||
|
||||
def get_category_by_title(self, title: str) -> Optional[Dict]:
|
||||
"""
|
||||
根据标题获取分类
|
||||
|
||||
Args:
|
||||
title: 分类标题
|
||||
|
||||
Returns:
|
||||
分类信息,如果不存在返回None
|
||||
"""
|
||||
try:
|
||||
category = self.category_table.get_category_by_title(title, self.user_id)
|
||||
if category and category.is_active:
|
||||
return asdict(category)
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get category by title {title}: {e}")
|
||||
return None
|
||||
|
||||
def create_category(self, title: str, ai_prompt: str = "", color: str = "#FF5733",
|
||||
is_cloud: bool = False) -> Dict:
|
||||
"""
|
||||
创建新分类
|
||||
|
||||
Args:
|
||||
title: 分类标题
|
||||
ai_prompt: AI识别提示词
|
||||
color: 展示颜色(hex格式)
|
||||
is_cloud: 是否为云端公共分类
|
||||
|
||||
Returns:
|
||||
创建的分类信息
|
||||
"""
|
||||
try:
|
||||
now = datetime.now().isoformat()
|
||||
new_category = ResourceCategory(
|
||||
id=str(uuid.uuid4()),
|
||||
title=title,
|
||||
ai_prompt=ai_prompt,
|
||||
color=color,
|
||||
is_cloud=is_cloud,
|
||||
user_id=self.user_id,
|
||||
is_active=True,
|
||||
created_at=now,
|
||||
updated_at=now
|
||||
)
|
||||
|
||||
category_id = self.category_table.create_category(new_category)
|
||||
|
||||
# 获取创建后的分类信息
|
||||
created_category = self.category_table.get_category_by_id(category_id)
|
||||
if created_category:
|
||||
logger.info(f"Created new category: {title} (ID: {category_id})")
|
||||
return asdict(created_category)
|
||||
else:
|
||||
raise Exception("Failed to retrieve created category")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create category {title}: {e}")
|
||||
raise
|
||||
|
||||
def update_category(self, category_id: str, title: str = None,
|
||||
ai_prompt: str = None, color: str = None,
|
||||
is_active: bool = None) -> Optional[Dict]:
|
||||
"""
|
||||
更新分类
|
||||
|
||||
Args:
|
||||
category_id: 分类ID
|
||||
title: 新标题
|
||||
ai_prompt: 新AI提示词
|
||||
color: 新颜色
|
||||
is_active: 新状态
|
||||
|
||||
Returns:
|
||||
更新后的分类信息,如果失败返回None
|
||||
"""
|
||||
try:
|
||||
# 构建更新字典
|
||||
updates = {}
|
||||
if title is not None:
|
||||
updates["title"] = title
|
||||
if ai_prompt is not None:
|
||||
updates["ai_prompt"] = ai_prompt
|
||||
if color is not None:
|
||||
updates["color"] = color
|
||||
if is_active is not None:
|
||||
updates["is_active"] = is_active
|
||||
|
||||
if not updates:
|
||||
logger.warning("No fields to update")
|
||||
return None
|
||||
|
||||
# 执行更新
|
||||
success = self.category_table.update_category(category_id, updates)
|
||||
if success:
|
||||
# 获取更新后的分类信息
|
||||
updated_category = self.category_table.get_category_by_id(category_id)
|
||||
if updated_category:
|
||||
logger.info(f"Updated category: {category_id}")
|
||||
return asdict(updated_category)
|
||||
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update category {category_id}: {e}")
|
||||
return None
|
||||
|
||||
def delete_category(self, category_id: str, hard_delete: bool = True) -> bool:
|
||||
"""
|
||||
删除分类
|
||||
|
||||
Args:
|
||||
category_id: 分类ID
|
||||
hard_delete: 是否硬删除(True)或软删除(False)
|
||||
|
||||
Returns:
|
||||
删除成功返回True
|
||||
"""
|
||||
try:
|
||||
success = self.category_table.delete_category(category_id, hard_delete=hard_delete)
|
||||
if success:
|
||||
delete_type = "hard deleted" if hard_delete else "soft deleted"
|
||||
logger.info(f"Category {delete_type}: {category_id}")
|
||||
return success
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete category {category_id}: {e}")
|
||||
return False
|
||||
|
||||
def search_categories(self, keyword: str, include_cloud: bool = True) -> List[Dict]:
|
||||
"""
|
||||
搜索分类
|
||||
|
||||
Args:
|
||||
keyword: 搜索关键词
|
||||
include_cloud: 是否包含云端公共分类
|
||||
|
||||
Returns:
|
||||
匹配的分类列表
|
||||
"""
|
||||
try:
|
||||
categories = self.category_table.search_categories(
|
||||
keyword,
|
||||
self.user_id,
|
||||
include_cloud=include_cloud
|
||||
)
|
||||
return [asdict(category) for category in categories]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to search categories with keyword {keyword}: {e}")
|
||||
return []
|
||||
|
||||
def get_categories_by_color(self, color: str, include_cloud: bool = True) -> List[Dict]:
|
||||
"""
|
||||
根据颜色获取分类
|
||||
|
||||
Args:
|
||||
color: 颜色值(hex格式)
|
||||
include_cloud: 是否包含云端公共分类
|
||||
|
||||
Returns:
|
||||
匹配的分类列表
|
||||
"""
|
||||
try:
|
||||
categories = self.category_table.get_categories_by_color(
|
||||
color,
|
||||
self.user_id,
|
||||
include_cloud=include_cloud
|
||||
)
|
||||
return [asdict(category) for category in categories]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get categories by color {color}: {e}")
|
||||
return []
|
||||
|
||||
def get_cloud_categories(self) -> List[Dict]:
|
||||
"""
|
||||
获取云端公共分类
|
||||
|
||||
Returns:
|
||||
云端分类列表
|
||||
"""
|
||||
try:
|
||||
categories = self.category_table.get_cloud_categories()
|
||||
return [asdict(category) for category in categories]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get cloud categories: {e}")
|
||||
return []
|
||||
|
||||
def get_category_count(self, include_cloud: bool = True) -> int:
|
||||
"""
|
||||
获取分类数量
|
||||
|
||||
Args:
|
||||
include_cloud: 是否包含云端公共分类
|
||||
|
||||
Returns:
|
||||
分类数量
|
||||
"""
|
||||
try:
|
||||
return self.category_table.get_category_count(self.user_id, include_cloud)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get category count: {e}")
|
||||
return 0
|
||||
|
||||
def activate_category(self, category_id: str) -> bool:
|
||||
"""
|
||||
激活分类
|
||||
|
||||
Args:
|
||||
category_id: 分类ID
|
||||
|
||||
Returns:
|
||||
激活成功返回True
|
||||
"""
|
||||
try:
|
||||
return self.category_table.activate_category(category_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to activate category {category_id}: {e}")
|
||||
return False
|
||||
|
||||
def deactivate_category(self, category_id: str) -> bool:
|
||||
"""
|
||||
停用分类
|
||||
|
||||
Args:
|
||||
category_id: 分类ID
|
||||
|
||||
Returns:
|
||||
停用成功返回True
|
||||
"""
|
||||
try:
|
||||
return self.category_table.deactivate_category(category_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to deactivate category {category_id}: {e}")
|
||||
return False
|
||||
|
||||
def batch_create_categories(self, categories_data: List[Dict]) -> Dict:
|
||||
"""
|
||||
批量创建分类
|
||||
|
||||
Args:
|
||||
categories_data: 分类数据列表,每个元素包含 title, ai_prompt, color 等字段
|
||||
|
||||
Returns:
|
||||
创建结果统计
|
||||
"""
|
||||
try:
|
||||
categories = []
|
||||
now = datetime.now().isoformat()
|
||||
|
||||
for data in categories_data:
|
||||
category = ResourceCategory(
|
||||
id=str(uuid.uuid4()),
|
||||
title=data.get("title", ""),
|
||||
ai_prompt=data.get("ai_prompt", ""),
|
||||
color=data.get("color", "#FF5733"),
|
||||
is_cloud=data.get("is_cloud", False),
|
||||
user_id=self.user_id,
|
||||
is_active=True,
|
||||
created_at=now,
|
||||
updated_at=now
|
||||
)
|
||||
categories.append(category)
|
||||
|
||||
result = self.category_table.batch_create_categories(categories)
|
||||
logger.info(f"Batch create completed: {result['success']} success, {result['failed']} failed")
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to batch create categories: {e}")
|
||||
return {
|
||||
"total": len(categories_data),
|
||||
"success": 0,
|
||||
"failed": len(categories_data),
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
|
||||
# 创建全局管理器实例(使用默认用户)
|
||||
resource_category_manager_v2 = ResourceCategoryManagerV2()
|
||||
Reference in New Issue
Block a user