some changes for the sql validation

This commit is contained in:
dal 2025-02-18 11:12:29 -07:00
parent 561c31965e
commit cf45f4eddd
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
9 changed files with 869 additions and 35 deletions

View File

@ -2,7 +2,7 @@
# see https://diesel.rs/guides/configuring-diesel-cli
[print_schema]
file = "src/database/schema.rs"
file = "libs/database/src/schema.rs"
custom_type_derives = ["diesel::query_builder::QueryId", "Clone"]
filter = { except_tables = ["asset_search", "terms_search"] }

View File

@ -6,7 +6,7 @@ dev:
supabase db reset
export DATABASE_URL=postgres://postgres:postgres@127.0.0.1:54322/postgres && \
diesel migration run && \
PGPASSWORD=postgres psql -h 127.0.0.1 -p 54322 -d postgres -U postgres -f src/database/seed.sql && \
PGPASSWORD=postgres psql -h 127.0.0.1 -p 54322 -d postgres -U postgres -f src/database_dep/seed.sql && \
export RUST_LOG=debug
export CARGO_INCREMENTAL=1
nice cargo watch -x run

667
api/src/database/schema.rs Normal file
View File

@ -0,0 +1,667 @@
// @generated automatically by Diesel CLI.
pub mod sql_types {
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "asset_permission_role_enum"))]
pub struct AssetPermissionRoleEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "asset_type_enum"))]
pub struct AssetTypeEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "data_source_onboarding_status_enum"))]
pub struct DataSourceOnboardingStatusEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "dataset_type_enum"))]
pub struct DatasetTypeEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "identity_type_enum"))]
pub struct IdentityTypeEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "message_feedback_enum"))]
pub struct MessageFeedbackEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "sharing_setting_enum"))]
pub struct SharingSettingEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "stored_values_status_enum"))]
pub struct StoredValuesStatusEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "team_role_enum"))]
pub struct TeamRoleEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "user_organization_role_enum"))]
pub struct UserOrganizationRoleEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "user_organization_status_enum"))]
pub struct UserOrganizationStatusEnum;
#[derive(diesel::query_builder::QueryId, Clone, diesel::sql_types::SqlType)]
#[diesel(postgres_type(name = "verification_enum"))]
pub struct VerificationEnum;
}
diesel::table! {
api_keys (id) {
id -> Uuid,
owner_id -> Uuid,
key -> Text,
organization_id -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::IdentityTypeEnum;
use super::sql_types::AssetTypeEnum;
use super::sql_types::AssetPermissionRoleEnum;
asset_permissions (identity_id, asset_id, asset_type, identity_type) {
identity_id -> Uuid,
identity_type -> IdentityTypeEnum,
asset_id -> Uuid,
asset_type -> AssetTypeEnum,
role -> AssetPermissionRoleEnum,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
created_by -> Uuid,
updated_by -> Uuid,
}
}
diesel::table! {
collections (id) {
id -> Uuid,
name -> Text,
description -> Nullable<Text>,
created_by -> Uuid,
updated_by -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
organization_id -> Uuid,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::AssetTypeEnum;
collections_to_assets (collection_id, asset_id, asset_type) {
collection_id -> Uuid,
asset_id -> Uuid,
asset_type -> AssetTypeEnum,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
created_by -> Uuid,
updated_by -> Uuid,
}
}
diesel::table! {
dashboard_files (id) {
id -> Uuid,
name -> Varchar,
file_name -> Varchar,
content -> Jsonb,
filter -> Nullable<Varchar>,
organization_id -> Uuid,
created_by -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
dashboard_versions (id) {
id -> Uuid,
dashboard_id -> Uuid,
config -> Jsonb,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
dashboards (id) {
id -> Uuid,
name -> Text,
description -> Nullable<Text>,
config -> Jsonb,
publicly_accessible -> Bool,
publicly_enabled_by -> Nullable<Uuid>,
public_expiry_date -> Nullable<Timestamptz>,
password_secret_id -> Nullable<Uuid>,
created_by -> Uuid,
updated_by -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
organization_id -> Uuid,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::DataSourceOnboardingStatusEnum;
data_sources (id) {
id -> Uuid,
name -> Text,
#[sql_name = "type"]
type_ -> Text,
secret_id -> Uuid,
onboarding_status -> DataSourceOnboardingStatusEnum,
onboarding_error -> Nullable<Text>,
organization_id -> Uuid,
created_by -> Uuid,
updated_by -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
env -> Varchar,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::StoredValuesStatusEnum;
dataset_columns (id) {
id -> Uuid,
dataset_id -> Uuid,
name -> Text,
#[sql_name = "type"]
type_ -> Text,
description -> Nullable<Text>,
nullable -> Bool,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
stored_values -> Nullable<Bool>,
stored_values_status -> Nullable<StoredValuesStatusEnum>,
stored_values_error -> Nullable<Text>,
stored_values_count -> Nullable<Int8>,
stored_values_last_synced -> Nullable<Timestamptz>,
semantic_type -> Nullable<Text>,
dim_type -> Nullable<Text>,
expr -> Nullable<Text>,
}
}
diesel::table! {
dataset_groups (id) {
id -> Uuid,
organization_id -> Uuid,
name -> Varchar,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
dataset_groups_permissions (id) {
id -> Uuid,
dataset_group_id -> Uuid,
permission_id -> Uuid,
permission_type -> Varchar,
organization_id -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
dataset_permissions (id) {
id -> Uuid,
organization_id -> Uuid,
dataset_id -> Uuid,
permission_id -> Uuid,
permission_type -> Varchar,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::DatasetTypeEnum;
datasets (id) {
id -> Uuid,
name -> Text,
database_name -> Text,
when_to_use -> Nullable<Text>,
when_not_to_use -> Nullable<Text>,
#[sql_name = "type"]
type_ -> DatasetTypeEnum,
definition -> Text,
schema -> Text,
enabled -> Bool,
imported -> Bool,
data_source_id -> Uuid,
organization_id -> Uuid,
created_by -> Uuid,
updated_by -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
model -> Nullable<Text>,
yml_file -> Nullable<Text>,
database_identifier -> Nullable<Text>,
}
}
diesel::table! {
datasets_to_dataset_groups (dataset_id, dataset_group_id) {
dataset_id -> Uuid,
dataset_group_id -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
datasets_to_permission_groups (dataset_id, permission_group_id) {
dataset_id -> Uuid,
permission_group_id -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
entity_relationship (primary_dataset_id, foreign_dataset_id) {
primary_dataset_id -> Uuid,
foreign_dataset_id -> Uuid,
relationship_type -> Text,
created_at -> Timestamptz,
}
}
diesel::table! {
messages (id) {
id -> Uuid,
request -> Text,
response -> Jsonb,
thread_id -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
created_by -> Uuid,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::MessageFeedbackEnum;
use super::sql_types::VerificationEnum;
messages_deprecated (id) {
id -> Uuid,
thread_id -> Uuid,
sent_by -> Uuid,
message -> Text,
responses -> Nullable<Jsonb>,
code -> Nullable<Text>,
context -> Nullable<Jsonb>,
title -> Nullable<Text>,
feedback -> Nullable<MessageFeedbackEnum>,
verification -> VerificationEnum,
dataset_id -> Nullable<Uuid>,
chart_config -> Nullable<Jsonb>,
chart_recommendations -> Nullable<Jsonb>,
time_frame -> Nullable<Text>,
data_metadata -> Nullable<Jsonb>,
draft_session_id -> Nullable<Uuid>,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
draft_state -> Nullable<Jsonb>,
summary_question -> Nullable<Text>,
sql_evaluation_id -> Nullable<Uuid>,
}
}
diesel::table! {
messages_to_files (id) {
id -> Uuid,
message_id -> Uuid,
file_id -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::VerificationEnum;
metric_files (id) {
id -> Uuid,
name -> Varchar,
file_name -> Varchar,
content -> Jsonb,
verification -> VerificationEnum,
evaluation_obj -> Nullable<Jsonb>,
evaluation_summary -> Nullable<Text>,
evaluation_score -> Nullable<Float8>,
organization_id -> Uuid,
created_by -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
organizations (id) {
id -> Uuid,
name -> Text,
domain -> Nullable<Text>,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
permission_groups (id) {
id -> Uuid,
name -> Text,
organization_id -> Uuid,
created_by -> Uuid,
updated_by -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::IdentityTypeEnum;
permission_groups_to_identities (permission_group_id, identity_id, identity_type) {
permission_group_id -> Uuid,
identity_id -> Uuid,
identity_type -> IdentityTypeEnum,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
created_by -> Uuid,
updated_by -> Uuid,
}
}
diesel::table! {
permission_groups_to_users (permission_group_id, user_id) {
permission_group_id -> Uuid,
user_id -> Uuid,
created_at -> Timestamptz,
}
}
diesel::table! {
sql_evaluations (id) {
id -> Uuid,
evaluation_obj -> Jsonb,
evaluation_summary -> Text,
score -> Text,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::SharingSettingEnum;
teams (id) {
id -> Uuid,
name -> Text,
organization_id -> Uuid,
sharing_setting -> SharingSettingEnum,
edit_sql -> Bool,
upload_csv -> Bool,
export_assets -> Bool,
email_slack_enabled -> Bool,
created_by -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::TeamRoleEnum;
teams_to_users (team_id, user_id) {
team_id -> Uuid,
user_id -> Uuid,
role -> TeamRoleEnum,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
terms (id) {
id -> Uuid,
name -> Text,
definition -> Nullable<Text>,
sql_snippet -> Nullable<Text>,
organization_id -> Uuid,
created_by -> Uuid,
updated_by -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
terms_to_datasets (term_id, dataset_id) {
term_id -> Uuid,
dataset_id -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
threads (id) {
id -> Uuid,
title -> Text,
organization_id -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
created_by -> Uuid,
}
}
diesel::table! {
threads_deprecated (id) {
id -> Uuid,
created_by -> Uuid,
updated_by -> Uuid,
publicly_accessible -> Bool,
publicly_enabled_by -> Nullable<Uuid>,
public_expiry_date -> Nullable<Timestamptz>,
password_secret_id -> Nullable<Uuid>,
state_message_id -> Nullable<Uuid>,
parent_thread_id -> Nullable<Uuid>,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
organization_id -> Uuid,
}
}
diesel::table! {
threads_to_dashboards (thread_id, dashboard_id) {
thread_id -> Uuid,
dashboard_id -> Uuid,
added_by -> Uuid,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::AssetTypeEnum;
user_favorites (user_id, asset_id, asset_type) {
user_id -> Uuid,
asset_id -> Uuid,
asset_type -> AssetTypeEnum,
order_index -> Int4,
created_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
}
}
diesel::table! {
users (id) {
id -> Uuid,
email -> Text,
name -> Nullable<Text>,
config -> Jsonb,
created_at -> Timestamptz,
updated_at -> Timestamptz,
attributes -> Jsonb,
}
}
diesel::table! {
use diesel::sql_types::*;
use super::sql_types::UserOrganizationRoleEnum;
use super::sql_types::SharingSettingEnum;
use super::sql_types::UserOrganizationStatusEnum;
users_to_organizations (user_id, organization_id) {
user_id -> Uuid,
organization_id -> Uuid,
role -> UserOrganizationRoleEnum,
sharing_setting -> SharingSettingEnum,
edit_sql -> Bool,
upload_csv -> Bool,
export_assets -> Bool,
email_slack_enabled -> Bool,
created_at -> Timestamptz,
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
created_by -> Uuid,
updated_by -> Uuid,
deleted_by -> Nullable<Uuid>,
status -> UserOrganizationStatusEnum,
}
}
diesel::joinable!(api_keys -> organizations (organization_id));
diesel::joinable!(api_keys -> users (owner_id));
diesel::joinable!(collections -> organizations (organization_id));
diesel::joinable!(dashboard_versions -> dashboards (dashboard_id));
diesel::joinable!(dashboards -> organizations (organization_id));
diesel::joinable!(data_sources -> organizations (organization_id));
diesel::joinable!(dataset_groups -> organizations (organization_id));
diesel::joinable!(dataset_groups_permissions -> dataset_groups (dataset_group_id));
diesel::joinable!(dataset_groups_permissions -> organizations (organization_id));
diesel::joinable!(dataset_permissions -> datasets (dataset_id));
diesel::joinable!(dataset_permissions -> organizations (organization_id));
diesel::joinable!(datasets -> data_sources (data_source_id));
diesel::joinable!(datasets -> organizations (organization_id));
diesel::joinable!(datasets_to_dataset_groups -> dataset_groups (dataset_group_id));
diesel::joinable!(datasets_to_dataset_groups -> datasets (dataset_id));
diesel::joinable!(datasets_to_permission_groups -> datasets (dataset_id));
diesel::joinable!(datasets_to_permission_groups -> permission_groups (permission_group_id));
diesel::joinable!(messages -> users (created_by));
diesel::joinable!(messages_deprecated -> datasets (dataset_id));
diesel::joinable!(messages_deprecated -> threads_deprecated (thread_id));
diesel::joinable!(messages_deprecated -> users (sent_by));
diesel::joinable!(messages_to_files -> messages (message_id));
diesel::joinable!(permission_groups -> organizations (organization_id));
diesel::joinable!(permission_groups_to_users -> permission_groups (permission_group_id));
diesel::joinable!(permission_groups_to_users -> users (user_id));
diesel::joinable!(teams -> organizations (organization_id));
diesel::joinable!(teams -> users (created_by));
diesel::joinable!(teams_to_users -> teams (team_id));
diesel::joinable!(teams_to_users -> users (user_id));
diesel::joinable!(terms -> organizations (organization_id));
diesel::joinable!(terms_to_datasets -> datasets (dataset_id));
diesel::joinable!(terms_to_datasets -> terms (term_id));
diesel::joinable!(threads -> organizations (organization_id));
diesel::joinable!(threads -> users (created_by));
diesel::joinable!(threads_deprecated -> organizations (organization_id));
diesel::joinable!(threads_to_dashboards -> dashboards (dashboard_id));
diesel::joinable!(threads_to_dashboards -> threads_deprecated (thread_id));
diesel::joinable!(threads_to_dashboards -> users (added_by));
diesel::joinable!(user_favorites -> users (user_id));
diesel::joinable!(users_to_organizations -> organizations (organization_id));
diesel::allow_tables_to_appear_in_same_query!(
api_keys,
asset_permissions,
collections,
collections_to_assets,
dashboard_files,
dashboard_versions,
dashboards,
data_sources,
dataset_columns,
dataset_groups,
dataset_groups_permissions,
dataset_permissions,
datasets,
datasets_to_dataset_groups,
datasets_to_permission_groups,
entity_relationship,
messages,
messages_deprecated,
messages_to_files,
metric_files,
organizations,
permission_groups,
permission_groups_to_identities,
permission_groups_to_users,
sql_evaluations,
teams,
teams_to_users,
terms,
terms_to_datasets,
threads,
threads_deprecated,
threads_to_dashboards,
user_favorites,
users,
users_to_organizations,
);

View File

@ -0,0 +1,63 @@
use anyhow::{anyhow, Result};
use tracing::debug;
use uuid::Uuid;
use crate::database_dep::{lib::get_pg_pool, schema::metric_files};
use crate::utils::query_engine::query_engine::query_engine;
use diesel::{ExpressionMethods, QueryDsl};
use diesel_async::RunQueryDsl;
/// Validates SQL query using existing query engine by attempting to run it
/// Returns Ok(()) if valid, Err with description if invalid
pub async fn validate_sql(sql: &str, dataset_id: &Uuid) -> Result<()> {
debug!("Validating SQL query for dataset {}", dataset_id);
if sql.trim().is_empty() {
return Err(anyhow!("SQL query cannot be empty"));
}
// Try to execute the query using query_engine
match query_engine(dataset_id, &sql.to_string()).await {
Ok(_) => Ok(()),
Err(e) => Err(anyhow!("SQL validation failed: {}", e)),
}
}
/// Validates existence of metric IDs in database
/// Returns Result with list of missing IDs if any
pub async fn validate_metric_ids(ids: &[Uuid]) -> Result<Vec<Uuid>> {
let mut conn = get_pg_pool().get().await?;
// Query for existing IDs
let existing_ids = metric_files::table
.filter(metric_files::id.eq_any(ids))
.filter(metric_files::deleted_at.is_null())
.select(metric_files::id)
.load::<Uuid>(&mut conn)
.await?;
// Find missing IDs by comparing with input IDs
let missing_ids: Vec<Uuid> = ids
.iter()
.filter(|id| !existing_ids.contains(id))
.cloned()
.collect();
Ok(missing_ids)
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_validate_sql_empty() {
let dataset_id = Uuid::new_v4();
let result = validate_sql("", &dataset_id).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("cannot be empty"));
}
// 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
}

View File

@ -8,6 +8,7 @@ use diesel_async::RunQueryDsl;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;
use tracing::{debug, error};
use crate::{
database_dep::{
@ -22,6 +23,7 @@ use crate::{
use super::{
file_types::{dashboard_yml::DashboardYml, file::FileEnum, metric_yml::MetricYml},
FileModificationTool,
common::{validate_sql, validate_metric_ids},
};
use litellm::ToolCall;
@ -106,6 +108,12 @@ impl ToolExecutor for CreateFilesTool {
match MetricYml::new(file.yml_content.clone()) {
Ok(metric_yml) => {
if let Some(metric_id) = &metric_yml.id {
// Validate SQL before creating the record
if let Err(e) = validate_sql(&metric_yml.sql, metric_id).await {
failed_files.push((file.name, format!("SQL validation failed: {}", e)));
continue;
}
let metric_file = MetricFile {
id: metric_id.clone(),
name: metric_yml.title.clone(),
@ -136,6 +144,27 @@ impl ToolExecutor for CreateFilesTool {
match DashboardYml::new(file.yml_content.clone()) {
Ok(dashboard_yml) => {
if let Some(dashboard_id) = &dashboard_yml.id {
// Collect and validate metric IDs from rows
let metric_ids: Vec<Uuid> = dashboard_yml.rows
.iter()
.flat_map(|row| row.items.iter())
.map(|item| item.id)
.collect();
if !metric_ids.is_empty() {
match validate_metric_ids(&metric_ids).await {
Ok(missing_ids) if !missing_ids.is_empty() => {
failed_files.push((file.name, format!("Referenced metrics not found: {:?}", missing_ids)));
continue;
}
Err(e) => {
failed_files.push((file.name, format!("Failed to validate metric IDs: {}", e)));
continue;
}
Ok(_) => (), // All metrics exist
}
}
let dashboard_file = DashboardFile {
id: dashboard_id.clone(),
name: dashboard_yml.name.clone(),

View File

@ -13,17 +13,17 @@ pub struct DashboardYml {
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Row {
items: Vec<RowItem>, // max number of items in a row is 4, min is 1
pub items: Vec<RowItem>, // max number of items in a row is 4, min is 1
#[serde(skip_serializing_if = "Option::is_none")]
row_height: Option<u32>, // max is 550, min is 320
pub row_height: Option<u32>, // max is 550, min is 320
#[serde(skip_serializing_if = "Option::is_none")]
column_sizes: Option<Vec<u32>>, // max sum of elements is 12 min is 3
pub column_sizes: Option<Vec<u32>>, // max sum of elements is 12 min is 3
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct RowItem {
// This id is the id of the metric or item reference that goes here in the dashboard.
id: Uuid,
pub id: Uuid,
}
impl DashboardYml {

View File

@ -5,6 +5,7 @@ pub mod open_files;
pub mod search_data_catalog;
pub mod search_files;
pub mod send_files_to_user;
pub mod common;
pub use modify_files::ModifyFilesTool;
pub use create_files::CreateFilesTool;

View File

@ -26,6 +26,7 @@ use super::{
metric_yml::MetricYml,
},
FileModificationTool,
common::{validate_sql, validate_metric_ids},
};
use litellm::ToolCall;
@ -628,6 +629,33 @@ async fn process_metric_file(
file_name = %modification.file_name,
"Successfully modified and validated metric file"
);
// Validate SQL if it was modified
let sql_changed = current_yml.sql != new_yml.sql;
if sql_changed {
if let Err(e) = validate_sql(&new_yml.sql, &file.id).await {
let error = format!("SQL validation failed: {}", e);
error!(
file_id = %file.id,
file_name = %modification.file_name,
error = %error,
"SQL validation error"
);
results.push(ModificationResult {
file_id: file.id,
file_type: "metric".to_string(),
file_name: modification.file_name.clone(),
success: false,
original_lines: original_lines.clone(),
adjusted_lines: adjusted_lines.clone(),
error: Some(error.clone()),
modification_type: "sql_validation".to_string(),
timestamp: Utc::now(),
duration,
});
return Err(anyhow::anyhow!(error));
}
}
// Update file record
file.content = serde_json::to_value(&new_yml)?;
@ -704,6 +732,12 @@ async fn process_dashboard_file(
modification: &FileModification,
duration: i64,
) -> Result<(DashboardFile, DashboardYml, Vec<ModificationResult>)> {
debug!(
file_id = %file.id,
file_name = %modification.file_name,
"Processing dashboard file modifications"
);
let mut results = Vec::new();
// Parse existing content
@ -775,32 +809,72 @@ async fn process_dashboard_file(
// Create and validate new YML object
match DashboardYml::new(modified_content) {
Ok(new_yml) => {
// Update file record
file.content = match serde_json::to_value(&new_yml) {
Ok(content) => content,
Err(e) => {
let error = format!("Failed to serialize modified dashboard YAML: {}", e);
error!(
file_id = %file.id,
file_name = %modification.file_name,
error = %error,
"YAML serialization error"
);
results.push(ModificationResult {
file_id: file.id,
file_type: "dashboard".to_string(),
file_name: modification.file_name.clone(),
success: false,
original_lines,
adjusted_lines: vec![],
error: Some(error.clone()),
modification_type: "serialization".to_string(),
timestamp: Utc::now(),
duration,
});
return Err(anyhow::anyhow!(error));
debug!(
file_id = %file.id,
file_name = %modification.file_name,
"Successfully modified and validated dashboard file"
);
// Collect all metric IDs from rows
let metric_ids: Vec<Uuid> = new_yml.rows
.iter()
.flat_map(|row| row.items.iter())
.map(|item| item.id)
.collect();
// Validate metric IDs if any exist
if !metric_ids.is_empty() {
match validate_metric_ids(&metric_ids).await {
Ok(missing_ids) if !missing_ids.is_empty() => {
let error = format!("Referenced metrics not found: {:?}", missing_ids);
error!(
file_id = %file.id,
file_name = %modification.file_name,
error = %error,
"Metric validation error"
);
results.push(ModificationResult {
file_id: file.id,
file_type: "dashboard".to_string(),
file_name: modification.file_name.clone(),
success: false,
original_lines: original_lines.clone(),
adjusted_lines: adjusted_lines.clone(),
error: Some(error.clone()),
modification_type: "metric_validation".to_string(),
timestamp: Utc::now(),
duration,
});
return Err(anyhow::anyhow!(error));
},
Err(e) => {
let error = format!("Failed to validate metric IDs: {}", e);
error!(
file_id = %file.id,
file_name = %modification.file_name,
error = %error,
"Metric validation error"
);
results.push(ModificationResult {
file_id: file.id,
file_type: "dashboard".to_string(),
file_name: modification.file_name.clone(),
success: false,
original_lines: original_lines.clone(),
adjusted_lines: adjusted_lines.clone(),
error: Some(error.clone()),
modification_type: "metric_validation".to_string(),
timestamp: Utc::now(),
duration,
});
return Err(anyhow::anyhow!(error));
},
Ok(_) => (), // All metrics exist
}
};
}
// Update file record
file.content = serde_json::to_value(&new_yml)?;
file.updated_at = Utc::now();
// Track successful modification
@ -820,7 +894,7 @@ async fn process_dashboard_file(
Ok((file, new_yml, results))
}
Err(e) => {
let error = format!("Failed to validate modified dashboard YAML: {}", e);
let error = format!("Failed to validate modified YAML: {}", e);
error!(
file_id = %file.id,
file_name = %modification.file_name,
@ -839,7 +913,7 @@ async fn process_dashboard_file(
timestamp: Utc::now(),
duration,
});
return Err(anyhow::anyhow!(error));
Err(anyhow::anyhow!(error))
}
}
}
@ -863,7 +937,7 @@ async fn process_dashboard_file(
timestamp: Utc::now(),
duration,
});
return Err(anyhow::anyhow!(error));
Err(anyhow::anyhow!(error))
}
}
}

View File

@ -255,7 +255,7 @@ impl ToolExecutor for SearchFilesTool {
},
"additionalProperties": false
},
"description": "Searches for existing metric and dashboard files. Only use when user explicitly asks about existing/previous/old files. Returns up to 10 most relevant files ordered by relevance."
"description": "Useful for when a user explicitly asks about old metrics or dashboars. Searches for existing metric and dashboard files. Only use when user explicitly asks about existing/previous/old files. Returns up to 10 most relevant files ordered by relevance."
})
}
}