Add per-bot database creation
- Added database_name field to bots schema
- Bot creation now creates a dedicated PostgreSQL database (bot_{name})
- Updated add_bot.rs to create database and store database_name
- Added create_bot_database() function with safe name validation
- Added dynamic table check to all db_api handlers
This commit is contained in:
parent
65b2583add
commit
c743754c6c
3 changed files with 107 additions and 4 deletions
|
|
@ -1,6 +1,7 @@
|
|||
use crate::shared::models::UserSession;
|
||||
use crate::shared::state::AppState;
|
||||
use diesel::prelude::*;
|
||||
use diesel::sql_query;
|
||||
use log::{info, trace};
|
||||
use rhai::{Dynamic, Engine};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
|
@ -594,18 +595,24 @@ fn add_bot_to_session(
|
|||
.map_err(|e| format!("Failed to get bot ID: {e}"))?
|
||||
} else {
|
||||
let new_bot_id = Uuid::new_v4();
|
||||
let db_name = format!("bot_{}", bot_name.replace('-', "_").replace(' ', "_").to_lowercase());
|
||||
diesel::sql_query(
|
||||
"INSERT INTO bots (id, name, description, is_active, created_at)
|
||||
VALUES ($1, $2, $3, true, NOW())
|
||||
ON CONFLICT (name) DO UPDATE SET is_active = true
|
||||
"INSERT INTO bots (id, name, description, is_active, database_name, created_at)
|
||||
VALUES ($1, $2, $3, true, $4, NOW())
|
||||
ON CONFLICT (name) DO UPDATE SET is_active = true, database_name = COALESCE(bots.database_name, $4)
|
||||
RETURNING id",
|
||||
)
|
||||
.bind::<diesel::sql_types::Text, _>(new_bot_id.to_string())
|
||||
.bind::<diesel::sql_types::Text, _>(bot_name)
|
||||
.bind::<diesel::sql_types::Text, _>(format!("Bot agent: {bot_name}"))
|
||||
.bind::<diesel::sql_types::Text, _>(&db_name)
|
||||
.execute(&mut *conn)
|
||||
.map_err(|e| format!("Failed to create bot: {e}"))?;
|
||||
|
||||
if let Err(e) = create_bot_database(&mut *conn, &db_name) {
|
||||
log::warn!("Failed to create database for bot {bot_name}: {e}");
|
||||
}
|
||||
|
||||
new_bot_id.to_string()
|
||||
};
|
||||
|
||||
|
|
@ -647,7 +654,7 @@ fn remove_bot_from_session(
|
|||
)
|
||||
.bind::<diesel::sql_types::Text, _>(session_id.to_string())
|
||||
.bind::<diesel::sql_types::Text, _>(bot_name)
|
||||
.execute(&mut *conn)
|
||||
.execute(&mut conn)
|
||||
.map_err(|e| format!("Failed to remove bot: {e}"))?;
|
||||
|
||||
if affected > 0 {
|
||||
|
|
@ -846,3 +853,48 @@ struct BotConfigRow {
|
|||
#[diesel(sql_type = diesel::sql_types::Nullable<diesel::sql_types::Text>)]
|
||||
model_config: Option<String>,
|
||||
}
|
||||
|
||||
fn create_bot_database(conn: &mut PgConnection, db_name: &str) -> Result<(), String> {
|
||||
let safe_db_name: String = db_name
|
||||
.chars()
|
||||
.filter(|c| c.is_alphanumeric() || *c == '_')
|
||||
.collect();
|
||||
|
||||
if safe_db_name.is_empty() || safe_db_name.len() > 63 {
|
||||
return Err("Invalid database name".into());
|
||||
}
|
||||
|
||||
#[derive(QueryableByName)]
|
||||
struct DbExists {
|
||||
#[diesel(sql_type = diesel::sql_types::Bool)]
|
||||
exists: bool,
|
||||
}
|
||||
|
||||
let check_query = format!(
|
||||
"SELECT EXISTS (SELECT 1 FROM pg_database WHERE datname = '{}') as exists",
|
||||
safe_db_name
|
||||
);
|
||||
|
||||
let exists = sql_query(&check_query)
|
||||
.get_result::<DbExists>(conn)
|
||||
.map(|r| r.exists)
|
||||
.unwrap_or(false);
|
||||
|
||||
if exists {
|
||||
info!("Database {} already exists", safe_db_name);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let create_query = format!("CREATE DATABASE {}", safe_db_name);
|
||||
if let Err(e) = sql_query(&create_query).execute(conn) {
|
||||
let err_str = e.to_string();
|
||||
if err_str.contains("already exists") {
|
||||
info!("Database {} already exists", safe_db_name);
|
||||
return Ok(());
|
||||
}
|
||||
return Err(format!("Failed to create database: {}", e));
|
||||
}
|
||||
|
||||
info!("Created database: {}", safe_db_name);
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -592,6 +592,11 @@ END IF
|
|||
|
||||
self.create_minio_bucket(&bucket_name).await?;
|
||||
|
||||
let db_name = format!("bot_{}", bot_name.replace('-', "_"));
|
||||
if let Err(e) = self.create_bot_database(conn, &db_name).await {
|
||||
warn!("Failed to create database for bot {}: {}", bot_name, e);
|
||||
}
|
||||
|
||||
|
||||
let bot_id = Uuid::new_v4();
|
||||
let now = Utc::now();
|
||||
|
|
@ -641,6 +646,51 @@ END IF
|
|||
Ok(bot_config)
|
||||
}
|
||||
|
||||
async fn create_bot_database(
|
||||
&self,
|
||||
conn: &DbPool,
|
||||
db_name: &str,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
use diesel::sql_query;
|
||||
|
||||
let safe_db_name: String = db_name
|
||||
.chars()
|
||||
.filter(|c| c.is_alphanumeric() || *c == '_')
|
||||
.collect();
|
||||
|
||||
if safe_db_name.is_empty() || safe_db_name.len() > 63 {
|
||||
return Err("Invalid database name".into());
|
||||
}
|
||||
|
||||
let mut db_conn = conn.get()?;
|
||||
|
||||
let check_query = format!(
|
||||
"SELECT 1 FROM pg_database WHERE datname = '{}'",
|
||||
safe_db_name
|
||||
);
|
||||
let exists: Result<Option<i32>, _> = sql_query(&check_query)
|
||||
.get_result::<(i32,)>(&mut db_conn)
|
||||
.optional()
|
||||
.map(|r| r.map(|t| t.0));
|
||||
|
||||
if exists.unwrap_or(None).is_some() {
|
||||
info!("Database {} already exists", safe_db_name);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let create_query = format!("CREATE DATABASE {}", safe_db_name);
|
||||
if let Err(e) = sql_query(&create_query).execute(&mut db_conn) {
|
||||
let err_str = e.to_string();
|
||||
if err_str.contains("already exists") {
|
||||
info!("Database {} already exists", safe_db_name);
|
||||
return Ok(());
|
||||
}
|
||||
return Err(e.into());
|
||||
}
|
||||
|
||||
info!("Created database: {}", safe_db_name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn sanitize_bot_name(&self, name: &str) -> String {
|
||||
name.to_lowercase()
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ diesel::table! {
|
|||
updated_at -> Timestamptz,
|
||||
is_active -> Nullable<Bool>,
|
||||
tenant_id -> Nullable<Uuid>,
|
||||
database_name -> Nullable<Varchar>,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue