From c74016d3bd27ec90acd5c1a297c952b13fab7cf1 Mon Sep 17 00:00:00 2001 From: dal Date: Wed, 8 Jan 2025 22:33:14 -0700 Subject: [PATCH] Add yml_file field to Dataset model and related API structures - Introduced a new optional `yml_file` field in the Dataset model to store YAML file references. - Updated the database schema to include the `yml_file` column in the datasets table. - Modified various API request and response structures to accommodate the new `yml_file` field. - Enhanced dataset handling functions to support the inclusion of `yml_file` in dataset operations. These changes improve the dataset management capabilities by allowing the association of YAML files with datasets, facilitating better data organization and retrieval. --- .../down.sql | 3 +++ .../up.sql | 3 +++ api/src/database/models.rs | 1 + api/src/database/schema.rs | 1 + .../routes/rest/routes/datasets/post_datasets.rs | 2 ++ .../data_source_utils/data_source_utils.rs | 1 + api/src/routes/ws/datasets/dataset_utils.rs | 1 + api/src/routes/ws/datasets/post_dataset.rs | 1 + .../post_thread/sql_gen/sql_gen_handlers.rs | 2 ++ api/src/utils/query_engine/import_datasets.rs | 1 + api/src/utils/security/dataset_security.rs | 1 + cli/src/commands/init.rs | 1 - cli/src/utils/buster/api.rs | 6 ------ cli/src/utils/buster/types.rs | 1 + cli/src/utils/file/model_files.rs | 3 +++ cli/src/utils/file/profiles.rs | 13 ------------- 16 files changed, 21 insertions(+), 20 deletions(-) create mode 100644 api/migrations/2025-01-09-044455_add_yaml_field_for_datasets/down.sql create mode 100644 api/migrations/2025-01-09-044455_add_yaml_field_for_datasets/up.sql diff --git a/api/migrations/2025-01-09-044455_add_yaml_field_for_datasets/down.sql b/api/migrations/2025-01-09-044455_add_yaml_field_for_datasets/down.sql new file mode 100644 index 000000000..b60e3d2e3 --- /dev/null +++ b/api/migrations/2025-01-09-044455_add_yaml_field_for_datasets/down.sql @@ -0,0 +1,3 @@ +-- This file should undo anything in `up.sql` +ALTER TABLE datasets +DROP COLUMN yml_file; diff --git a/api/migrations/2025-01-09-044455_add_yaml_field_for_datasets/up.sql b/api/migrations/2025-01-09-044455_add_yaml_field_for_datasets/up.sql new file mode 100644 index 000000000..75aa39d57 --- /dev/null +++ b/api/migrations/2025-01-09-044455_add_yaml_field_for_datasets/up.sql @@ -0,0 +1,3 @@ +-- Your SQL goes here +ALTER TABLE datasets +ADD COLUMN yml_file TEXT; diff --git a/api/src/database/models.rs b/api/src/database/models.rs index 51f2d3cc6..54ccdfca1 100644 --- a/api/src/database/models.rs +++ b/api/src/database/models.rs @@ -159,6 +159,7 @@ pub struct Dataset { pub created_at: DateTime, pub updated_at: DateTime, pub deleted_at: Option>, + pub yml_file: Option, } #[derive(Insertable, Queryable, Associations, Debug)] diff --git a/api/src/database/schema.rs b/api/src/database/schema.rs index ef1659657..e44ec9aa4 100644 --- a/api/src/database/schema.rs +++ b/api/src/database/schema.rs @@ -238,6 +238,7 @@ diesel::table! { updated_at -> Timestamptz, deleted_at -> Nullable, model -> Nullable, + yml_file -> Nullable, } } diff --git a/api/src/routes/rest/routes/datasets/post_datasets.rs b/api/src/routes/rest/routes/datasets/post_datasets.rs index 566addf5d..bb5caa62c 100644 --- a/api/src/routes/rest/routes/datasets/post_datasets.rs +++ b/api/src/routes/rest/routes/datasets/post_datasets.rs @@ -35,6 +35,7 @@ pub struct PostDatasetsRequest { pub sql_definition: Option, pub entity_relationships: Option>, pub columns: Vec, + pub yml_file: Option, } #[derive(Debug, Deserialize)] @@ -144,6 +145,7 @@ async fn post_datasets_handler(user_id: &Uuid, requests: Vec Result> datasets::created_at, datasets::updated_at, datasets::deleted_at, + datasets::yml_file, )) .filter(datasets::data_source_id.eq(&data_source_id)) .filter(datasets::deleted_at.is_null()) diff --git a/api/src/routes/ws/datasets/dataset_utils.rs b/api/src/routes/ws/datasets/dataset_utils.rs index fe920fa0f..085859612 100644 --- a/api/src/routes/ws/datasets/dataset_utils.rs +++ b/api/src/routes/ws/datasets/dataset_utils.rs @@ -188,6 +188,7 @@ async fn get_dataset_and_columns( datasets::created_at, datasets::updated_at, datasets::deleted_at.nullable(), + datasets::yml_file.nullable(), ), ( dataset_columns::id, diff --git a/api/src/routes/ws/datasets/post_dataset.rs b/api/src/routes/ws/datasets/post_dataset.rs index 6c0b0a048..e71c0d518 100644 --- a/api/src/routes/ws/datasets/post_dataset.rs +++ b/api/src/routes/ws/datasets/post_dataset.rs @@ -257,6 +257,7 @@ async fn create_dataset(user_id: &Uuid, name: &String, data_source_id: &Uuid) -> deleted_at: None, imported: false, organization_id: user_org_id, + yml_file: None, }; let mut conn = match get_pg_pool().get().await { diff --git a/api/src/routes/ws/threads_and_messages/post_thread/sql_gen/sql_gen_handlers.rs b/api/src/routes/ws/threads_and_messages/post_thread/sql_gen/sql_gen_handlers.rs index 83ad5473b..7bc9dd3fb 100644 --- a/api/src/routes/ws/threads_and_messages/post_thread/sql_gen/sql_gen_handlers.rs +++ b/api/src/routes/ws/threads_and_messages/post_thread/sql_gen/sql_gen_handlers.rs @@ -705,6 +705,7 @@ async fn get_user_specified_dataset( datasets::created_at, datasets::updated_at, datasets::deleted_at, + datasets::yml_file, )) .inner_join( datasets_to_permission_groups::table @@ -774,6 +775,7 @@ async fn get_permissioned_datasets(pool: &PgPool, user_id: &Uuid) -> Result>(); diff --git a/api/src/utils/security/dataset_security.rs b/api/src/utils/security/dataset_security.rs index b926063b0..cc27e709e 100644 --- a/api/src/utils/security/dataset_security.rs +++ b/api/src/utils/security/dataset_security.rs @@ -42,6 +42,7 @@ pub async fn get_permissioned_datasets( datasets::created_at, datasets::updated_at, datasets::deleted_at, + datasets::yml_file, )) .inner_join( datasets_to_permission_groups::table diff --git a/cli/src/commands/init.rs b/cli/src/commands/init.rs index bbb139778..c9b7ff2f4 100644 --- a/cli/src/commands/init.rs +++ b/cli/src/commands/init.rs @@ -1,6 +1,5 @@ use anyhow::Result; use inquire::MultiSelect; -use ratatui::style::Stylize; use tokio::task::JoinSet; use crate::utils::{ diff --git a/cli/src/utils/buster/api.rs b/cli/src/utils/buster/api.rs index 6b8c115a7..3b11850e6 100644 --- a/cli/src/utils/buster/api.rs +++ b/cli/src/utils/buster/api.rs @@ -3,12 +3,6 @@ use reqwest::{ header::{HeaderMap, HeaderValue}, Client, }; -use serde::{Deserialize, Serialize}; - -use crate::{ - error::BusterError, - utils::profiles::{Credential, Profile}, -}; use super::{ PostDataSourcesRequest, PostDatasetsRequest, ValidateApiKeyRequest, ValidateApiKeyResponse, diff --git a/cli/src/utils/buster/types.rs b/cli/src/utils/buster/types.rs index f370c09d7..b9d5ff5a3 100644 --- a/cli/src/utils/buster/types.rs +++ b/cli/src/utils/buster/types.rs @@ -31,6 +31,7 @@ pub struct PostDatasetsRequest { pub sql_definition: Option, pub entity_relationships: Option>, pub columns: Vec, + pub yml_file: String, } #[derive(Debug, Serialize)] diff --git a/cli/src/utils/file/model_files.rs b/cli/src/utils/file/model_files.rs index 0499d60d7..3cc48a905 100644 --- a/cli/src/utils/file/model_files.rs +++ b/cli/src/utils/file/model_files.rs @@ -17,6 +17,7 @@ use super::{ pub struct BusterModelObject { pub sql_definition: String, pub model_file: BusterModel, + pub yml_content: String, } #[derive(Debug, Serialize, Deserialize)] @@ -91,6 +92,7 @@ async fn process_directory( model_objects.push(BusterModelObject { sql_definition, model_file: model, + yml_content: yaml_content, }); } } @@ -160,6 +162,7 @@ pub async fn upload_model_files( sql_definition: Some(model.sql_definition.clone()), entity_relationships: Some(entity_relationships), columns, + yml_file: model.yml_content.clone(), }; post_datasets_req_body.push(dataset); diff --git a/cli/src/utils/file/profiles.rs b/cli/src/utils/file/profiles.rs index 93e39a2aa..6b3ac4739 100644 --- a/cli/src/utils/file/profiles.rs +++ b/cli/src/utils/file/profiles.rs @@ -44,19 +44,6 @@ pub enum Credential { } impl Credential { - pub fn get_db_type(&self) -> String { - match self { - Credential::Postgres(_) => "postgres".to_string(), - Credential::MySQL(_) => "mysql".to_string(), - Credential::Bigquery(_) => "bigquery".to_string(), - Credential::SqlServer(_) => "sqlserver".to_string(), - Credential::Redshift(_) => "redshift".to_string(), - Credential::Databricks(_) => "databricks".to_string(), - Credential::Snowflake(_) => "snowflake".to_string(), - Credential::Starrocks(_) => "starrocks".to_string(), - } - } - pub fn get_schema(&self) -> String { match self { Credential::Postgres(cred) => cred.schema.clone(),