You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
333 lines
14 KiB
333 lines
14 KiB
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
|
|
from sqlalchemy.orm import DeclarativeBase
|
|
from sqlalchemy import text
|
|
from config import settings
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
"""SQLAlchemy ORM 基类，所有数据库模型均继承此类。"""
|
|
pass
|
|
|
|
|
|
# 异步数据库引擎，连接池大小 20，最大溢出 40，启用连接健康检查
|
|
async_engine = create_async_engine(
|
|
settings.DATABASE_URL,
|
|
pool_size=20,
|
|
max_overflow=40,
|
|
pool_pre_ping=True,
|
|
pool_recycle=3600,
|
|
echo=False,
|
|
)
|
|
|
|
# 异步数据库会话工厂，用于创建数据库会话实例
|
|
AsyncSessionLocal = async_sessionmaker(
|
|
async_engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
)
|
|
|
|
|
|
async def init_db():
|
|
"""初始化数据库：创建所有表并执行增量迁移。"""
|
|
async with async_engine.begin() as conn:
|
|
from models import Base as MBase
|
|
await conn.run_sync(MBase.metadata.create_all)
|
|
|
|
await _run_migrations()
|
|
|
|
|
|
async def _run_migrations():
|
|
"""执行数据库增量迁移：创建缺失的表并安全添加新字段。"""
|
|
async with async_engine.begin() as conn:
|
|
# --- 创建可能缺失的新表 ---
|
|
await conn.execute(text("""
|
|
CREATE TABLE IF NOT EXISTS agent_configs (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
name VARCHAR(200) NOT NULL,
|
|
description TEXT,
|
|
system_prompt TEXT,
|
|
model VARCHAR(100) DEFAULT '',
|
|
model_instance_id UUID,
|
|
embedding_model_id UUID,
|
|
temperature FLOAT DEFAULT 0.7,
|
|
tools JSONB DEFAULT '[]',
|
|
status VARCHAR(20) DEFAULT 'active',
|
|
creator_id UUID REFERENCES users(id),
|
|
created_at TIMESTAMP DEFAULT now(),
|
|
updated_at TIMESTAMP DEFAULT now()
|
|
)
|
|
"""))
|
|
await conn.execute(text("""
|
|
CREATE TABLE IF NOT EXISTS model_providers (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
name VARCHAR(100) NOT NULL,
|
|
provider_type VARCHAR(50) NOT NULL,
|
|
base_url VARCHAR(500),
|
|
api_key VARCHAR(500),
|
|
extra_config JSONB DEFAULT '{}',
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
created_at TIMESTAMP DEFAULT now()
|
|
)
|
|
"""))
|
|
await conn.execute(text("""
|
|
CREATE TABLE IF NOT EXISTS model_instances (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
provider_id UUID NOT NULL REFERENCES model_providers(id) ON DELETE CASCADE,
|
|
model_name VARCHAR(200) NOT NULL,
|
|
model_type VARCHAR(20) NOT NULL,
|
|
display_name VARCHAR(200),
|
|
capabilities JSONB DEFAULT '{}',
|
|
default_params JSONB DEFAULT '{}',
|
|
is_default BOOLEAN DEFAULT FALSE,
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
created_at TIMESTAMP DEFAULT now()
|
|
)
|
|
"""))
|
|
await conn.execute(text("""
|
|
CREATE TABLE IF NOT EXISTS custom_tools (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
name VARCHAR(100) UNIQUE NOT NULL,
|
|
description TEXT,
|
|
schema_json JSONB NOT NULL DEFAULT '{}',
|
|
endpoint_url VARCHAR(500) DEFAULT '',
|
|
method VARCHAR(10) DEFAULT 'GET',
|
|
path VARCHAR(500) DEFAULT '',
|
|
headers_json JSONB DEFAULT '{}',
|
|
auth_type VARCHAR(20) DEFAULT 'none',
|
|
auth_config JSONB DEFAULT '{}',
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
created_by UUID REFERENCES users(id),
|
|
created_at TIMESTAMP DEFAULT now(),
|
|
updated_at TIMESTAMP DEFAULT now()
|
|
)
|
|
"""))
|
|
await conn.execute(text("""
|
|
CREATE TABLE IF NOT EXISTS mcp_services (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
name VARCHAR(100) NOT NULL,
|
|
endpoint_url VARCHAR(500) NOT NULL,
|
|
status VARCHAR(20) DEFAULT 'active',
|
|
config JSONB DEFAULT '{}',
|
|
last_heartbeat TIMESTAMP,
|
|
created_at TIMESTAMP DEFAULT now(),
|
|
updated_at TIMESTAMP DEFAULT now()
|
|
)
|
|
"""))
|
|
await conn.execute(text("""
|
|
CREATE TABLE IF NOT EXISTS audit_logs (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
user_id UUID REFERENCES users(id),
|
|
action VARCHAR(50) NOT NULL,
|
|
resource_type VARCHAR(50) NOT NULL,
|
|
resource_id VARCHAR(100),
|
|
detail JSONB DEFAULT '{}',
|
|
ip_address VARCHAR(50),
|
|
user_agent VARCHAR(500),
|
|
created_at TIMESTAMP DEFAULT now()
|
|
)
|
|
"""))
|
|
await conn.execute(text("""
|
|
CREATE TABLE IF NOT EXISTS notification_templates (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
name VARCHAR(100) UNIQUE NOT NULL,
|
|
channel VARCHAR(20) NOT NULL,
|
|
event_type VARCHAR(50) NOT NULL,
|
|
subject_template VARCHAR(500),
|
|
body_template TEXT,
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
created_at TIMESTAMP DEFAULT now(),
|
|
updated_at TIMESTAMP DEFAULT now()
|
|
)
|
|
"""))
|
|
await conn.execute(text("""
|
|
CREATE TABLE IF NOT EXISTS system_metrics (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
metric_name VARCHAR(50) NOT NULL,
|
|
metric_value FLOAT,
|
|
tags JSONB DEFAULT '{}',
|
|
recorded_at TIMESTAMP DEFAULT now()
|
|
)
|
|
"""))
|
|
await conn.execute(text("""
|
|
CREATE TABLE IF NOT EXISTS flow_executions (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
flow_id UUID NOT NULL REFERENCES flow_definitions(id),
|
|
session_id VARCHAR(100),
|
|
status VARCHAR(20) DEFAULT 'running',
|
|
input_data JSONB DEFAULT '{}',
|
|
output_data JSONB DEFAULT '{}',
|
|
error_message TEXT,
|
|
started_by UUID REFERENCES users(id),
|
|
started_at TIMESTAMP DEFAULT now(),
|
|
finished_at TIMESTAMP,
|
|
duration_ms INTEGER
|
|
)
|
|
"""))
|
|
await conn.execute(text("""
|
|
CREATE TABLE IF NOT EXISTS memory_messages (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
flow_execution_id UUID NOT NULL REFERENCES flow_executions(id) ON DELETE CASCADE,
|
|
node_id VARCHAR(100) NOT NULL,
|
|
role VARCHAR(20) NOT NULL,
|
|
content TEXT NOT NULL,
|
|
metadata JSONB DEFAULT '{}',
|
|
created_at TIMESTAMP DEFAULT now()
|
|
)
|
|
"""))
|
|
|
|
# --- 为已有表添加缺失列 ---
|
|
await conn.execute(text(
|
|
"ALTER TABLE flow_definitions ADD COLUMN IF NOT EXISTS published_version_id UUID REFERENCES flow_versions(id)"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE flow_definitions ADD COLUMN IF NOT EXISTS draft_definition_json JSONB"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE flow_definitions ADD COLUMN IF NOT EXISTS flow_mode VARCHAR(20) DEFAULT 'chatflow'"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE flow_definitions ADD COLUMN IF NOT EXISTS published_to_wecom BOOLEAN DEFAULT FALSE"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE flow_definitions ADD COLUMN IF NOT EXISTS published_to_web BOOLEAN DEFAULT FALSE"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE agent_configs ADD COLUMN IF NOT EXISTS model_instance_id UUID"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE agent_configs ADD COLUMN IF NOT EXISTS embedding_model_id UUID"
|
|
))
|
|
# mcp_services 表可能存在旧版本，补齐所有缺失列
|
|
await conn.execute(text(
|
|
"ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS name VARCHAR(100) NOT NULL DEFAULT ''"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS endpoint_url VARCHAR(500) NOT NULL DEFAULT ''"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS status VARCHAR(20) DEFAULT 'active'"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS config JSONB DEFAULT '{}'"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS last_heartbeat TIMESTAMP"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS created_at TIMESTAMP DEFAULT now()"
|
|
))
|
|
await conn.execute(text(
|
|
"ALTER TABLE mcp_services ADD COLUMN IF NOT EXISTS updated_at TIMESTAMP DEFAULT now()"
|
|
))
|
|
# --- 其他新表的列级迁移（防止旧版本表缺少字段） ---
|
|
# model_providers
|
|
for col_sql in [
|
|
"name VARCHAR(100) NOT NULL DEFAULT ''",
|
|
"provider_type VARCHAR(50) NOT NULL DEFAULT 'openai_compatible'",
|
|
"base_url VARCHAR(500) DEFAULT ''",
|
|
"api_key VARCHAR(500) DEFAULT ''",
|
|
"extra_config JSONB DEFAULT '{}'",
|
|
"is_active BOOLEAN DEFAULT TRUE",
|
|
"created_at TIMESTAMP DEFAULT now()",
|
|
]:
|
|
await conn.execute(text(f"ALTER TABLE model_providers ADD COLUMN IF NOT EXISTS {col_sql}"))
|
|
# model_instances
|
|
for col_sql in [
|
|
"id UUID PRIMARY KEY DEFAULT gen_random_uuid()",
|
|
"provider_id UUID NOT NULL REFERENCES model_providers(id) ON DELETE CASCADE",
|
|
"model_name VARCHAR(200) NOT NULL DEFAULT ''",
|
|
"model_type VARCHAR(20) NOT NULL DEFAULT 'llm'",
|
|
"display_name VARCHAR(200) DEFAULT ''",
|
|
"capabilities JSONB DEFAULT '{}'",
|
|
"default_params JSONB DEFAULT '{}'",
|
|
"is_default BOOLEAN DEFAULT FALSE",
|
|
"is_active BOOLEAN DEFAULT TRUE",
|
|
"created_at TIMESTAMP DEFAULT now()",
|
|
]:
|
|
await conn.execute(text(f"ALTER TABLE model_instances ADD COLUMN IF NOT EXISTS {col_sql}"))
|
|
# custom_tools
|
|
for col_sql in [
|
|
"name VARCHAR(100) UNIQUE NOT NULL DEFAULT ''",
|
|
"description TEXT",
|
|
"schema_json JSONB NOT NULL DEFAULT '{}'",
|
|
"endpoint_url VARCHAR(500) DEFAULT ''",
|
|
"method VARCHAR(10) DEFAULT 'GET'",
|
|
"path VARCHAR(500) DEFAULT ''",
|
|
"headers_json JSONB DEFAULT '{}'",
|
|
"auth_type VARCHAR(20) DEFAULT 'none'",
|
|
"auth_config JSONB DEFAULT '{}'",
|
|
"is_active BOOLEAN DEFAULT TRUE",
|
|
"created_by UUID REFERENCES users(id)",
|
|
"created_at TIMESTAMP DEFAULT now()",
|
|
"updated_at TIMESTAMP DEFAULT now()",
|
|
]:
|
|
await conn.execute(text(f"ALTER TABLE custom_tools ADD COLUMN IF NOT EXISTS {col_sql}"))
|
|
# audit_logs
|
|
for col_sql in [
|
|
"user_id UUID REFERENCES users(id)",
|
|
"action VARCHAR(50) NOT NULL DEFAULT ''",
|
|
"resource_type VARCHAR(50) NOT NULL DEFAULT ''",
|
|
"resource_id VARCHAR(100)",
|
|
"detail JSONB DEFAULT '{}'",
|
|
"ip_address VARCHAR(50)",
|
|
"user_agent VARCHAR(500)",
|
|
"created_at TIMESTAMP DEFAULT now()",
|
|
]:
|
|
await conn.execute(text(f"ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS {col_sql}"))
|
|
# notification_templates
|
|
for col_sql in [
|
|
"name VARCHAR(100) UNIQUE NOT NULL DEFAULT ''",
|
|
"channel VARCHAR(20) NOT NULL DEFAULT ''",
|
|
"event_type VARCHAR(50) NOT NULL DEFAULT ''",
|
|
"subject_template VARCHAR(500)",
|
|
"body_template TEXT",
|
|
"is_active BOOLEAN DEFAULT TRUE",
|
|
"created_at TIMESTAMP DEFAULT now()",
|
|
"updated_at TIMESTAMP DEFAULT now()",
|
|
]:
|
|
await conn.execute(text(f"ALTER TABLE notification_templates ADD COLUMN IF NOT EXISTS {col_sql}"))
|
|
# system_metrics
|
|
for col_sql in [
|
|
"metric_name VARCHAR(50) NOT NULL DEFAULT ''",
|
|
"metric_value FLOAT",
|
|
"tags JSONB DEFAULT '{}'",
|
|
"recorded_at TIMESTAMP DEFAULT now()",
|
|
]:
|
|
await conn.execute(text(f"ALTER TABLE system_metrics ADD COLUMN IF NOT EXISTS {col_sql}"))
|
|
# flow_executions
|
|
for col_sql in [
|
|
"flow_id UUID NOT NULL REFERENCES flow_definitions(id)",
|
|
"session_id VARCHAR(100)",
|
|
"status VARCHAR(20) DEFAULT 'running'",
|
|
"input_data JSONB DEFAULT '{}'",
|
|
"output_data JSONB DEFAULT '{}'",
|
|
"error_message TEXT",
|
|
"started_by UUID REFERENCES users(id)",
|
|
"started_at TIMESTAMP DEFAULT now()",
|
|
"finished_at TIMESTAMP",
|
|
"duration_ms INTEGER",
|
|
]:
|
|
await conn.execute(text(f"ALTER TABLE flow_executions ADD COLUMN IF NOT EXISTS {col_sql}"))
|
|
# memory_messages
|
|
for col_sql in [
|
|
"flow_execution_id UUID NOT NULL REFERENCES flow_executions(id) ON DELETE CASCADE",
|
|
"node_id VARCHAR(100) NOT NULL DEFAULT ''",
|
|
"role VARCHAR(20) NOT NULL DEFAULT 'user'",
|
|
"content TEXT NOT NULL DEFAULT ''",
|
|
"metadata JSONB DEFAULT '{}'",
|
|
"created_at TIMESTAMP DEFAULT now()",
|
|
]:
|
|
await conn.execute(text(f"ALTER TABLE memory_messages ADD COLUMN IF NOT EXISTS {col_sql}"))
|
|
|
|
|
|
async def get_db():
|
|
"""FastAPI 依赖注入函数，提供数据库会话，自动提交或回滚事务。"""
|
|
async with AsyncSessionLocal() as session:
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|
|
finally:
|
|
await session.close()
|