Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

modifications finalizations

This commit is contained in:
dal 2025-04-04 13:07:13 -06:00
parent 773dc9edee
commit 597803ab73
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
6 changed files with 403 additions and 81 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
api/libs/agents/src/tools/categories/file_tools/common.rs
View File

@ -1074,6 +1074,17 @@ pub fn apply_modifications_to_content(
let mut modified_content = content.to_string();
for modification in modifications {
// If content_to_replace is empty, append the new content to the end of the file
if modification.content_to_replace.is_empty() {
// Add a newline at the end if the content doesn't end with one
if !modified_content.ends_with('\n') {
modified_content.push('\n');
}
// Append the new content
modified_content.push_str(&modification.new_content);
continue;
}
// Check if the content to replace exists in the file
if !modified_content.contains(&modification.content_to_replace) {
return Err(anyhow::anyhow!(
@ -1141,6 +1152,30 @@ mod tests {
assert!(err_msg.contains("2 occurrences"));
}
#[test]
fn test_apply_modifications_empty_content_to_replace() {
let original_content = "name: test_metric\ntype: counter\ndescription: A test metric";
// Test appending content when content_to_replace is empty
let mods = vec![Modification {
content_to_replace: "".to_string(),
new_content: "additional_field: true".to_string(),
}];
let result = apply_modifications_to_content(original_content, &mods, "test.yml").unwrap();
assert_eq!(
result,
"name: test_metric\ntype: counter\ndescription: A test metric\nadditional_field: true"
);
// Test appending content when original content doesn't end with newline
let original_content_no_newline = "name: test_metric\ntype: counter\ndescription: A test metric";
let result = apply_modifications_to_content(original_content_no_newline, &mods, "test.yml").unwrap();
assert_eq!(
result,
"name: test_metric\ntype: counter\ndescription: A test metric\nadditional_field: true"
);
}
// Note: We'll need integration tests with a real database for testing actual SQL validation
// Unit tests can only cover basic cases like empty SQL

32
api/libs/agents/src/tools/categories/file_tools/modify_dashboards.rs
View File

@ -272,7 +272,7 @@ impl ToolExecutor for ModifyDashboardFilesTool {
async fn get_modify_dashboards_description() -> String {
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
return "Modifies existing dashboard configuration files by replacing specified content with new content".to_string();
return "Modifies existing dashboard configuration files by replacing specified content with new content. If the dashboard fundamentally changes, the name should be updated to reflect the changes. However, if the core dashboard topic remains the same, the name should stay unchanged.".to_string();
}
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
@ -280,14 +280,14 @@ async fn get_modify_dashboards_description() -> String {
Ok(message) => message,
Err(e) => {
eprintln!("Failed to get prompt system message: {}", e);
"Modifies existing dashboard configuration files by replacing specified content with new content".to_string()
"Modifies existing dashboard configuration files by replacing specified content with new content. If the dashboard fundamentally changes, the name should be updated to reflect the changes. However, if the core dashboard topic remains the same, the name should stay unchanged.".to_string()
}
}
}
async fn get_modify_dashboards_yml_description() -> String {
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
return DASHBOARD_YML_SCHEMA.to_string();
return format!("{}. When making significant changes to a dashboard, update the name to reflect these changes. If the dashboard's core topic remains the same, keep the original name.", DASHBOARD_YML_SCHEMA);
}
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
@ -295,7 +295,7 @@ async fn get_modify_dashboards_yml_description() -> String {
Ok(message) => message,
Err(e) => {
eprintln!("Failed to get prompt system message: {}", e);
DASHBOARD_YML_SCHEMA.to_string()
format!("{}. When making significant changes to a dashboard, update the name to reflect these changes. If the dashboard's core topic remains the same, keep the original name.", DASHBOARD_YML_SCHEMA)
}
}
}
@ -347,7 +347,7 @@ async fn get_modify_dashboards_modifications_description() -> String {
async fn get_modify_dashboards_new_content_description() -> String {
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
return "The new content to replace the existing content with".to_string();
return "The new content to replace the existing content with. If fundamentally changing the dashboard's purpose or focus, update the name property accordingly. If the core topic remains the same, maintain the original name.".to_string();
}
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
@ -355,14 +355,14 @@ async fn get_modify_dashboards_new_content_description() -> String {
Ok(message) => message,
Err(e) => {
eprintln!("Failed to get prompt system message: {}", e);
"The new content to replace the existing content with".to_string()
"The new content to replace the existing content with. If fundamentally changing the dashboard's purpose or focus, update the name property accordingly. If the core topic remains the same, maintain the original name.".to_string()
}
}
}
async fn get_modify_dashboards_content_to_replace_description() -> String {
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
return "The exact content in the file that should be replaced. Must match exactly."
return "The exact content in the file that should be replaced. Must match exactly and be specific enough to only match once. Use an empty string to append the new content to the end of the file."
.to_string();
}
@ -371,7 +371,7 @@ async fn get_modify_dashboards_content_to_replace_description() -> String {
Ok(message) => message,
Err(e) => {
eprintln!("Failed to get prompt system message: {}", e);
"The exact content in the file that should be replaced. Must match exactly.".to_string()
"The exact content in the file that should be replaced. Must match exactly and be specific enough to only match once. Use an empty string to append the new content to the end of the file.".to_string()
}
}
}
@ -487,4 +487,20 @@ mod tests {
let result = serde_json::from_str::<ModifyFilesParams>(&missing_args);
assert!(result.is_err());
}
#[test]
fn test_apply_modifications_append() {
let original_content = "name: test_dashboard\ntype: dashboard\ndescription: A test dashboard";
// Test appending content with empty content_to_replace
let mods = vec![Modification {
content_to_replace: "".to_string(),
new_content: "\nvisibility: public".to_string(),
}];
let result = apply_modifications_to_content(original_content, &mods, "test.yml").unwrap();
assert_eq!(
result,
"name: test_dashboard\ntype: dashboard\ndescription: A test dashboard\nvisibility: public"
);
}
}

16
api/libs/agents/src/tools/categories/file_tools/modify_metrics.rs
View File

@ -357,7 +357,7 @@ impl ToolExecutor for ModifyMetricFilesTool {
async fn get_modify_metrics_description() -> String {
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
return "Modifies existing metric configuration files by replacing specified content with new content".to_string();
return "Modifies existing metric configuration files by replacing specified content with new content. When modifying a metric's SQL, make sure to also update the name and/or time frame to reflect the changes.".to_string();
}
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
@ -365,14 +365,14 @@ async fn get_modify_metrics_description() -> String {
Ok(message) => message,
Err(e) => {
eprintln!("Failed to get prompt system message: {}", e);
"Modifies existing metric configuration files by replacing specified content with new content".to_string()
"Modifies existing metric configuration files by replacing specified content with new content. When modifying a metric's SQL, make sure to also update the name and/or time frame to reflect the changes.".to_string()
}
}
}
async fn get_modify_metrics_yml_description() -> String {
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
return METRIC_YML_SCHEMA.to_string();
return format!("{}. Remember to update the metric name and/or time frame properties when modifying SQL to ensure they reflect the changes.", METRIC_YML_SCHEMA);
}
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
@ -380,7 +380,7 @@ async fn get_modify_metrics_yml_description() -> String {
Ok(message) => message,
Err(e) => {
eprintln!("Failed to get prompt system message: {}", e);
METRIC_YML_SCHEMA.to_string()
format!("{}. Remember to update the metric name and/or time frame properties when modifying SQL to ensure they reflect the changes.", METRIC_YML_SCHEMA)
}
}
}
@ -417,7 +417,7 @@ async fn get_modify_metrics_modifications_description() -> String {
async fn get_modify_metrics_new_content_description() -> String {
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
return "The new content to replace the existing content with".to_string();
return "The new content to replace the existing content with. If modifying SQL, ensure name and time frame properties are also updated to reflect the changes.".to_string();
}
let client = BraintrustClient::new(None, "96af8b2b-cf3c-494f-9092-44eb3d5b96ff").unwrap();
@ -425,14 +425,14 @@ async fn get_modify_metrics_new_content_description() -> String {
Ok(message) => message,
Err(e) => {
eprintln!("Failed to get prompt system message: {}", e);
"The new content to replace the existing content with".to_string()
"The new content to replace the existing content with. If modifying SQL, ensure name and time frame properties are also updated to reflect the changes.".to_string()
}
}
}
async fn get_modify_metrics_content_to_replace_description() -> String {
if env::var("USE_BRAINTRUST_PROMPTS").is_err() {
return "The exact content in the file that should be replaced. Must match exactly."
return "The exact content in the file that should be replaced. Must match exactly and be specific enough to only match once. Use an empty string to append the new content to the end of the file."
.to_string();
}
@ -441,7 +441,7 @@ async fn get_modify_metrics_content_to_replace_description() -> String {
Ok(message) => message,
Err(e) => {
eprintln!("Failed to get prompt system message: {}", e);
"The exact content in the file that should be replaced. Must match exactly.".to_string()
"The exact content in the file that should be replaced. Must match exactly and be specific enough to only match once. Use an empty string to append the new content to the end of the file.".to_string()
}
}
}

135
api/libs/database/src/helpers/chats.rs
View File

@ -1,14 +1,14 @@
use crate::enums::{AssetPermissionRole, AssetType};
use anyhow::Result;
use diesel::JoinOnDsl;
use diesel::{BoolExpressionMethods, ExpressionMethods, QueryDsl, Queryable};
use diesel::{JoinOnDsl, NullableExpressionMethods};
use diesel_async::RunQueryDsl;
use tokio::try_join;
use uuid::Uuid;
use crate::models::{AssetPermission, Chat};
use crate::pool::get_pg_pool;
use crate::schema::{asset_permissions, chats};
use crate::schema::{asset_permissions, chats, collections_to_assets};
/// Fetches a single chat by ID that hasn't been deleted
///
@ -72,6 +72,43 @@ async fn fetch_chat_permission(id: &Uuid, user_id: &Uuid) -> Result<Option<Asset
Ok(permission)
}
/// Helper function to fetch collection-based permissions for a chat
///
/// Checks if the user has access to the chat through collections
/// by finding collections that contain this chat and checking
/// if the user has permissions on those collections
async fn fetch_collection_permissions_for_chat(
id: &Uuid,
user_id: &Uuid,
) -> Result<Option<AssetPermissionRole>> {
let mut conn = get_pg_pool().get().await?;
// Find collections containing this chat
// then join with asset_permissions to find user's permissions on those collections
let permissions = asset_permissions::table
.inner_join(
collections_to_assets::table.on(asset_permissions::asset_id
.eq(collections_to_assets::collection_id)
.and(asset_permissions::asset_type.eq(AssetType::Collection))
.and(collections_to_assets::asset_id.eq(id))
.and(collections_to_assets::asset_type.eq(AssetType::Chat))
.and(collections_to_assets::deleted_at.is_null())),
)
.filter(asset_permissions::identity_id.eq(user_id))
.filter(asset_permissions::deleted_at.is_null())
.select(asset_permissions::role)
.load::<AssetPermissionRole>(&mut conn)
.await?;
// Return any collection-based permission
if permissions.is_empty() {
Ok(None)
} else {
// Just take the first one since any collection permission is sufficient
Ok(permissions.into_iter().next())
}
}
#[derive(Queryable)]
pub struct ChatWithPermission {
pub chat: Chat,
@ -82,8 +119,12 @@ pub async fn fetch_chat_with_permission(
id: &Uuid,
user_id: &Uuid,
) -> Result<Option<ChatWithPermission>> {
// Run both queries concurrently
let (chat, permission) = try_join!(fetch_chat_helper(id), fetch_chat_permission(id, user_id))?;
// Run all queries concurrently
let (chat, direct_permission, collection_permission) = try_join!(
fetch_chat_helper(id),
fetch_chat_permission(id, user_id),
fetch_collection_permissions_for_chat(id, user_id)
)?;
// If the chat doesn't exist, return None
let chat = match chat {
@ -91,7 +132,16 @@ pub async fn fetch_chat_with_permission(
None => return Ok(None),
};
Ok(Some(ChatWithPermission { chat, permission }))
// If collection permission exists, use it; otherwise use direct permission
let effective_permission = match collection_permission {
Some(collection) => Some(collection),
None => direct_permission,
};
Ok(Some(ChatWithPermission {
chat,
permission: effective_permission,
}))
}
/// Fetches multiple chats with their permissions in a single operation
@ -112,29 +162,74 @@ pub async fn fetch_chats_with_permissions(
let mut conn = get_pg_pool().get().await?;
// Use a LEFT JOIN to fetch chats and their permissions in a single query
let results = chats::table
.left_join(
asset_permissions::table.on(asset_permissions::asset_id
.eq(chats::id)
.and(asset_permissions::asset_type.eq(AssetType::Chat))
.and(asset_permissions::identity_id.eq(user_id))
.and(asset_permissions::deleted_at.is_null())),
)
// 1. Fetch all chats
let chats_data = chats::table
.filter(chats::id.eq_any(ids))
.filter(chats::deleted_at.is_null())
.select((chats::all_columns, asset_permissions::role.nullable()))
.load::<(Chat, Option<AssetPermissionRole>)>(&mut conn)
.load::<Chat>(&mut conn)
.await?;
if results.is_empty() {
if chats_data.is_empty() {
return Ok(Vec::new());
}
// Create ChatWithPermission objects
let result = results
// 2. Fetch direct permissions for these chats
let direct_permissions = asset_permissions::table
.filter(asset_permissions::asset_id.eq_any(ids))
.filter(asset_permissions::asset_type.eq(AssetType::Chat))
.filter(asset_permissions::identity_id.eq(user_id))
.filter(asset_permissions::deleted_at.is_null())
.select((asset_permissions::asset_id, asset_permissions::role))
.load::<(Uuid, AssetPermissionRole)>(&mut conn)
.await?;
// 3. Fetch collection-based permissions for these chats
let collection_permissions = asset_permissions::table
.inner_join(
collections_to_assets::table.on(asset_permissions::asset_id
.eq(collections_to_assets::collection_id)
.and(asset_permissions::asset_type.eq(AssetType::Collection))
.and(collections_to_assets::asset_id.eq_any(ids))
.and(collections_to_assets::asset_type.eq(AssetType::Chat))
.and(collections_to_assets::deleted_at.is_null())),
)
.filter(asset_permissions::identity_id.eq(user_id))
.filter(asset_permissions::deleted_at.is_null())
.select((collections_to_assets::asset_id, asset_permissions::role))
.load::<(Uuid, AssetPermissionRole)>(&mut conn)
.await?;
// Create maps for easier lookup
let mut direct_permission_map = std::collections::HashMap::new();
for (asset_id, role) in direct_permissions {
direct_permission_map.insert(asset_id, role);
}
// Create map for collection permissions (just take first one for each asset)
let mut collection_permission_map = std::collections::HashMap::new();
for (asset_id, role) in collection_permissions {
// Only insert if not already present (first one wins)
collection_permission_map.entry(asset_id).or_insert(role);
}
// Create ChatWithPermission objects with effective permissions
let result = chats_data
.into_iter()
.map(|(chat, permission)| ChatWithPermission { chat, permission })
.map(|chat| {
let direct_permission = direct_permission_map.get(&chat.id).cloned();
let collection_permission = collection_permission_map.get(&chat.id).cloned();
// Use collection permission if it exists, otherwise use direct permission
let effective_permission = match collection_permission {
Some(collection) => Some(collection),
None => direct_permission,
};
ChatWithPermission {
chat,
permission: effective_permission,
}
})
.collect();
Ok(result)

134
api/libs/database/src/helpers/dashboard_files.rs
View File

@ -8,7 +8,7 @@ use crate::enums::{AssetPermissionRole, AssetType};
use crate::models::{AssetPermission, DashboardFile};
use crate::pool::get_pg_pool;
use crate::schema::{asset_permissions, dashboard_files};
use crate::schema::{asset_permissions, dashboard_files, collections_to_assets};
/// Fetches a single dashboard file by ID that hasn't been deleted
///
@ -90,6 +90,44 @@ async fn fetch_dashboard_permission(id: &Uuid, user_id: &Uuid) -> Result<Option<
Ok(permission)
}
/// Helper function to fetch collection-based permissions for a dashboard file
///
/// Checks if the user has access to the dashboard file through collections
/// by finding collections that contain this dashboard file and checking
/// if the user has permissions on those collections
async fn fetch_collection_permissions_for_dashboard(
id: &Uuid,
user_id: &Uuid
) -> Result<Option<AssetPermissionRole>> {
let mut conn = get_pg_pool().get().await?;
// Find collections containing this dashboard file
// then join with asset_permissions to find user's permissions on those collections
let permissions = asset_permissions::table
.inner_join(
collections_to_assets::table.on(
asset_permissions::asset_id.eq(collections_to_assets::collection_id)
.and(asset_permissions::asset_type.eq(AssetType::Collection))
.and(collections_to_assets::asset_id.eq(id))
.and(collections_to_assets::asset_type.eq(AssetType::DashboardFile))
.and(collections_to_assets::deleted_at.is_null())
)
)
.filter(asset_permissions::identity_id.eq(user_id))
.filter(asset_permissions::deleted_at.is_null())
.select(asset_permissions::role)
.load::<AssetPermissionRole>(&mut conn)
.await?;
// Return any collection-based permission
if permissions.is_empty() {
Ok(None)
} else {
// Just take the first one since any collection permission is sufficient
Ok(permissions.into_iter().next())
}
}
#[derive(Queryable)]
pub struct DashboardFileWithPermission {
pub dashboard_file: DashboardFile,
@ -100,10 +138,11 @@ pub async fn fetch_dashboard_file_with_permission(
id: &Uuid,
user_id: &Uuid,
) -> Result<Option<DashboardFileWithPermission>> {
// Run both queries concurrently
let (dashboard_file, permission) = try_join!(
// Run all queries concurrently
let (dashboard_file, direct_permission, collection_permission) = try_join!(
fetch_dashboard(id),
fetch_dashboard_permission(id, user_id)
fetch_dashboard_permission(id, user_id),
fetch_collection_permissions_for_dashboard(id, user_id)
)?;
// If the dashboard file doesn't exist, return None
@ -112,9 +151,15 @@ pub async fn fetch_dashboard_file_with_permission(
None => return Ok(None),
};
// If collection permission exists, use it; otherwise use direct permission
let effective_permission = match collection_permission {
Some(collection) => Some(collection),
None => direct_permission
};
Ok(Some(DashboardFileWithPermission {
dashboard_file,
permission,
permission: effective_permission,
}))
}
@ -136,27 +181,19 @@ pub async fn fetch_dashboard_files_with_permissions(
let mut conn = get_pg_pool().get().await?;
// Use a LEFT JOIN to fetch dashboard files and their permissions in a single query
let mut results = dashboard_files::table
.left_join(
asset_permissions::table.on(asset_permissions::asset_id
.eq(dashboard_files::id)
.and(asset_permissions::asset_type.eq(AssetType::Dashboard))
.and(asset_permissions::identity_id.eq(user_id))
.and(asset_permissions::deleted_at.is_null()))
)
// 1. Fetch all dashboard files
let mut dashboard_files = dashboard_files::table
.filter(dashboard_files::id.eq_any(ids))
.filter(dashboard_files::deleted_at.is_null())
.select((dashboard_files::all_columns, asset_permissions::role.nullable()))
.load::<(DashboardFile, Option<AssetPermissionRole>)>(&mut conn)
.load::<DashboardFile>(&mut conn)
.await?;
if results.is_empty() {
if dashboard_files.is_empty() {
return Ok(Vec::new());
}
// Ensure all rows have IDs (for backwards compatibility)
for (dashboard_file, _) in &mut results {
for dashboard_file in &mut dashboard_files {
for (index, row) in dashboard_file.content.rows.iter_mut().enumerate() {
if row.id == 0 {
row.id = (index as u32) + 1;
@ -164,12 +201,63 @@ pub async fn fetch_dashboard_files_with_permissions(
}
}
// Create DashboardFileWithPermission objects
let result = results
// 2. Fetch direct permissions for these dashboard files
let direct_permissions = asset_permissions::table
.filter(asset_permissions::asset_id.eq_any(ids))
.filter(asset_permissions::asset_type.eq(AssetType::DashboardFile))
.filter(asset_permissions::identity_id.eq(user_id))
.filter(asset_permissions::deleted_at.is_null())
.select((asset_permissions::asset_id, asset_permissions::role))
.load::<(Uuid, AssetPermissionRole)>(&mut conn)
.await?;
// 3. Fetch collection-based permissions for these dashboard files
let collection_permissions = asset_permissions::table
.inner_join(
collections_to_assets::table.on(
asset_permissions::asset_id.eq(collections_to_assets::collection_id)
.and(asset_permissions::asset_type.eq(AssetType::Collection))
.and(collections_to_assets::asset_id.eq_any(ids))
.and(collections_to_assets::asset_type.eq(AssetType::DashboardFile))
.and(collections_to_assets::deleted_at.is_null())
)
)
.filter(asset_permissions::identity_id.eq(user_id))
.filter(asset_permissions::deleted_at.is_null())
.select((collections_to_assets::asset_id, asset_permissions::role))
.load::<(Uuid, AssetPermissionRole)>(&mut conn)
.await?;
// Create maps for easier lookup
let mut direct_permission_map = std::collections::HashMap::new();
for (asset_id, role) in direct_permissions {
direct_permission_map.insert(asset_id, role);
}
// Create map for collection permissions (just take first one for each asset)
let mut collection_permission_map = std::collections::HashMap::new();
for (asset_id, role) in collection_permissions {
// Only insert if not already present (first one wins)
collection_permission_map.entry(asset_id).or_insert(role);
}
// Create DashboardFileWithPermission objects with effective permissions
let result = dashboard_files
.into_iter()
.map(|(dashboard_file, permission)| DashboardFileWithPermission {
dashboard_file,
permission,
.map(|dashboard_file| {
let direct_permission = direct_permission_map.get(&dashboard_file.id).cloned();
let collection_permission = collection_permission_map.get(&dashboard_file.id).cloned();
// Use collection permission if it exists, otherwise use direct permission
let effective_permission = match collection_permission {
Some(collection) => Some(collection),
None => direct_permission
};
DashboardFileWithPermission {
dashboard_file,
permission: effective_permission,
}
})
.collect();

132
api/libs/database/src/helpers/metric_files.rs
View File

@ -8,7 +8,7 @@ use crate::enums::{AssetPermissionRole, AssetType};
use crate::models::{AssetPermission, MetricFile};
use crate::pool::get_pg_pool;
use crate::schema::{asset_permissions, metric_files};
use crate::schema::{asset_permissions, metric_files, collections_to_assets};
/// Fetches a single metric file by ID that hasn't been deleted
///
@ -72,6 +72,44 @@ async fn fetch_permission(id: &Uuid, user_id: &Uuid) -> Result<Option<AssetPermi
Ok(permission)
}
/// Helper function to fetch collection-based permissions for a metric file
///
/// Checks if the user has access to the metric file through collections
/// by finding collections that contain this metric file and checking
/// if the user has permissions on those collections
async fn fetch_collection_permissions_for_metric(
id: &Uuid,
user_id: &Uuid
) -> Result<Option<AssetPermissionRole>> {
let mut conn = get_pg_pool().get().await?;
// Find collections containing this metric file
// then join with asset_permissions to find user's permissions on those collections
let permissions = asset_permissions::table
.inner_join(
collections_to_assets::table.on(
asset_permissions::asset_id.eq(collections_to_assets::collection_id)
.and(asset_permissions::asset_type.eq(AssetType::Collection))
.and(collections_to_assets::asset_id.eq(id))
.and(collections_to_assets::asset_type.eq(AssetType::MetricFile))
.and(collections_to_assets::deleted_at.is_null())
)
)
.filter(asset_permissions::identity_id.eq(user_id))
.filter(asset_permissions::deleted_at.is_null())
.select(asset_permissions::role)
.load::<AssetPermissionRole>(&mut conn)
.await?;
// Return any collection-based permission (no need to determine highest)
if permissions.is_empty() {
Ok(None)
} else {
// Just take the first one since any collection permission is sufficient
Ok(permissions.into_iter().next())
}
}
#[derive(Queryable)]
pub struct MetricFileWithPermission {
pub metric_file: MetricFile,
@ -82,10 +120,11 @@ pub async fn fetch_metric_file_with_permissions(
id: &Uuid,
user_id: &Uuid,
) -> Result<Option<MetricFileWithPermission>> {
// Run both queries concurrently
let (metric_file, permission) = try_join!(
// Run all queries concurrently
let (metric_file, direct_permission, collection_permission) = try_join!(
fetch_metric(id),
fetch_permission(id, user_id)
fetch_permission(id, user_id),
fetch_collection_permissions_for_metric(id, user_id)
)?;
// If the metric file doesn't exist, return None
@ -94,9 +133,15 @@ pub async fn fetch_metric_file_with_permissions(
None => return Ok(None),
};
// If collection permission exists, use it; otherwise use direct permission
let effective_permission = match collection_permission {
Some(collection) => Some(collection),
None => direct_permission
};
Ok(Some(MetricFileWithPermission {
metric_file,
permission,
permission: effective_permission,
}))
}
@ -118,31 +163,74 @@ pub async fn fetch_metric_files_with_permissions(
let mut conn = get_pg_pool().get().await?;
// Use a LEFT JOIN to fetch metric files and their permissions in a single query
let results = metric_files::table
.left_join(
asset_permissions::table.on(asset_permissions::asset_id
.eq(metric_files::id)
.and(asset_permissions::asset_type.eq(AssetType::MetricFile))
.and(asset_permissions::identity_id.eq(user_id))
.and(asset_permissions::deleted_at.is_null()))
)
// 1. Fetch all metric files
let metric_files = metric_files::table
.filter(metric_files::id.eq_any(ids))
.filter(metric_files::deleted_at.is_null())
.select((metric_files::all_columns, asset_permissions::role.nullable()))
.load::<(MetricFile, Option<AssetPermissionRole>)>(&mut conn)
.load::<MetricFile>(&mut conn)
.await?;
if results.is_empty() {
if metric_files.is_empty() {
return Ok(Vec::new());
}
// Create MetricFileWithPermission objects
let result = results
// 2. Fetch direct permissions for these metric files
let direct_permissions = asset_permissions::table
.filter(asset_permissions::asset_id.eq_any(ids))
.filter(asset_permissions::asset_type.eq(AssetType::MetricFile))
.filter(asset_permissions::identity_id.eq(user_id))
.filter(asset_permissions::deleted_at.is_null())
.select((asset_permissions::asset_id, asset_permissions::role))
.load::<(Uuid, AssetPermissionRole)>(&mut conn)
.await?;
// 3. Fetch collection-based permissions for these metric files
let collection_permissions = asset_permissions::table
.inner_join(
collections_to_assets::table.on(
asset_permissions::asset_id.eq(collections_to_assets::collection_id)
.and(asset_permissions::asset_type.eq(AssetType::Collection))
.and(collections_to_assets::asset_id.eq_any(ids))
.and(collections_to_assets::asset_type.eq(AssetType::MetricFile))
.and(collections_to_assets::deleted_at.is_null())
)
)
.filter(asset_permissions::identity_id.eq(user_id))
.filter(asset_permissions::deleted_at.is_null())
.select((collections_to_assets::asset_id, asset_permissions::role))
.load::<(Uuid, AssetPermissionRole)>(&mut conn)
.await?;
// Create maps for easier lookup
let mut direct_permission_map = std::collections::HashMap::new();
for (asset_id, role) in direct_permissions {
direct_permission_map.insert(asset_id, role);
}
// Create map for collection permissions (just take first one for each asset)
let mut collection_permission_map = std::collections::HashMap::new();
for (asset_id, role) in collection_permissions {
// Only insert if not already present (first one wins)
collection_permission_map.entry(asset_id).or_insert(role);
}
// Create MetricFileWithPermission objects with effective permissions
let result = metric_files
.into_iter()
.map(|(metric_file, permission)| MetricFileWithPermission {
metric_file,
permission,
.map(|metric_file| {
let direct_permission = direct_permission_map.get(&metric_file.id).cloned();
let collection_permission = collection_permission_map.get(&metric_file.id).cloned();
// Use collection permission if it exists, otherwise use direct permission
let effective_permission = match collection_permission {
Some(collection) => Some(collection),
None => direct_permission
};
MetricFileWithPermission {
metric_file,
permission: effective_permission,
}
})
.collect();