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::models::UserSession;
|
||||||
use crate::shared::state::AppState;
|
use crate::shared::state::AppState;
|
||||||
use diesel::prelude::*;
|
use diesel::prelude::*;
|
||||||
|
use diesel::sql_query;
|
||||||
use log::{info, trace};
|
use log::{info, trace};
|
||||||
use rhai::{Dynamic, Engine};
|
use rhai::{Dynamic, Engine};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
@ -594,18 +595,24 @@ fn add_bot_to_session(
|
||||||
.map_err(|e| format!("Failed to get bot ID: {e}"))?
|
.map_err(|e| format!("Failed to get bot ID: {e}"))?
|
||||||
} else {
|
} else {
|
||||||
let new_bot_id = Uuid::new_v4();
|
let new_bot_id = Uuid::new_v4();
|
||||||
|
let db_name = format!("bot_{}", bot_name.replace('-', "_").replace(' ', "_").to_lowercase());
|
||||||
diesel::sql_query(
|
diesel::sql_query(
|
||||||
"INSERT INTO bots (id, name, description, is_active, created_at)
|
"INSERT INTO bots (id, name, description, is_active, database_name, created_at)
|
||||||
VALUES ($1, $2, $3, true, NOW())
|
VALUES ($1, $2, $3, true, $4, NOW())
|
||||||
ON CONFLICT (name) DO UPDATE SET is_active = true
|
ON CONFLICT (name) DO UPDATE SET is_active = true, database_name = COALESCE(bots.database_name, $4)
|
||||||
RETURNING id",
|
RETURNING id",
|
||||||
)
|
)
|
||||||
.bind::<diesel::sql_types::Text, _>(new_bot_id.to_string())
|
.bind::<diesel::sql_types::Text, _>(new_bot_id.to_string())
|
||||||
.bind::<diesel::sql_types::Text, _>(bot_name)
|
.bind::<diesel::sql_types::Text, _>(bot_name)
|
||||||
.bind::<diesel::sql_types::Text, _>(format!("Bot agent: {bot_name}"))
|
.bind::<diesel::sql_types::Text, _>(format!("Bot agent: {bot_name}"))
|
||||||
|
.bind::<diesel::sql_types::Text, _>(&db_name)
|
||||||
.execute(&mut *conn)
|
.execute(&mut *conn)
|
||||||
.map_err(|e| format!("Failed to create bot: {e}"))?;
|
.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()
|
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, _>(session_id.to_string())
|
||||||
.bind::<diesel::sql_types::Text, _>(bot_name)
|
.bind::<diesel::sql_types::Text, _>(bot_name)
|
||||||
.execute(&mut *conn)
|
.execute(&mut conn)
|
||||||
.map_err(|e| format!("Failed to remove bot: {e}"))?;
|
.map_err(|e| format!("Failed to remove bot: {e}"))?;
|
||||||
|
|
||||||
if affected > 0 {
|
if affected > 0 {
|
||||||
|
|
@ -846,3 +853,48 @@ struct BotConfigRow {
|
||||||
#[diesel(sql_type = diesel::sql_types::Nullable<diesel::sql_types::Text>)]
|
#[diesel(sql_type = diesel::sql_types::Nullable<diesel::sql_types::Text>)]
|
||||||
model_config: Option<String>,
|
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?;
|
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 bot_id = Uuid::new_v4();
|
||||||
let now = Utc::now();
|
let now = Utc::now();
|
||||||
|
|
@ -641,6 +646,51 @@ END IF
|
||||||
Ok(bot_config)
|
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 {
|
fn sanitize_bot_name(&self, name: &str) -> String {
|
||||||
name.to_lowercase()
|
name.to_lowercase()
|
||||||
|
|
|
||||||
|
|
@ -20,6 +20,7 @@ diesel::table! {
|
||||||
updated_at -> Timestamptz,
|
updated_at -> Timestamptz,
|
||||||
is_active -> Nullable<Bool>,
|
is_active -> Nullable<Bool>,
|
||||||
tenant_id -> Nullable<Uuid>,
|
tenant_id -> Nullable<Uuid>,
|
||||||
|
database_name -> Nullable<Varchar>,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue