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

add to collections

This commit is contained in:
dal 2025-03-20 09:57:01 -06:00
parent b8e4a95316
commit 5b484510c3
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
4 changed files with 70 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -558,6 +558,7 @@ pub async fn process_metric_file(
tool_call_id: String,
file_name: String,
yml_content: String,
user_id: &Uuid,
) -> Result<
(
MetricFile,
@ -599,7 +600,7 @@ pub async fn process_metric_file(
name: file_name.clone(),
file_name: file_name.clone(),
content: metric_yml.clone(),
created_by: Uuid::new_v4(),
created_by: user_id.clone(),
verification: Verification::NotRequested,
evaluation_obj: None,
evaluation_summary: None,

5
api/libs/agents/src/tools/categories/file_tools/create_dashboards.rs
View File

@ -71,6 +71,7 @@ impl FileModificationTool for CreateDashboardFilesTool {}
async fn process_dashboard_file(
tool_call_id: String,
file: DashboardFileParams,
user_id: &Uuid,
) -> Result<(DashboardFile, DashboardYml), String> {
debug!("Processing dashboard file creation: {}", file.name);
@ -106,7 +107,7 @@ async fn process_dashboard_file(
content: dashboard_yml.clone(),
filter: None,
organization_id: Uuid::new_v4(),
created_by: Uuid::new_v4(),
created_by: user_id.clone(),
created_at: Utc::now(),
updated_at: Utc::now(),
deleted_at: None,
@ -152,7 +153,7 @@ impl ToolExecutor for CreateDashboardFilesTool {
// First pass - validate and prepare all records
for file in files {
match process_dashboard_file(tool_call_id.clone(), file.clone()).await {
match process_dashboard_file(tool_call_id.clone(), file.clone(), &self.agent.get_user_id()).await {
Ok((dashboard_file, dashboard_yml)) => {
dashboard_records.push(dashboard_file);
dashboard_ymls.push(dashboard_yml);

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

@ -5,10 +5,10 @@ use async_trait::async_trait;
use braintrust::{get_prompt_system_message, BraintrustClient};
use chrono::Utc;
use database::{
pool::get_pg_pool,
schema::{metric_files, asset_permissions},
enums::{AssetPermissionRole, AssetType, IdentityType},
models::AssetPermission,
enums::{AssetType, IdentityType, AssetPermissionRole},
pool::get_pg_pool,
schema::{asset_permissions, metric_files},
};
use diesel::insert_into;
use diesel_async::RunQueryDsl;
@ -23,13 +23,16 @@ use crate::{
tools::{
file_tools::{
common::{process_metric_file, METRIC_YML_SCHEMA},
file_types::{file::FileWithId},
file_types::file::FileWithId,
},
ToolExecutor,
},
};
use super::{common::{validate_sql, generate_deterministic_uuid}, FileModificationTool};
use super::{
common::{generate_deterministic_uuid, validate_sql},
FileModificationTool,
};
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct MetricFileParams {
@ -99,7 +102,14 @@ impl ToolExecutor for CreateMetricFilesTool {
let mut results_vec = vec![];
// First pass - validate and prepare all records
for file in files {
match process_metric_file(tool_call_id.clone(), file.name.clone(), file.yml_content.clone()).await {
match process_metric_file(
tool_call_id.clone(),
file.name.clone(),
file.yml_content.clone(),
&self.agent.get_user_id(),
)
.await
{
Ok((metric_file, metric_yml, message, results)) => {
metric_records.push(metric_file);
metric_ymls.push(metric_yml);
@ -127,7 +137,7 @@ impl ToolExecutor for CreateMetricFilesTool {
Ok(_) => {
// Get the user ID from the agent state
let user_id = self.agent.get_user_id();
// Create asset permissions for each metric file
let now = Utc::now();
let asset_permissions: Vec<AssetPermission> = metric_records
@ -145,7 +155,7 @@ impl ToolExecutor for CreateMetricFilesTool {
updated_by: user_id,
})
.collect();
// Insert asset permissions
match insert_into(asset_permissions::table)
.values(&asset_permissions)
@ -153,15 +163,18 @@ impl ToolExecutor for CreateMetricFilesTool {
.await
{
Ok(_) => {
tracing::debug!("Successfully inserted asset permissions for {} metric files", asset_permissions.len());
},
tracing::debug!(
"Successfully inserted asset permissions for {} metric files",
asset_permissions.len()
);
}
Err(e) => {
tracing::error!("Error inserting asset permissions: {}", e);
// Continue with the process even if permissions failed
// We'll still return the created files
}
}
for (i, yml) in metric_ymls.into_iter().enumerate() {
created_files.push(FileWithId {
id: metric_records[i].id,

72
api/libs/handlers/src/collections/get_collection_handler.rs
View File

@ -1,12 +1,21 @@
use anyhow::{anyhow, Result};
use chrono::{DateTime, Utc};
use database::{collections::fetch_collection, enums::{AssetPermissionRole, AssetType, IdentityType}, pool::get_pg_pool, schema::{asset_permissions, collections, collections_to_assets, dashboard_files, metric_files, users}};
use diesel::{ExpressionMethods, JoinOnDsl, QueryDsl, Queryable};
use database::{
collections::fetch_collection,
enums::{AssetPermissionRole, AssetType, IdentityType},
pool::get_pg_pool,
schema::{
asset_permissions, collections, collections_to_assets, dashboard_files, metric_files, users,
},
};
use diesel::{ExpressionMethods, JoinOnDsl, NullableExpressionMethods, QueryDsl, Queryable};
use diesel_async::RunQueryDsl;
use tracing;
use uuid::Uuid;
use crate::collections::types::{AssetUser, BusterShareIndividual, CollectionAsset, CollectionState, GetCollectionRequest};
use crate::collections::types::{
AssetUser, BusterShareIndividual, CollectionAsset, CollectionState, GetCollectionRequest,
};
#[derive(Queryable)]
struct AssetPermissionInfo {
@ -22,7 +31,7 @@ struct AssetQueryResult {
id: Uuid,
name: String,
user_name: Option<String>,
email: String,
email: Option<String>,
created_at: DateTime<Utc>,
updated_at: DateTime<Utc>,
asset_type: AssetType,
@ -41,18 +50,16 @@ struct AssetQueryResult {
fn format_assets(assets: Vec<AssetQueryResult>) -> Vec<CollectionAsset> {
assets
.into_iter()
.map(|asset| {
CollectionAsset {
id: asset.id,
name: asset.name,
created_by: AssetUser {
name: asset.user_name,
email: asset.email,
},
created_at: asset.created_at,
updated_at: asset.updated_at,
asset_type: asset.asset_type,
}
.map(|asset| CollectionAsset {
id: asset.id,
name: asset.name,
created_by: AssetUser {
name: asset.user_name,
email: asset.email.unwrap_or("chad@buster.so".to_string()),
},
created_at: asset.created_at,
updated_at: asset.updated_at,
asset_type: asset.asset_type,
})
.collect()
}
@ -66,12 +73,12 @@ pub async fn get_collection_handler(
Some(collection) => collection,
None => return Err(anyhow!("Collection not found")),
};
let mut conn = match get_pg_pool().get().await {
Ok(conn) => conn,
Err(e) => return Err(anyhow!("Failed to get database connection: {}", e)),
};
// Query individual permissions for this collection
let individual_permissions_query = asset_permissions::table
.inner_join(users::table.on(users::id.eq(asset_permissions::identity_id)))
@ -87,11 +94,12 @@ pub async fn get_collection_handler(
))
.load::<AssetPermissionInfo>(&mut conn)
.await;
// For collections, we'll default public fields to false/none
// since the schema doesn't have these fields yet
let public_info: Result<(bool, Option<Uuid>, Option<DateTime<Utc>>), anyhow::Error> = Ok((false, None, None));
let public_info: Result<(bool, Option<Uuid>, Option<DateTime<Utc>>), anyhow::Error> =
Ok((false, None, None));
// Convert AssetPermissionInfo to BusterShareIndividual
let individual_permissions = match individual_permissions_query {
Ok(permissions) => {
@ -112,7 +120,7 @@ pub async fn get_collection_handler(
}
Err(_) => None,
};
// Get public access info
let (publicly_accessible, public_enabled_by, public_expiry_date) = match public_info {
Ok((accessible, enabled_by_id, expiry)) => {
@ -127,7 +135,7 @@ pub async fn get_collection_handler(
} else {
None
};
(accessible, enabled_by_email, expiry)
}
Err(_) => (false, None, None),
@ -136,7 +144,7 @@ pub async fn get_collection_handler(
// Query for metric assets in the collection
let metric_assets_result = collections_to_assets::table
.inner_join(metric_files::table.on(metric_files::id.eq(collections_to_assets::asset_id)))
.inner_join(users::table.on(users::id.eq(metric_files::created_by)))
.left_join(users::table.on(users::id.eq(metric_files::created_by)))
.filter(collections_to_assets::collection_id.eq(req.id))
.filter(collections_to_assets::asset_type.eq(AssetType::MetricFile))
.filter(collections_to_assets::deleted_at.is_null())
@ -144,8 +152,8 @@ pub async fn get_collection_handler(
.select((
metric_files::id,
metric_files::name,
users::name,
users::email,
users::name.nullable(),
users::email.nullable(),
metric_files::created_at,
metric_files::updated_at,
collections_to_assets::asset_type,
@ -155,8 +163,10 @@ pub async fn get_collection_handler(
// Query for dashboard assets in the collection
let dashboard_assets_result = collections_to_assets::table
.inner_join(dashboard_files::table.on(dashboard_files::id.eq(collections_to_assets::asset_id)))
.inner_join(users::table.on(users::id.eq(dashboard_files::created_by)))
.inner_join(
dashboard_files::table.on(dashboard_files::id.eq(collections_to_assets::asset_id)),
)
.left_join(users::table.on(users::id.eq(dashboard_files::created_by)))
.filter(collections_to_assets::collection_id.eq(req.id))
.filter(collections_to_assets::asset_type.eq(AssetType::DashboardFile))
.filter(collections_to_assets::deleted_at.is_null())
@ -164,8 +174,8 @@ pub async fn get_collection_handler(
.select((
dashboard_files::id,
dashboard_files::name,
users::name,
users::email,
users::name.nullable(),
users::email.nullable(),
dashboard_files::created_at,
dashboard_files::updated_at,
collections_to_assets::asset_type,
@ -194,7 +204,7 @@ pub async fn get_collection_handler(
// Combine and format the assets
let mut combined_assets = metric_assets;
combined_assets.extend(dashboard_assets);
// Only include assets in the response if we found some
let formatted_assets = if combined_assets.is_empty() {
None