use sea_orm::{ConnectionTrait, DatabaseConnection};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tracing::{error, info};
#[derive(Error, Debug)]
pub enum SchemaError {
#[error("数据库查询失败: {0}")]
QueryFailed(String),
#[error("Schema 操作失败: {0}")]
OperationFailed(String),
#[error("Schema 不存在: {0}")]
SchemaNotFound(String),
#[error("Schema 已存在: {0}")]
SchemaExists(String),
#[error("无效的输入: {0}")]
InvalidInput(String),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SchemaInfo {
pub schema_name: String,
pub schema_owner: Option<String>,
pub schema_comment: Option<String>,
}
pub struct SchemaService;
impl SchemaService {
pub async fn list_schemas(db: &DatabaseConnection) -> Result<Vec<SchemaInfo>, SchemaError> {
let backend = db.get_database_backend();
if backend != sea_orm::DatabaseBackend::Postgres {
return Err(SchemaError::QueryFailed(
format!("当前仅支持 PostgreSQL 数据库,当前数据库类型: {:?}", backend)
));
}
let _ = db.execute_unprepared("SELECT 1").await.map_err(|e| {
SchemaError::QueryFailed(format!("数据库连接不可用: {}", e))
})?;
Err(SchemaError::QueryFailed(
"Schema 列表查询功能需要 sqlx 连接池支持。\
建议:1) 使用 sqlx 连接池直接查询,或 2) 通过数据库管理工具查询".to_string(),
))
}
pub async fn get_schema_info(
db: &DatabaseConnection,
schema_name: &str,
) -> Result<SchemaInfo, SchemaError> {
if schema_name.is_empty() {
return Err(SchemaError::InvalidInput("Schema 名称不能为空".to_string()));
}
let backend = db.get_database_backend();
if backend != sea_orm::DatabaseBackend::Postgres {
return Err(SchemaError::QueryFailed(
format!("当前仅支持 PostgreSQL 数据库,当前数据库类型: {:?}", backend)
));
}
let _ = db.execute_unprepared("SELECT 1").await.map_err(|e| {
SchemaError::QueryFailed(format!("数据库连接不可用: {}", e))
})?;
Err(SchemaError::QueryFailed(
format!(
"Schema 信息查询功能需要 sqlx 连接池支持。\
建议:1) 使用 sqlx 连接池直接查询,或 2) 通过数据库管理工具查询。\
要查询的 Schema: {}",
schema_name
)
))
}
pub async fn schema_exists(
db: &DatabaseConnection,
schema_name: &str,
) -> Result<bool, SchemaError> {
match Self::get_schema_info(db, schema_name).await {
Ok(_) => Ok(true),
Err(SchemaError::SchemaNotFound(_)) => Ok(false),
Err(e) => Err(e),
}
}
pub async fn create_schema(
db: &DatabaseConnection,
schema_name: &str,
if_not_exists: bool,
) -> Result<(), SchemaError> {
if Self::schema_exists(db, schema_name).await? {
if if_not_exists {
info!("Schema {} 已存在,跳过创建", schema_name);
return Ok(());
} else {
return Err(SchemaError::SchemaExists(format!(
"Schema {} 已存在",
schema_name
)));
}
}
let sql = format!("CREATE SCHEMA IF NOT EXISTS {}", schema_name);
db.execute_unprepared(sql.as_str()).await.map_err(|e| {
error!("创建 Schema 失败: {}", e);
SchemaError::OperationFailed(format!("创建 Schema 失败: {}", e))
})?;
info!("✓ 成功创建 Schema: {}", schema_name);
Ok(())
}
pub async fn drop_schema(
db: &DatabaseConnection,
schema_name: &str,
if_exists: bool,
cascade: bool,
) -> Result<(), SchemaError> {
if !Self::schema_exists(db, schema_name).await? {
if if_exists {
info!("Schema {} 不存在,跳过删除", schema_name);
return Ok(());
} else {
return Err(SchemaError::SchemaNotFound(format!(
"Schema {} 不存在",
schema_name
)));
}
}
let cascade_str = if cascade { " CASCADE" } else { "" };
let sql = format!("DROP SCHEMA IF EXISTS {}{}", schema_name, cascade_str);
db.execute_unprepared(sql.as_str()).await.map_err(|e| {
error!("删除 Schema 失败: {}", e);
SchemaError::OperationFailed(format!("删除 Schema 失败: {}", e))
})?;
info!("✓ 成功删除 Schema: {}", schema_name);
Ok(())
}
pub async fn set_schema_comment(
db: &DatabaseConnection,
schema_name: &str,
comment: &str,
) -> Result<(), SchemaError> {
if !Self::schema_exists(db, schema_name).await? {
return Err(SchemaError::SchemaNotFound(format!(
"Schema {} 不存在",
schema_name
)));
}
let sql = format!("COMMENT ON SCHEMA {} IS '{}'", schema_name, comment);
db.execute_unprepared(sql.as_str()).await.map_err(|e| {
error!("设置 Schema 注释失败: {}", e);
SchemaError::OperationFailed(format!("设置 Schema 注释失败: {}", e))
})?;
info!("✓ 成功设置 Schema {} 的注释", schema_name);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_schema_info_serialization() {
let info = SchemaInfo {
schema_name: "test_schema".to_string(),
schema_owner: Some("postgres".to_string()),
schema_comment: Some("测试 Schema".to_string()),
};
let json = serde_json::to_string(&info).unwrap();
assert!(json.contains("test_schema"));
}
}