2025-03-17 14:58:05 +08:00

275 lines
10 KiB
Python

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
数据库会话管理模块
提供数据库会话获取和管理功能
"""
from contextlib import contextmanager
from config.database import DBConfig, CacheConfig
def get_session():
"""
获取数据库会话
Returns:
Session: 数据库会话对象
"""
return DBConfig.get_session()
@contextmanager
def session_scope():
"""
会话上下文管理器
提供自动提交和异常回滚功能
Yields:
Session: 数据库会话对象
"""
session = get_session()
try:
yield session
session.commit()
except Exception as e:
session.rollback()
raise e
finally:
session.close()
def initialize_database():
"""
初始化数据库
创建所有表并初始化基础数据
"""
# 创建所有表
DBConfig.init_db()
# 初始化基础数据
_initialize_component_categories()
_initialize_component_types()
_initialize_system_components()
print("数据库初始化完成")
def _initialize_component_categories():
"""初始化组件类别"""
from data.models.component import ComponentCategory
# 检查是否已存在数据
if ComponentCategory.query.count() > 0:
return
# 定义基础组件类别
categories = [
{"name": "基础组件", "code": "basic", "description": "基础功能组件", "sort_order": 1},
{"name": "流程组件", "code": "flow", "description": "流程控制组件", "sort_order": 2},
{"name": "机器人组件", "code": "robot", "description": "机器人控制组件", "sort_order": 3},
{"name": "库位组件", "code": "storage", "description": "库位管理组件", "sort_order": 4},
{"name": "设备组件", "code": "device", "description": "设备控制组件", "sort_order": 5},
{"name": "HTTP组件", "code": "http", "description": "HTTP请求组件", "sort_order": 6},
{"name": "脚本组件", "code": "script", "description": "脚本执行组件", "sort_order": 7},
{"name": "子任务组件", "code": "subtask", "description": "子任务组件", "sort_order": 8},
{"name": "任务组件", "code": "task", "description": "任务管理组件", "sort_order": 9}
]
# 批量创建
with session_scope() as session:
for category_data in categories:
category = ComponentCategory(**category_data)
session.add(category)
# 缓存组件类别
_cache_component_categories()
def _initialize_component_types():
"""初始化组件类型"""
from data.models.component import ComponentType, ComponentCategory
# 检查是否已存在数据
if ComponentType.query.count() > 0:
return
# 获取组件类别
categories = {category.code: category.id for category in ComponentCategory.query.all()}
# 定义基础组件类型
types = [
# 基础组件
{"name": "变量赋值", "code": "variable_assign", "category_id": categories["basic"], "description": "为变量赋值", "icon": "variable", "sort_order": 1},
{"name": "条件判断", "code": "condition", "category_id": categories["basic"], "description": "条件判断", "icon": "condition", "sort_order": 2},
{"name": "延时等待", "code": "delay", "category_id": categories["basic"], "description": "延时等待", "icon": "delay", "sort_order": 3},
# 流程组件
{"name": "If条件", "code": "if", "category_id": categories["flow"], "description": "If条件判断", "icon": "if", "sort_order": 1},
{"name": "If-Else条件", "code": "if_else", "category_id": categories["flow"], "description": "If-Else条件判断", "icon": "if_else", "sort_order": 2},
{"name": "循环", "code": "loop", "category_id": categories["flow"], "description": "循环执行", "icon": "loop", "sort_order": 3},
{"name": "并行执行", "code": "parallel", "category_id": categories["flow"], "description": "并行执行多个分支", "icon": "parallel", "sort_order": 4},
# 机器人组件
{"name": "选择机器人", "code": "select_robot", "category_id": categories["robot"], "description": "选择执行机器人", "icon": "robot", "sort_order": 1},
{"name": "机器人移动", "code": "robot_move", "category_id": categories["robot"], "description": "控制机器人移动", "icon": "move", "sort_order": 2},
{"name": "获取机器人状态", "code": "robot_status", "category_id": categories["robot"], "description": "获取机器人状态", "icon": "status", "sort_order": 3},
# HTTP组件
{"name": "HTTP请求", "code": "http_request", "category_id": categories["http"], "description": "发送HTTP请求", "icon": "http", "sort_order": 1},
{"name": "API调用", "code": "api_call", "category_id": categories["http"], "description": "调用系统API", "icon": "api", "sort_order": 2},
# 脚本组件
{"name": "JavaScript脚本", "code": "javascript", "category_id": categories["script"], "description": "执行JavaScript脚本", "icon": "script", "sort_order": 1},
# 子任务组件
{"name": "执行子任务", "code": "execute_subtask", "category_id": categories["subtask"], "description": "执行子任务", "icon": "subtask", "sort_order": 1}
]
# 批量创建
with session_scope() as session:
for type_data in types:
component_type = ComponentType(**type_data)
session.add(component_type)
# 缓存组件类型
_cache_component_types()
def _initialize_system_components():
"""初始化系统组件"""
from data.models.component import Component, ComponentType
# 检查是否已存在数据
if Component.query.count() > 0:
return
# 获取组件类型
types = {component_type.code: component_type.id for component_type in ComponentType.query.all()}
# 定义系统组件
components = [
{
"name": "变量赋值",
"code": "variable_assign",
"type_id": types["variable_assign"],
"description": "为变量赋值",
"is_system": True,
"config_schema": {
"type": "object",
"properties": {
"variable_name": {"type": "string", "title": "变量名"},
"value_type": {"type": "string", "enum": ["string", "number", "boolean", "object"], "title": "值类型"},
"value": {"type": "string", "title": ""}
},
"required": ["variable_name", "value_type", "value"]
},
"input_schema": {},
"output_schema": {
"type": "object",
"properties": {
"result": {"type": "boolean", "title": "执行结果"}
}
}
},
{
"name": "条件判断",
"code": "condition",
"type_id": types["condition"],
"description": "条件判断",
"is_system": True,
"config_schema": {
"type": "object",
"properties": {
"condition": {"type": "string", "title": "条件表达式"}
},
"required": ["condition"]
},
"input_schema": {},
"output_schema": {
"type": "object",
"properties": {
"result": {"type": "boolean", "title": "判断结果"}
}
}
},
{
"name": "If条件",
"code": "if",
"type_id": types["if"],
"description": "If条件判断",
"is_system": True,
"config_schema": {
"type": "object",
"properties": {
"condition": {"type": "string", "title": "条件表达式"}
},
"required": ["condition"]
},
"input_schema": {},
"output_schema": {
"type": "object",
"properties": {
"result": {"type": "boolean", "title": "判断结果"}
}
}
}
]
# 批量创建
with session_scope() as session:
for component_data in components:
component = Component(**component_data)
session.add(component)
# 缓存系统组件
_cache_system_components()
def _cache_component_categories():
"""缓存组件类别"""
from data.models.component import ComponentCategory
categories = ComponentCategory.query.all()
categories_dict = {category.id: {
"id": category.id,
"name": category.name,
"code": category.code,
"description": category.description,
"sort_order": category.sort_order
} for category in categories}
# 使用Redis缓存
CacheConfig.set("component_categories", categories_dict, expire=86400) # 缓存24小时
def _cache_component_types():
"""缓存组件类型"""
from data.models.component import ComponentType
types = ComponentType.query.all()
types_dict = {type.id: {
"id": type.id,
"name": type.name,
"code": type.code,
"category_id": type.category_id,
"description": type.description,
"icon": type.icon,
"sort_order": type.sort_order
} for type in types}
# 使用Redis缓存
CacheConfig.set("component_types", types_dict, expire=86400) # 缓存24小时
def _cache_system_components():
"""缓存系统组件"""
from data.models.component import Component
components = Component.query.filter(Component.is_system == True).all()
components_dict = {component.id: {
"id": component.id,
"name": component.name,
"code": component.code,
"type_id": component.type_id,
"description": component.description,
"config_schema": component.config_schema,
"input_schema": component.input_schema,
"output_schema": component.output_schema
} for component in components}
# 使用Redis缓存
CacheConfig.set("system_components", components_dict, expire=86400) # 缓存24小时