mirror of https://github.com/buster-so/buster.git
add to collections
This commit is contained in:
parent
b8e4a95316
commit
5b484510c3
|
@ -558,6 +558,7 @@ pub async fn process_metric_file(
|
||||||
tool_call_id: String,
|
tool_call_id: String,
|
||||||
file_name: String,
|
file_name: String,
|
||||||
yml_content: String,
|
yml_content: String,
|
||||||
|
user_id: &Uuid,
|
||||||
) -> Result<
|
) -> Result<
|
||||||
(
|
(
|
||||||
MetricFile,
|
MetricFile,
|
||||||
|
@ -599,7 +600,7 @@ pub async fn process_metric_file(
|
||||||
name: file_name.clone(),
|
name: file_name.clone(),
|
||||||
file_name: file_name.clone(),
|
file_name: file_name.clone(),
|
||||||
content: metric_yml.clone(),
|
content: metric_yml.clone(),
|
||||||
created_by: Uuid::new_v4(),
|
created_by: user_id.clone(),
|
||||||
verification: Verification::NotRequested,
|
verification: Verification::NotRequested,
|
||||||
evaluation_obj: None,
|
evaluation_obj: None,
|
||||||
evaluation_summary: None,
|
evaluation_summary: None,
|
||||||
|
|
|
@ -71,6 +71,7 @@ impl FileModificationTool for CreateDashboardFilesTool {}
|
||||||
async fn process_dashboard_file(
|
async fn process_dashboard_file(
|
||||||
tool_call_id: String,
|
tool_call_id: String,
|
||||||
file: DashboardFileParams,
|
file: DashboardFileParams,
|
||||||
|
user_id: &Uuid,
|
||||||
) -> Result<(DashboardFile, DashboardYml), String> {
|
) -> Result<(DashboardFile, DashboardYml), String> {
|
||||||
debug!("Processing dashboard file creation: {}", file.name);
|
debug!("Processing dashboard file creation: {}", file.name);
|
||||||
|
|
||||||
|
@ -106,7 +107,7 @@ async fn process_dashboard_file(
|
||||||
content: dashboard_yml.clone(),
|
content: dashboard_yml.clone(),
|
||||||
filter: None,
|
filter: None,
|
||||||
organization_id: Uuid::new_v4(),
|
organization_id: Uuid::new_v4(),
|
||||||
created_by: Uuid::new_v4(),
|
created_by: user_id.clone(),
|
||||||
created_at: Utc::now(),
|
created_at: Utc::now(),
|
||||||
updated_at: Utc::now(),
|
updated_at: Utc::now(),
|
||||||
deleted_at: None,
|
deleted_at: None,
|
||||||
|
@ -152,7 +153,7 @@ impl ToolExecutor for CreateDashboardFilesTool {
|
||||||
|
|
||||||
// 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()).await {
|
match process_dashboard_file(tool_call_id.clone(), file.clone(), &self.agent.get_user_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);
|
||||||
|
|
|
@ -5,10 +5,10 @@ use async_trait::async_trait;
|
||||||
use braintrust::{get_prompt_system_message, BraintrustClient};
|
use braintrust::{get_prompt_system_message, BraintrustClient};
|
||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
use database::{
|
use database::{
|
||||||
pool::get_pg_pool,
|
enums::{AssetPermissionRole, AssetType, IdentityType},
|
||||||
schema::{metric_files, asset_permissions},
|
|
||||||
models::AssetPermission,
|
models::AssetPermission,
|
||||||
enums::{AssetType, IdentityType, AssetPermissionRole},
|
pool::get_pg_pool,
|
||||||
|
schema::{asset_permissions, metric_files},
|
||||||
};
|
};
|
||||||
use diesel::insert_into;
|
use diesel::insert_into;
|
||||||
use diesel_async::RunQueryDsl;
|
use diesel_async::RunQueryDsl;
|
||||||
|
@ -23,13 +23,16 @@ use crate::{
|
||||||
tools::{
|
tools::{
|
||||||
file_tools::{
|
file_tools::{
|
||||||
common::{process_metric_file, METRIC_YML_SCHEMA},
|
common::{process_metric_file, METRIC_YML_SCHEMA},
|
||||||
file_types::{file::FileWithId},
|
file_types::file::FileWithId,
|
||||||
},
|
},
|
||||||
ToolExecutor,
|
ToolExecutor,
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
use super::{common::{validate_sql, generate_deterministic_uuid}, FileModificationTool};
|
use super::{
|
||||||
|
common::{generate_deterministic_uuid, validate_sql},
|
||||||
|
FileModificationTool,
|
||||||
|
};
|
||||||
|
|
||||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||||
pub struct MetricFileParams {
|
pub struct MetricFileParams {
|
||||||
|
@ -99,7 +102,14 @@ impl ToolExecutor for CreateMetricFilesTool {
|
||||||
let mut results_vec = vec![];
|
let mut results_vec = 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_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)) => {
|
Ok((metric_file, metric_yml, message, results)) => {
|
||||||
metric_records.push(metric_file);
|
metric_records.push(metric_file);
|
||||||
metric_ymls.push(metric_yml);
|
metric_ymls.push(metric_yml);
|
||||||
|
@ -153,8 +163,11 @@ impl ToolExecutor for CreateMetricFilesTool {
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(_) => {
|
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) => {
|
Err(e) => {
|
||||||
tracing::error!("Error inserting asset permissions: {}", e);
|
tracing::error!("Error inserting asset permissions: {}", e);
|
||||||
// Continue with the process even if permissions failed
|
// Continue with the process even if permissions failed
|
||||||
|
|
|
@ -1,12 +1,21 @@
|
||||||
use anyhow::{anyhow, Result};
|
use anyhow::{anyhow, Result};
|
||||||
use chrono::{DateTime, Utc};
|
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 database::{
|
||||||
use diesel::{ExpressionMethods, JoinOnDsl, QueryDsl, Queryable};
|
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 diesel_async::RunQueryDsl;
|
||||||
use tracing;
|
use tracing;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::collections::types::{AssetUser, BusterShareIndividual, CollectionAsset, CollectionState, GetCollectionRequest};
|
use crate::collections::types::{
|
||||||
|
AssetUser, BusterShareIndividual, CollectionAsset, CollectionState, GetCollectionRequest,
|
||||||
|
};
|
||||||
|
|
||||||
#[derive(Queryable)]
|
#[derive(Queryable)]
|
||||||
struct AssetPermissionInfo {
|
struct AssetPermissionInfo {
|
||||||
|
@ -22,7 +31,7 @@ struct AssetQueryResult {
|
||||||
id: Uuid,
|
id: Uuid,
|
||||||
name: String,
|
name: String,
|
||||||
user_name: Option<String>,
|
user_name: Option<String>,
|
||||||
email: String,
|
email: Option<String>,
|
||||||
created_at: DateTime<Utc>,
|
created_at: DateTime<Utc>,
|
||||||
updated_at: DateTime<Utc>,
|
updated_at: DateTime<Utc>,
|
||||||
asset_type: AssetType,
|
asset_type: AssetType,
|
||||||
|
@ -41,18 +50,16 @@ struct AssetQueryResult {
|
||||||
fn format_assets(assets: Vec<AssetQueryResult>) -> Vec<CollectionAsset> {
|
fn format_assets(assets: Vec<AssetQueryResult>) -> Vec<CollectionAsset> {
|
||||||
assets
|
assets
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|asset| {
|
.map(|asset| CollectionAsset {
|
||||||
CollectionAsset {
|
id: asset.id,
|
||||||
id: asset.id,
|
name: asset.name,
|
||||||
name: asset.name,
|
created_by: AssetUser {
|
||||||
created_by: AssetUser {
|
name: asset.user_name,
|
||||||
name: asset.user_name,
|
email: asset.email.unwrap_or("chad@buster.so".to_string()),
|
||||||
email: asset.email,
|
},
|
||||||
},
|
created_at: asset.created_at,
|
||||||
created_at: asset.created_at,
|
updated_at: asset.updated_at,
|
||||||
updated_at: asset.updated_at,
|
asset_type: asset.asset_type,
|
||||||
asset_type: asset.asset_type,
|
|
||||||
}
|
|
||||||
})
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
@ -90,7 +97,8 @@ pub async fn get_collection_handler(
|
||||||
|
|
||||||
// For collections, we'll default public fields to false/none
|
// For collections, we'll default public fields to false/none
|
||||||
// since the schema doesn't have these fields yet
|
// 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
|
// Convert AssetPermissionInfo to BusterShareIndividual
|
||||||
let individual_permissions = match individual_permissions_query {
|
let individual_permissions = match individual_permissions_query {
|
||||||
|
@ -136,7 +144,7 @@ pub async fn get_collection_handler(
|
||||||
// Query for metric assets in the collection
|
// Query for metric assets in the collection
|
||||||
let metric_assets_result = collections_to_assets::table
|
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(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::collection_id.eq(req.id))
|
||||||
.filter(collections_to_assets::asset_type.eq(AssetType::MetricFile))
|
.filter(collections_to_assets::asset_type.eq(AssetType::MetricFile))
|
||||||
.filter(collections_to_assets::deleted_at.is_null())
|
.filter(collections_to_assets::deleted_at.is_null())
|
||||||
|
@ -144,8 +152,8 @@ pub async fn get_collection_handler(
|
||||||
.select((
|
.select((
|
||||||
metric_files::id,
|
metric_files::id,
|
||||||
metric_files::name,
|
metric_files::name,
|
||||||
users::name,
|
users::name.nullable(),
|
||||||
users::email,
|
users::email.nullable(),
|
||||||
metric_files::created_at,
|
metric_files::created_at,
|
||||||
metric_files::updated_at,
|
metric_files::updated_at,
|
||||||
collections_to_assets::asset_type,
|
collections_to_assets::asset_type,
|
||||||
|
@ -155,8 +163,10 @@ pub async fn get_collection_handler(
|
||||||
|
|
||||||
// Query for dashboard assets in the collection
|
// Query for dashboard assets in the collection
|
||||||
let dashboard_assets_result = collections_to_assets::table
|
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(
|
||||||
.inner_join(users::table.on(users::id.eq(dashboard_files::created_by)))
|
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::collection_id.eq(req.id))
|
||||||
.filter(collections_to_assets::asset_type.eq(AssetType::DashboardFile))
|
.filter(collections_to_assets::asset_type.eq(AssetType::DashboardFile))
|
||||||
.filter(collections_to_assets::deleted_at.is_null())
|
.filter(collections_to_assets::deleted_at.is_null())
|
||||||
|
@ -164,8 +174,8 @@ pub async fn get_collection_handler(
|
||||||
.select((
|
.select((
|
||||||
dashboard_files::id,
|
dashboard_files::id,
|
||||||
dashboard_files::name,
|
dashboard_files::name,
|
||||||
users::name,
|
users::name.nullable(),
|
||||||
users::email,
|
users::email.nullable(),
|
||||||
dashboard_files::created_at,
|
dashboard_files::created_at,
|
||||||
dashboard_files::updated_at,
|
dashboard_files::updated_at,
|
||||||
collections_to_assets::asset_type,
|
collections_to_assets::asset_type,
|
||||||
|
|
Loading…
Reference in New Issue