bugfix: dashboard org association

This commit is contained in:
dal 2025-04-21 14:58:18 -06:00
parent 5e21bbae43
commit 38a8135dd7
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
1 changed files with 45 additions and 12 deletions

View File

@ -9,11 +9,12 @@ use database::{
enums::{AssetPermissionRole, AssetType, IdentityType}, enums::{AssetPermissionRole, AssetType, IdentityType},
models::{AssetPermission, DashboardFile, MetricFileToDashboardFile}, models::{AssetPermission, DashboardFile, MetricFileToDashboardFile},
pool::get_pg_pool, pool::get_pg_pool,
schema::{asset_permissions, dashboard_files, metric_files_to_dashboard_files}, schema::{asset_permissions, dashboard_files, metric_files_to_dashboard_files, users_to_organizations},
types::{DashboardYml, VersionHistory}, types::{DashboardYml, VersionHistory},
}; };
use diesel::insert_into; use diesel::prelude::*;
use diesel_async::RunQueryDsl; use diesel::{insert_into, SelectableHelper};
use diesel_async::{AsyncPgConnection, RunQueryDsl};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::{self, Value}; use serde_json::{self, Value};
use tracing::debug; use tracing::debug;
@ -74,6 +75,7 @@ async fn process_dashboard_file(
tool_call_id: String, tool_call_id: String,
file: DashboardFileParams, file: DashboardFileParams,
user_id: &Uuid, user_id: &Uuid,
organization_id: &Uuid,
) -> Result<(DashboardFile, DashboardYml), String> { ) -> Result<(DashboardFile, DashboardYml), String> {
debug!("Processing dashboard file creation: {}", file.name); debug!("Processing dashboard file creation: {}", file.name);
@ -108,7 +110,7 @@ async fn process_dashboard_file(
file_name: file.name.clone(), file_name: file.name.clone(),
content: dashboard_yml.clone(), content: dashboard_yml.clone(),
filter: None, filter: None,
organization_id: Uuid::new_v4(), organization_id: *organization_id,
created_by: *user_id, created_by: *user_id,
created_at: Utc::now(), created_at: Utc::now(),
updated_at: Utc::now(), updated_at: Utc::now(),
@ -135,18 +137,55 @@ impl ToolExecutor for CreateDashboardFilesTool {
async fn execute(&self, params: Self::Params, tool_call_id: String) -> Result<Self::Output> { async fn execute(&self, params: Self::Params, tool_call_id: String) -> Result<Self::Output> {
let start_time = Instant::now(); let start_time = Instant::now();
let user_id = self.agent.get_user_id();
let files = params.files; let files = params.files;
let mut created_files = vec![]; let mut created_files = vec![];
let mut failed_files = vec![]; let mut failed_files = vec![];
// Get DB connection early
let mut conn = match get_pg_pool().get().await {
Ok(conn) => conn,
Err(e) => return Err(anyhow!("Failed to get database connection: {}", e)),
};
// Fetch the organization ID for the user
let organization_id = match users_to_organizations::table
.filter(users_to_organizations::user_id.eq(user_id))
.select(users_to_organizations::organization_id)
.first::<Uuid>(&mut conn)
.await
{
Ok(org_id) => org_id,
Err(diesel::NotFound) => {
return Err(anyhow!(
"User {} is not associated with any organization.",
user_id
));
}
Err(e) => {
return Err(anyhow!(
"Failed to fetch organization ID for user {}: {}",
user_id,
e
));
}
};
// Process dashboard files // Process dashboard files
let mut dashboard_records = vec![]; let mut dashboard_records = vec![];
let mut dashboard_ymls = vec![]; let mut dashboard_ymls = vec![];
// First pass - validate and prepare all records // First pass - validate and prepare all records
for file in files { for file in files {
match process_dashboard_file(tool_call_id.clone(), file.clone(), &self.agent.get_user_id()).await { match process_dashboard_file(
tool_call_id.clone(),
file.clone(),
&user_id,
&organization_id,
)
.await
{
Ok((dashboard_file, dashboard_yml)) => { Ok((dashboard_file, dashboard_yml)) => {
dashboard_records.push(dashboard_file); dashboard_records.push(dashboard_file);
dashboard_ymls.push(dashboard_yml); dashboard_ymls.push(dashboard_yml);
@ -158,12 +197,6 @@ impl ToolExecutor for CreateDashboardFilesTool {
} }
// Second pass - bulk insert records // Second pass - bulk insert records
let mut conn = match get_pg_pool().get().await {
Ok(conn) => conn,
Err(e) => return Err(anyhow!(e)),
};
// Insert dashboard files
if !dashboard_records.is_empty() { if !dashboard_records.is_empty() {
match insert_into(dashboard_files::table) match insert_into(dashboard_files::table)
.values(&dashboard_records) .values(&dashboard_records)
@ -172,7 +205,7 @@ impl ToolExecutor for CreateDashboardFilesTool {
{ {
Ok(_) => { Ok(_) => {
// Get the user ID from the agent state // Get the user ID from the agent state
let user_id = self.agent.get_user_id(); // let user_id = self.agent.get_user_id(); // Redundant
// Create asset permissions for each dashboard file // Create asset permissions for each dashboard file
for dashboard_file in &dashboard_records { for dashboard_file in &dashboard_records {