From af3262d0c1cb46e6b8934e3813c7ce1c87d7ae40 Mon Sep 17 00:00:00 2001 From: dal Date: Mon, 5 May 2025 18:05:51 -0600 Subject: [PATCH] refactor for modular and testable code --- cli/Cargo.toml | 1 + cli/cli/Cargo.toml | 1 + cli/cli/src/commands/init.rs | 1246 ++++++++++++++-------------------- 3 files changed, 505 insertions(+), 743 deletions(-) diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 4cf8a6a1b..9e60d6215 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -48,6 +48,7 @@ chrono = "0.4" # Specify the version here semver = "1.0.19" crossterm = "0.29" # Add crossterm explicitly rustyline = "15.0.0" +once_cell = "1.19.0" # Keep dev-dependencies separate if they aren't shared # tempfile = "3.16.0" \ No newline at end of file diff --git a/cli/cli/Cargo.toml b/cli/cli/Cargo.toml index 827beb4fd..2bd75858e 100644 --- a/cli/cli/Cargo.toml +++ b/cli/cli/Cargo.toml @@ -42,6 +42,7 @@ zip = { workspace = true } glob = { workspace = true } walkdir = { workspace = true } semver = { workspace = true } +once_cell = { workspace = true } # Add the shared query engine library query_engine = { workspace = true } diff --git a/cli/cli/src/commands/init.rs b/cli/cli/src/commands/init.rs index 666cc2d04..2a5e81ef9 100644 --- a/cli/cli/src/commands/init.rs +++ b/cli/cli/src/commands/init.rs @@ -2,23 +2,24 @@ use anyhow::Result; use colored::*; use indicatif::{ProgressBar, ProgressStyle}; use inquire::{validator::Validation, Confirm, Password, Select, Text}; +use once_cell::sync::Lazy; use query_engine::credentials::{ BigqueryCredentials, Credential, DatabricksCredentials, MySqlCredentials, PostgresCredentials, RedshiftCredentials, SnowflakeCredentials, SqlServerCredentials, }; use regex::Regex; -use serde::{Deserialize, Serialize}; +use serde::Deserialize; use serde_yaml; -use std::error::Error; use std::fs; use std::path::{Path, PathBuf}; use std::time::Duration; // Update imports to use new modules via mod.rs re-exports use crate::utils::{ - BusterConfig, ProjectContext, // Use re-exported items directly buster::{BusterClient, PostDataSourcesRequest}, file::buster_credentials::get_and_validate_buster_credentials, + BusterConfig, + ProjectContext, // Use re-exported items directly }; #[derive(Debug, Clone)] @@ -46,8 +47,6 @@ impl std::fmt::Display for DatabaseType { } } -// Using shared RedshiftCredentials from query_engine now, no need for local definition - // Helper struct to parse dbt_project.yml #[derive(Debug, Deserialize)] struct DbtProject { @@ -62,8 +61,11 @@ fn parse_dbt_project(base_dir: &Path) -> Result> { if dbt_project_path.exists() && dbt_project_path.is_file() { println!( "{}", - format!("Found {}, attempting to read config...", dbt_project_path.display()) - .dimmed() + format!( + "Found {}, attempting to read config...", + dbt_project_path.display() + ) + .dimmed() ); match fs::read_to_string(&dbt_project_path) { Ok(content) => { @@ -77,13 +79,18 @@ fn parse_dbt_project(base_dir: &Path) -> Result> { Ok(None) } } - }, + } Err(e) => { - eprintln!( - "{}", - format!("Warning: Failed to read {}: {}. Proceeding without dbt project info.", dbt_project_path.display(), e).yellow() - ); - Ok(None) + eprintln!( + "{}", + format!( + "Warning: Failed to read {}: {}. Proceeding without dbt project info.", + dbt_project_path.display(), + e + ) + .yellow() + ); + Ok(None) } } } else { @@ -91,6 +98,123 @@ fn parse_dbt_project(base_dir: &Path) -> Result> { } } +// --- Input Helper Functions --- + +static NAME_REGEX: Lazy = Lazy::new(|| Regex::new(r"^[a-zA-Z0-9_-]+$").unwrap()); + +fn prompt_required_text(prompt: &str, help_message: Option<&str>) -> Result { + let mut text_prompt = Text::new(prompt); + if let Some(help) = help_message { + text_prompt = text_prompt.with_help_message(help); + } + let prompt_string = prompt.to_string(); // Clone prompt + text_prompt + .with_validator(move |input: &str| { // Add move keyword + if input.trim().is_empty() { + // Use the cloned prompt_string here + Ok(Validation::Invalid(format!("{} cannot be empty", prompt_string.replace(":", "")).into())) + } else { + Ok(Validation::Valid) + } + }) + .prompt() + .map_err(Into::into) // Convert inquire::Error to anyhow::Error +} + +fn prompt_validated_name(prompt: &str, suggested_name: Option<&str>) -> Result { + let mut name_prompt = Text::new(prompt) + .with_help_message("Only alphanumeric characters, dash (-) and underscore (_) allowed"); + + if let Some(s_name) = suggested_name { + name_prompt = name_prompt.with_default(s_name); + } + + name_prompt + .with_validator(move |input: &str| { + if input.trim().is_empty() { + Ok(Validation::Invalid("Name cannot be empty".into())) + } else if NAME_REGEX.is_match(input) { + Ok(Validation::Valid) + } else { + Ok(Validation::Invalid( + "Name must contain only alphanumeric characters, dash (-) or underscore (_)" + .into(), + )) + } + }) + .prompt() + .map_err(Into::into) +} + +fn prompt_password(prompt: &str) -> Result { + Password::new(prompt) + .with_validator(|input: &str| { + if input.trim().is_empty() { + Ok(Validation::Invalid("Password cannot be empty".into())) + } else { + Ok(Validation::Valid) + } + }) + .without_confirmation() + .prompt() + .map_err(Into::into) +} + +fn prompt_u16_with_default(prompt: &str, default: &str, help_message: Option<&str>) -> Result { + let mut port_prompt = Text::new(prompt).with_default(default); + if let Some(help) = help_message { + port_prompt = port_prompt.with_help_message(help); + } + let port_str = port_prompt + .with_validator(|input: &str| match input.parse::() { + Ok(_) => Ok(Validation::Valid), + Err(_) => Ok(Validation::Invalid( + "Port must be a valid number between 1 and 65535".into(), + )), + }) + .prompt()?; + port_str.parse::().map_err(Into::into) // Convert parse error +} + +// --- End Input Helper Functions --- + +// --- API Interaction Helper --- + +async fn create_data_source_with_progress( + client: &BusterClient, + request: PostDataSourcesRequest, +) -> Result<()> { + let spinner = ProgressBar::new_spinner(); + spinner.set_style( + ProgressStyle::default_spinner() + .tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ ") + .template("{spinner:.green} {msg}") + .unwrap(), + ); + spinner.set_message("Sending credentials to Buster API..."); + spinner.enable_steady_tick(Duration::from_millis(100)); + + match client.post_data_sources(request).await { + Ok(_) => { + spinner.finish_with_message( + "✓ Data source created successfully!" + .green() + .bold() + .to_string(), + ); + Ok(()) + } + Err(e) => { + spinner.finish_with_message("✗ Failed to create data source".red().bold().to_string()); + println!("\nError: {}", e); + println!("Please check your credentials and try again."); + Err(anyhow::anyhow!("Failed to create data source: {}", e)) + } + } +} + +// --- End API Interaction Helper --- + pub async fn init(destination_path: Option<&str>) -> Result<()> { println!("{}", "Initializing Buster...".bold().green()); @@ -124,18 +248,21 @@ pub async fn init(destination_path: Option<&str>) -> Result<()> { } } - // --- Try to parse dbt_project.yml --- + // --- Try to parse dbt_project.yml --- let dbt_config = parse_dbt_project(&dest_path)?; // Extract suggested name (if available) let suggested_name = dbt_config.as_ref().and_then(|c| c.name.as_deref()); if let Some(name) = suggested_name { - println!( + println!( "{}", - format!("Suggesting data source name '{}' from dbt_project.yml", name.cyan()) - .dimmed() + format!( + "Suggesting data source name '{}' from dbt_project.yml", + name.cyan() + ) + .dimmed() ); } - // --- End dbt_project.yml parsing --- + // --- End dbt_project.yml parsing --- // Check for Buster credentials with progress indicator let spinner = ProgressBar::new_spinner(); @@ -175,29 +302,74 @@ pub async fn init(destination_path: Option<&str>) -> Result<()> { let db_type = Select::new("Select your database type:", db_types).prompt()?; - println!("You selected: {}", db_type.to_string().cyan()); + println!( + "{}", + format!("You selected: {}", db_type.to_string().cyan()).dimmed() + ); match db_type { DatabaseType::Redshift => { - setup_redshift(buster_creds.url, buster_creds.api_key, &config_path, suggested_name).await + setup_redshift( + buster_creds.url, + buster_creds.api_key, + &config_path, + suggested_name, + ) + .await } DatabaseType::Postgres => { - setup_postgres(buster_creds.url, buster_creds.api_key, &config_path, suggested_name).await + setup_postgres( + buster_creds.url, + buster_creds.api_key, + &config_path, + suggested_name, + ) + .await } DatabaseType::BigQuery => { - setup_bigquery(buster_creds.url, buster_creds.api_key, &config_path, suggested_name).await + setup_bigquery( + buster_creds.url, + buster_creds.api_key, + &config_path, + suggested_name, + ) + .await } DatabaseType::Snowflake => { - setup_snowflake(buster_creds.url, buster_creds.api_key, &config_path, suggested_name).await + setup_snowflake( + buster_creds.url, + buster_creds.api_key, + &config_path, + suggested_name, + ) + .await } DatabaseType::MySql => { - setup_mysql(buster_creds.url, buster_creds.api_key, &config_path, suggested_name).await + setup_mysql( + buster_creds.url, + buster_creds.api_key, + &config_path, + suggested_name, + ) + .await } DatabaseType::SqlServer => { - setup_sqlserver(buster_creds.url, buster_creds.api_key, &config_path, suggested_name).await + setup_sqlserver( + buster_creds.url, + buster_creds.api_key, + &config_path, + suggested_name, + ) + .await } DatabaseType::Databricks => { - setup_databricks(buster_creds.url, buster_creds.api_key, &config_path, suggested_name).await + setup_databricks( + buster_creds.url, + buster_creds.api_key, + &config_path, + suggested_name, + ) + .await } } } @@ -214,26 +386,27 @@ async fn setup_redshift( let name_regex = Regex::new(r"^[a-zA-Z0-9_-]+$")?; let mut name_prompt = Text::new("Enter a unique name for this data source:") .with_help_message("Only alphanumeric characters, dash (-) and underscore (_) allowed"); - + // Set default if provided if let Some(s_name) = suggested_name { name_prompt = name_prompt.with_default(s_name); } - let name = name_prompt.with_validator(move |input: &str| { - if input.trim().is_empty() { - return Ok(Validation::Invalid("Name cannot be empty".into())); - } - if name_regex.is_match(input) { - Ok(Validation::Valid) - } else { - Ok(Validation::Invalid( - "Name must contain only alphanumeric characters, dash (-) or underscore (_)" - .into(), - )) - } - }) - .prompt()?; + let name = name_prompt + .with_validator(move |input: &str| { + if input.trim().is_empty() { + return Ok(Validation::Invalid("Name cannot be empty".into())); + } + if name_regex.is_match(input) { + Ok(Validation::Valid) + } else { + Ok(Validation::Invalid( + "Name must contain only alphanumeric characters, dash (-) or underscore (_)" + .into(), + )) + } + }) + .prompt()?; // Collect host let host = Text::new("Enter the Redshift host:") @@ -291,13 +464,14 @@ async fn setup_redshift( .prompt()?; // Collect schema (required) - let schema = Text::new("Enter the default Redshift schema (optional):") + let schema = Text::new("Enter the default Redshift schema:") + .with_validator(|input: &str| { + if input.trim().is_empty() { + return Ok(Validation::Invalid("Schema cannot be empty".into())); + } + Ok(Validation::Valid) + }) .prompt()?; - let schema_opt = if schema.trim().is_empty() { - None - } else { - Some(schema.clone()) - }; // Show summary and confirm println!("\n{}", "Connection Summary:".bold()); @@ -307,9 +481,7 @@ async fn setup_redshift( println!("Username: {}", username.cyan()); println!("Password: {}", "********".cyan()); println!("Default Database: {}", database.cyan()); - if let Some(s) = &schema_opt { - println!("Default Schema: {}", s.cyan()); - } + println!("Default Schema: {}", schema.cyan()); let confirm = Confirm::new("Do you want to create this data source?") .with_default(true) @@ -327,7 +499,7 @@ async fn setup_redshift( username, password, default_database: database.clone(), - default_schema: schema_opt.clone(), + default_schema: Some(schema.clone()), }; let credential = Credential::Redshift(redshift_creds); let request = PostDataSourcesRequest { @@ -336,49 +508,14 @@ async fn setup_redshift( }; // Send to API with progress indicator - let spinner = ProgressBar::new_spinner(); - spinner.set_style( - ProgressStyle::default_spinner() - .tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ ") - .template("{spinner:.green} {msg}") - .unwrap(), - ); - spinner.set_message("Sending credentials to Buster API..."); - spinner.enable_steady_tick(Duration::from_millis(100)); - let client = BusterClient::new(buster_url, buster_api_key)?; + create_data_source_with_progress(&client, request).await?; - match client.post_data_sources(request).await { - Ok(_) => { - spinner.finish_with_message( - "✓ Data source created successfully!" - .green() - .bold() - .to_string(), - ); - println!( - "\nData source '{}' is now available for use with Buster.", - name.cyan() - ); + // Create buster.yml file + create_buster_config_file(config_path, &name, &database, Some(&schema))?; - // Create buster.yml file - create_buster_config_file( - config_path, - &name, - &database, - schema_opt.as_deref(), - )?; - - println!("You can now use this data source with other Buster commands."); - Ok(()) - } - Err(e) => { - spinner.finish_with_message("✗ Failed to create data source".red().bold().to_string()); - println!("\nError: {}", e); - println!("Please check your credentials and try again."); - Err(anyhow::anyhow!("Failed to create data source: {}", e)) - } - } + println!("You can now use this data source with other Buster commands."); + Ok(()) } async fn setup_postgres( @@ -389,94 +526,21 @@ async fn setup_postgres( ) -> Result<()> { println!("{}", "Setting up PostgreSQL connection...".bold().green()); - // Collect name (with validation) - let name_regex = Regex::new(r"^[a-zA-Z0-9_-]+$")?; - let mut name_prompt = Text::new("Enter a unique name for this data source:") - .with_help_message("Only alphanumeric characters, dash (-) and underscore (_) allowed"); - - // Set default if provided - if let Some(s_name) = suggested_name { - name_prompt = name_prompt.with_default(s_name); - } - - let name = name_prompt.with_validator(move |input: &str| { - if input.trim().is_empty() { - return Ok(Validation::Invalid("Name cannot be empty".into())); - } - if name_regex.is_match(input) { - Ok(Validation::Valid) - } else { - Ok(Validation::Invalid( - "Name must contain only alphanumeric characters, dash (-) or underscore (_)" - .into(), - )) - } - }) - .prompt()?; - - // Collect host - let host = Text::new("Enter the PostgreSQL host:") - .with_help_message("Example: localhost or db.example.com") - .with_validator(|input: &str| { - if input.trim().is_empty() { - return Ok(Validation::Invalid("Host cannot be empty".into())); - } - Ok(Validation::Valid) - }) - .prompt()?; - - // Collect port (with validation) - let port_str = Text::new("Enter the PostgreSQL port:") - .with_default("5432") // Default Postgres port is 5432 - .with_help_message("Default PostgreSQL port is 5432") - .with_validator(|input: &str| match input.parse::() { - Ok(_) => Ok(Validation::Valid), - Err(_) => Ok(Validation::Invalid( - "Port must be a valid number between 1 and 65535".into(), - )), - }) - .prompt()?; - let port = port_str.parse::()?; - - // Collect username - let username = Text::new("Enter the PostgreSQL username:") - .with_validator(|input: &str| { - if input.trim().is_empty() { - return Ok(Validation::Invalid("Username cannot be empty".into())); - } - Ok(Validation::Valid) - }) - .prompt()?; - - // Collect password (masked) - let password = Password::new("Enter the PostgreSQL password:") - .with_validator(|input: &str| { - if input.trim().is_empty() { - return Ok(Validation::Invalid("Password cannot be empty".into())); - } - Ok(Validation::Valid) - }) - .without_confirmation() - .prompt()?; - - // Collect database (required) - let database = Text::new("Enter the default PostgreSQL database name:") - .with_validator(|input: &str| { - if input.trim().is_empty() { - return Ok(Validation::Invalid("Database cannot be empty".into())); - } - Ok(Validation::Valid) - }) - .prompt()?; - - // Collect schema (required) - let schema = Text::new("Enter the default PostgreSQL schema (optional):") - .prompt()?; - let schema_opt = if schema.trim().is_empty() { - None - } else { - Some(schema.clone()) - }; + // Collect fields using helpers + let name = prompt_validated_name("Enter a unique name for this data source:", suggested_name)?; + let host = prompt_required_text( + "Enter the PostgreSQL host:", + Some("Example: localhost or db.example.com"), + )?; + let port = prompt_u16_with_default( + "Enter the PostgreSQL port:", + "5432", + Some("Default PostgreSQL port is 5432"), + )?; + let username = prompt_required_text("Enter the PostgreSQL username:", None)?; + let password = prompt_password("Enter the PostgreSQL password:")?; + let database = prompt_required_text("Enter the default PostgreSQL database name:", None)?; + let schema = prompt_required_text("Enter the default PostgreSQL schema:", None)?; // Show summary and confirm println!("\n{}", "Connection Summary:".bold()); @@ -486,9 +550,7 @@ async fn setup_postgres( println!("Username: {}", username.cyan()); println!("Password: {}", "********".cyan()); println!("Default Database: {}", database.cyan()); - if let Some(s) = &schema_opt { - println!("Default Schema: {}", s.cyan()); - } + println!("Default Schema: {}", schema.cyan()); let confirm = Confirm::new("Do you want to create this data source?") .with_default(true) @@ -499,14 +561,14 @@ async fn setup_postgres( return Ok(()); } - // Create API request + // Create credentials and request let postgres_creds = PostgresCredentials { host, port, username, password, default_database: database.clone(), - default_schema: schema_opt.clone(), + default_schema: Some(schema.clone()), jump_host: None, ssh_username: None, ssh_private_key: None, @@ -517,50 +579,19 @@ async fn setup_postgres( credential, }; - // Send to API with progress indicator - let spinner = ProgressBar::new_spinner(); - spinner.set_style( - ProgressStyle::default_spinner() - .tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ ") - .template("{spinner:.green} {msg}") - .unwrap(), - ); - spinner.set_message("Sending credentials to Buster API..."); - spinner.enable_steady_tick(Duration::from_millis(100)); - + // Send to API using helper let client = BusterClient::new(buster_url, buster_api_key)?; + create_data_source_with_progress(&client, request).await?; - match client.post_data_sources(request).await { - Ok(_) => { - spinner.finish_with_message( - "✓ Data source created successfully!" - .green() - .bold() - .to_string(), - ); - println!( - "\nData source '{}' is now available for use with Buster.", - name.cyan() - ); + // If successful, proceed to create config file + println!( + "\nData source '{}' is now available for use with Buster.", + name.cyan() + ); + create_buster_config_file(config_path, &name, &database, Some(&schema))?; + println!("You can now use this data source with other Buster commands."); - // Create buster.yml file - create_buster_config_file( - config_path, - &name, - &database, - schema_opt.as_deref(), - )?; - - println!("You can now use this data source with other Buster commands."); - Ok(()) - } - Err(e) => { - spinner.finish_with_message("✗ Failed to create data source".red().bold().to_string()); - println!("\nError: {}", e); - println!("Please check your credentials and try again."); - Err(anyhow::anyhow!("Failed to create data source: {}", e)) - } - } + Ok(()) } async fn setup_bigquery( @@ -571,53 +602,15 @@ async fn setup_bigquery( ) -> Result<()> { println!("{}", "Setting up BigQuery connection...".bold().green()); - // Collect name (with validation) - let name_regex = Regex::new(r"^[a-zA-Z0-9_-]+$")?; - let mut name_prompt = Text::new("Enter a unique name for this data source:") - .with_help_message("Only alphanumeric characters, dash (-) and underscore (_) allowed"); + // Collect fields using helpers + let name = prompt_validated_name("Enter a unique name for this data source:", suggested_name)?; + let project_id = prompt_required_text( + "Enter the default Google Cloud project ID:", + Some("Example: my-project-123456"), + )?; + let dataset_id = prompt_required_text("Enter the default BigQuery dataset ID:", None)?; - // Set default if provided - if let Some(s_name) = suggested_name { - name_prompt = name_prompt.with_default(s_name); - } - - let name = name_prompt.with_validator(move |input: &str| { - if input.trim().is_empty() { - return Ok(Validation::Invalid("Name cannot be empty".into())); - } - if name_regex.is_match(input) { - Ok(Validation::Valid) - } else { - Ok(Validation::Invalid( - "Name must contain only alphanumeric characters, dash (-) or underscore (_)" - .into(), - )) - } - }) - .prompt()?; - - // Collect project ID - let project_id = Text::new("Enter the default Google Cloud project ID:") - .with_help_message("Example: my-project-123456") - .with_validator(|input: &str| { - if input.trim().is_empty() { - return Ok(Validation::Invalid("Project ID cannot be empty".into())); - } - Ok(Validation::Valid) - }) - .prompt()?; - - // Collect dataset ID (required) - let dataset_id = Text::new("Enter the default BigQuery dataset ID:") - .with_validator(|input: &str| { - if input.trim().is_empty() { - return Ok(Validation::Invalid("Dataset ID cannot be empty".into())); - } - Ok(Validation::Valid) - }) - .prompt()?; - - // Collect credentials JSON + // Collect credentials JSON (specific to BigQuery) println!( "\n{}", "BigQuery requires a service account credentials JSON file.".bold() @@ -626,42 +619,43 @@ async fn setup_bigquery( "You can create one in the Google Cloud Console under IAM & Admin > Service Accounts." ); - let credentials_path = Text::new("Enter the path to your credentials JSON file:") + let credentials_path_str = Text::new("Enter the path to your credentials JSON file:") .with_help_message("Example: /path/to/credentials.json") .with_validator(|input: &str| { let path = Path::new(input); if !path.exists() { - return Ok(Validation::Invalid("File does not exist".into())); + Ok(Validation::Invalid("File does not exist".into())) + } else if !path.is_file() { + Ok(Validation::Invalid("Path is not a file".into())) + } else { + Ok(Validation::Valid) } - if !path.is_file() { - return Ok(Validation::Invalid("Path is not a file".into())); - } - Ok(Validation::Valid) }) .prompt()?; - // Read credentials file - let credentials_content = match fs::read_to_string(&credentials_path) { - Ok(content) => content, - Err(e) => { - return Err(anyhow::anyhow!("Failed to read credentials file: {}", e)); - } - }; - - // Parse JSON to ensure it's valid and convert to serde_json::Value - let credentials_json: serde_json::Value = match serde_json::from_str(&credentials_content) { - Ok(json) => json, - Err(e) => { - return Err(anyhow::anyhow!("Invalid JSON in credentials file: {}", e)); - } - }; + // Read and parse credentials file + let credentials_content = fs::read_to_string(&credentials_path_str).map_err(|e| { + anyhow::anyhow!( + "Failed to read credentials file '{}': {}", + credentials_path_str, + e + ) + })?; + let credentials_json: serde_json::Value = + serde_json::from_str(&credentials_content).map_err(|e| { + anyhow::anyhow!( + "Invalid JSON in credentials file '{}': {}", + credentials_path_str, + e + ) + })?; // Show summary and confirm println!("\n{}", "Connection Summary:".bold()); println!("Name: {}", name.cyan()); println!("Default Project ID: {}", project_id.cyan()); println!("Default Dataset ID: {}", dataset_id.cyan()); - println!("Credentials: {}", credentials_path.cyan()); + println!("Credentials: {}", credentials_path_str.cyan()); let confirm = Confirm::new("Do you want to create this data source?") .with_default(true) @@ -672,11 +666,11 @@ async fn setup_bigquery( return Ok(()); } - // Create API request + // Create credentials and request let bigquery_creds = BigqueryCredentials { default_project_id: project_id.clone(), default_dataset_id: dataset_id.clone(), - credentials_json: credentials_json, + credentials_json, }; let credential = Credential::Bigquery(bigquery_creds); let request = PostDataSourcesRequest { @@ -684,50 +678,20 @@ async fn setup_bigquery( credential, }; - // Send to API with progress indicator - let spinner = ProgressBar::new_spinner(); - spinner.set_style( - ProgressStyle::default_spinner() - .tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ ") - .template("{spinner:.green} {msg}") - .unwrap(), - ); - spinner.set_message("Sending credentials to Buster API..."); - spinner.enable_steady_tick(Duration::from_millis(100)); - + // Send to API using helper let client = BusterClient::new(buster_url, buster_api_key)?; + create_data_source_with_progress(&client, request).await?; - match client.post_data_sources(request).await { - Ok(_) => { - spinner.finish_with_message( - "✓ Data source created successfully!" - .green() - .bold() - .to_string(), - ); - println!( - "\nData source '{}' is now available for use with Buster.", - name.cyan() - ); + // If successful, proceed to create config file + println!( + "\nData source '{}' is now available for use with Buster.", + name.cyan() + ); + // Map project_id to database, dataset_id to schema for config + create_buster_config_file(config_path, &name, &project_id, Some(&dataset_id))?; + println!("You can now use this data source with other Buster commands."); - // Create buster.yml file - create_buster_config_file( - config_path, - &name, - &project_id, - Some(&dataset_id), - )?; - - println!("You can now use this data source with other Buster commands."); - Ok(()) - } - Err(e) => { - spinner.finish_with_message("✗ Failed to create data source".red().bold().to_string()); - println!("\nError: {}", e); - println!("Please check your credentials and try again."); - Err(anyhow::anyhow!("Failed to create data source: {}", e)) - } - } + Ok(()) } async fn setup_mysql( @@ -736,69 +700,26 @@ async fn setup_mysql( config_path: &Path, suggested_name: Option<&str>, ) -> Result<()> { - println!("{}", "Setting up MySQL/MariaDB connection...".bold().green()); + println!( + "{}", + "Setting up MySQL/MariaDB connection...".bold().green() + ); - // Collect name - let name_regex = Regex::new(r"^[a-zA-Z0-9_-]+$")?; - let mut name_prompt = Text::new("Enter a unique name for this data source:") - .with_help_message("Only alphanumeric characters, dash (-) and underscore (_) allowed"); - - // Set default if provided - if let Some(s_name) = suggested_name { - name_prompt = name_prompt.with_default(s_name); - } - - let name = name_prompt.with_validator(move |input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Name cannot be empty".into())) } - else if name_regex.is_match(input) { Ok(Validation::Valid) } - else { Ok(Validation::Invalid("Name must contain only alphanumeric characters, dash (-) or underscore (_)".into())) } - }) - .prompt()?; - - // Collect host - let host = Text::new("Enter the MySQL/MariaDB host:") - .with_help_message("Example: localhost or db.example.com") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Host cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect port - let port_str = Text::new("Enter the MySQL/MariaDB port:") - .with_default("3306") - .with_help_message("Default MySQL/MariaDB port is 3306") - .with_validator(|input: &str| match input.parse::() { - Ok(_) => Ok(Validation::Valid), - Err(_) => Ok(Validation::Invalid("Port must be a valid number between 1 and 65535".into())), - }) - .prompt()?; - let port = port_str.parse::()?; - - // Collect username - let username = Text::new("Enter the MySQL/MariaDB username:") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Username cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect password - let password = Password::new("Enter the MySQL/MariaDB password:") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Password cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .without_confirmation() - .prompt()?; - - // Collect database (required) - let database = Text::new("Enter the default MySQL/MariaDB database name:") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Database cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; + // Collect fields using helpers + let name = prompt_validated_name("Enter a unique name for this data source:", suggested_name)?; + let host = prompt_required_text( + "Enter the MySQL/MariaDB host:", + Some("Example: localhost or db.example.com"), + )?; + let port = prompt_u16_with_default( + "Enter the MySQL/MariaDB port:", + "3306", + Some("Default MySQL/MariaDB port is 3306"), + )?; + let username = prompt_required_text("Enter the MySQL/MariaDB username:", None)?; + let password = prompt_password("Enter the MySQL/MariaDB password:")?; + let database = prompt_required_text("Enter the default MySQL/MariaDB database name:", None)?; + // No schema for MySQL // Show summary and confirm println!("\n{}", "Connection Summary:".bold()); @@ -818,41 +739,36 @@ async fn setup_mysql( return Ok(()); } - // Create API request + // Create credentials and request let mysql_creds = MySqlCredentials { host, port, username, password, default_database: database.clone(), - jump_host: None, // Not prompted for simplicity + jump_host: None, ssh_username: None, ssh_private_key: None, }; let credential = Credential::MySql(mysql_creds); - let request = PostDataSourcesRequest { name: name.clone(), credential }; + let request = PostDataSourcesRequest { + name: name.clone(), + credential, + }; - // Send to API - let spinner = ProgressBar::new_spinner(); - spinner.set_style(ProgressStyle::default_spinner().tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ ").template("{spinner:.green} {msg}").unwrap()); - spinner.set_message("Sending credentials to Buster API..."); - spinner.enable_steady_tick(Duration::from_millis(100)); + // Send to API using helper let client = BusterClient::new(buster_url, buster_api_key)?; + create_data_source_with_progress(&client, request).await?; - match client.post_data_sources(request).await { - Ok(_) => { - spinner.finish_with_message("✓ Data source created successfully!".green().bold().to_string()); - println!("\nData source '{}' is now available for use with Buster.", name.cyan()); - create_buster_config_file(config_path, &name, &database, None)?; // MySQL doesn't have a top-level schema concept like others - println!("You can now use this data source with other Buster commands."); - Ok(()) - } - Err(e) => { - spinner.finish_with_message("✗ Failed to create data source".red().bold().to_string()); - println!("\nError: {}", e); - Err(anyhow::anyhow!("Failed to create data source: {}", e)) - } - } + // If successful, proceed to create config file + println!( + "\nData source '{}' is now available for use with Buster.", + name.cyan() + ); + create_buster_config_file(config_path, &name, &database, None)?; + println!("You can now use this data source with other Buster commands."); + + Ok(()) } async fn setup_sqlserver( @@ -863,77 +779,21 @@ async fn setup_sqlserver( ) -> Result<()> { println!("{}", "Setting up SQL Server connection...".bold().green()); - // Collect name - let name_regex = Regex::new(r"^[a-zA-Z0-9_-]+$")?; - let mut name_prompt = Text::new("Enter a unique name for this data source:") - .with_help_message("Only alphanumeric characters, dash (-) and underscore (_) allowed"); - - // Set default if provided - if let Some(s_name) = suggested_name { - name_prompt = name_prompt.with_default(s_name); - } - - let name = name_prompt.with_validator(move |input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Name cannot be empty".into())) } - else if name_regex.is_match(input) { Ok(Validation::Valid) } - else { Ok(Validation::Invalid("Name must contain only alphanumeric characters, dash (-) or underscore (_)".into())) } - }) - .prompt()?; - - // Collect host - let host = Text::new("Enter the SQL Server host:") - .with_help_message("Example: server.database.windows.net or localhost") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Host cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect port - let port_str = Text::new("Enter the SQL Server port:") - .with_default("1433") - .with_help_message("Default SQL Server port is 1433") - .with_validator(|input: &str| match input.parse::() { - Ok(_) => Ok(Validation::Valid), - Err(_) => Ok(Validation::Invalid("Port must be a valid number between 1 and 65535".into())), - }) - .prompt()?; - let port = port_str.parse::()?; - - // Collect username - let username = Text::new("Enter the SQL Server username:") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Username cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect password - let password = Password::new("Enter the SQL Server password:") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Password cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .without_confirmation() - .prompt()?; - - // Collect database (required) - let database = Text::new("Enter the default SQL Server database name:") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Database cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect schema (optional) - let schema = Text::new("Enter the default SQL Server schema (optional, default: dbo):") - // No validator needed for optional field - .prompt()?; - let schema_opt = if schema.trim().is_empty() { - None // Or perhaps Some("dbo".to_string()) depending on desired default behavior - } else { - Some(schema.clone()) - }; + // Collect fields using helpers + let name = prompt_validated_name("Enter a unique name for this data source:", suggested_name)?; + let host = prompt_required_text( + "Enter the SQL Server host:", + Some("Example: server.database.windows.net or localhost"), + )?; + let port = prompt_u16_with_default( + "Enter the SQL Server port:", + "1433", + Some("Default SQL Server port is 1433"), + )?; + let username = prompt_required_text("Enter the SQL Server username:", None)?; + let password = prompt_password("Enter the SQL Server password:")?; + let database = prompt_required_text("Enter the default SQL Server database name:", None)?; + let schema = prompt_required_text("Enter the default SQL Server schema:", None)?; // Show summary and confirm println!("\n{}", "Connection Summary:".bold()); @@ -943,9 +803,7 @@ async fn setup_sqlserver( println!("Username: {}", username.cyan()); println!("Password: {}", "********".cyan()); println!("Default Database: {}", database.cyan()); - if let Some(s) = &schema_opt { - println!("Default Schema: {}", s.cyan()); - } + println!("Default Schema: {}", schema.cyan()); let confirm = Confirm::new("Do you want to create this data source?") .with_default(true) @@ -956,42 +814,37 @@ async fn setup_sqlserver( return Ok(()); } - // Create API request + // Create credentials and request let sqlserver_creds = SqlServerCredentials { host, port, username, password, default_database: database.clone(), - default_schema: schema_opt.clone(), - jump_host: None, // Not prompted for simplicity + default_schema: Some(schema.clone()), + jump_host: None, ssh_username: None, ssh_private_key: None, }; let credential = Credential::SqlServer(sqlserver_creds); - let request = PostDataSourcesRequest { name: name.clone(), credential }; + let request = PostDataSourcesRequest { + name: name.clone(), + credential, + }; - // Send to API - let spinner = ProgressBar::new_spinner(); - spinner.set_style(ProgressStyle::default_spinner().tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ ").template("{spinner:.green} {msg}").unwrap()); - spinner.set_message("Sending credentials to Buster API..."); - spinner.enable_steady_tick(Duration::from_millis(100)); + // Send to API using helper let client = BusterClient::new(buster_url, buster_api_key)?; + create_data_source_with_progress(&client, request).await?; - match client.post_data_sources(request).await { - Ok(_) => { - spinner.finish_with_message("✓ Data source created successfully!".green().bold().to_string()); - println!("\nData source '{}' is now available for use with Buster.", name.cyan()); - create_buster_config_file(config_path, &name, &database, schema_opt.as_deref())?; - println!("You can now use this data source with other Buster commands."); - Ok(()) - } - Err(e) => { - spinner.finish_with_message("✗ Failed to create data source".red().bold().to_string()); - println!("\nError: {}", e); - Err(anyhow::anyhow!("Failed to create data source: {}", e)) - } - } + // If successful, proceed to create config file + println!( + "\nData source '{}' is now available for use with Buster.", + name.cyan() + ); + create_buster_config_file(config_path, &name, &database, Some(&schema))?; + println!("You can now use this data source with other Buster commands."); + + Ok(()) } async fn setup_databricks( @@ -1002,62 +855,20 @@ async fn setup_databricks( ) -> Result<()> { println!("{}", "Setting up Databricks connection...".bold().green()); - // Collect name - let name_regex = Regex::new(r"^[a-zA-Z0-9_-]+$")?; - let mut name_prompt = Text::new("Enter a unique name for this data source:") - .with_help_message("Only alphanumeric characters, dash (-) and underscore (_) allowed"); - - // Set default if provided - if let Some(s_name) = suggested_name { - name_prompt = name_prompt.with_default(s_name); - } - - let name = name_prompt.with_validator(move |input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Name cannot be empty".into())) } - else if name_regex.is_match(input) { Ok(Validation::Valid) } - else { Ok(Validation::Invalid("Name must contain only alphanumeric characters, dash (-) or underscore (_)".into())) } - }) - .prompt()?; - - // Collect host - let host = Text::new("Enter the Databricks host:") - .with_help_message("Example: adb-xxxxxxxxxxxx.xx.azuredatabricks.net") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Host cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect API key - let api_key = Password::new("Enter the Databricks API key (Personal Access Token):") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("API key cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .without_confirmation() - .prompt()?; - - // Collect Warehouse ID - let warehouse_id = Text::new("Enter the Databricks SQL Warehouse HTTP Path:") - .with_help_message("Example: /sql/1.0/warehouses/xxxxxxxxxxxx") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Warehouse ID cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect default catalog (required) - let catalog = Text::new("Enter the default Databricks catalog:") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Catalog cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect default schema (optional) - let schema = Text::new("Enter the default Databricks schema (optional):") - .prompt()?; - let schema_opt = if schema.trim().is_empty() { None } else { Some(schema) }; + // Collect fields using helpers + let name = prompt_validated_name("Enter a unique name for this data source:", suggested_name)?; + let host = prompt_required_text( + "Enter the Databricks host:", + Some("Example: adb-xxxxxxxxxxxx.xx.azuredatabricks.net"), + )?; + // Databricks uses API key and warehouse ID instead of user/pass/port + let api_key = prompt_password("Enter the Databricks API key (Personal Access Token):")?; + let warehouse_id = prompt_required_text( + "Enter the Databricks SQL Warehouse HTTP Path:", + Some("Example: /sql/1.0/warehouses/xxxxxxxxxxxx"), + )?; + let catalog = prompt_required_text("Enter the default Databricks catalog:", None)?; + let schema = prompt_required_text("Enter the default Databricks schema:", None)?; // Show summary and confirm println!("\n{}", "Connection Summary:".bold()); @@ -1066,9 +877,7 @@ async fn setup_databricks( println!("API Key: {}", "********".cyan()); println!("Warehouse ID: {}", warehouse_id.cyan()); println!("Default Catalog: {}", catalog.cyan()); - if let Some(s) = &schema_opt { - println!("Default Schema: {}", s.cyan()); - } + println!("Default Schema: {}", schema.cyan()); let confirm = Confirm::new("Do you want to create this data source?") .with_default(true) @@ -1079,39 +888,34 @@ async fn setup_databricks( return Ok(()); } - // Create API request + // Create credentials and request let databricks_creds = DatabricksCredentials { host, api_key, warehouse_id, default_catalog: catalog.clone(), - default_schema: schema_opt.clone(), + default_schema: Some(schema.clone()), }; let credential = Credential::Databricks(databricks_creds); - let request = PostDataSourcesRequest { name: name.clone(), credential }; + let request = PostDataSourcesRequest { + name: name.clone(), + credential, + }; - // Send to API - let spinner = ProgressBar::new_spinner(); - spinner.set_style(ProgressStyle::default_spinner().tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ ").template("{spinner:.green} {msg}").unwrap()); - spinner.set_message("Sending credentials to Buster API..."); - spinner.enable_steady_tick(Duration::from_millis(100)); + // Send to API using helper let client = BusterClient::new(buster_url, buster_api_key)?; + create_data_source_with_progress(&client, request).await?; - match client.post_data_sources(request).await { - Ok(_) => { - spinner.finish_with_message("✓ Data source created successfully!".green().bold().to_string()); - println!("\nData source '{}' is now available for use with Buster.", name.cyan()); - // Map catalog to database, schema to schema for buster.yml - create_buster_config_file(config_path, &name, &catalog, schema_opt.as_deref())?; - println!("You can now use this data source with other Buster commands."); - Ok(()) - } - Err(e) => { - spinner.finish_with_message("✗ Failed to create data source".red().bold().to_string()); - println!("\nError: {}", e); - Err(anyhow::anyhow!("Failed to create data source: {}", e)) - } - } + // If successful, proceed to create config file + println!( + "\nData source '{}' is now available for use with Buster.", + name.cyan() + ); + // Map catalog to database, schema to schema for buster.yml + create_buster_config_file(config_path, &name, &catalog, Some(&schema))?; + println!("You can now use this data source with other Buster commands."); + + Ok(()) } async fn setup_snowflake( @@ -1122,76 +926,27 @@ async fn setup_snowflake( ) -> Result<()> { println!("{}", "Setting up Snowflake connection...".bold().green()); - // Collect name - let name_regex = Regex::new(r"^[a-zA-Z0-9_-]+$")?; - let mut name_prompt = Text::new("Enter a unique name for this data source:") - .with_help_message("Only alphanumeric characters, dash (-) and underscore (_) allowed"); - - // Set default if provided - if let Some(s_name) = suggested_name { - name_prompt = name_prompt.with_default(s_name); - } - - let name = name_prompt.with_validator(move |input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Name cannot be empty".into())) } - else if name_regex.is_match(input) { Ok(Validation::Valid) } - else { Ok(Validation::Invalid("Name must contain only alphanumeric characters, dash (-) or underscore (_)".into())) } - }) - .prompt()?; - - // Collect account ID - let account_id = Text::new("Enter the Snowflake account identifier:") - .with_help_message("Example: xy12345.us-east-1") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Account identifier cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect Warehouse ID - let warehouse_id = Text::new("Enter the Snowflake warehouse name:") - .with_help_message("Example: COMPUTE_WH") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Warehouse name cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect username - let username = Text::new("Enter the Snowflake username:") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Username cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect password - let password = Password::new("Enter the Snowflake password:") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Password cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .without_confirmation() - .prompt()?; - - // Collect role (optional) - let role = Text::new("Enter the Snowflake role (optional):") - .prompt()?; - let role_opt = if role.trim().is_empty() { None } else { Some(role) }; - - // Collect database (required) - let database = Text::new("Enter the default Snowflake database name:") - .with_validator(|input: &str| { - if input.trim().is_empty() { Ok(Validation::Invalid("Database cannot be empty".into())) } - else { Ok(Validation::Valid) } - }) - .prompt()?; - - // Collect schema (optional) - let schema = Text::new("Enter the default Snowflake schema (optional):") - .prompt()?; - let schema_opt = if schema.trim().is_empty() { None } else { Some(schema) }; - + // Collect fields using helpers + let name = prompt_validated_name("Enter a unique name for this data source:", suggested_name)?; + let account_id = prompt_required_text( + "Enter the Snowflake account identifier:", + Some("Example: xy12345.us-east-1"), + )?; + let warehouse_id = prompt_required_text( + "Enter the Snowflake warehouse name:", + Some("Example: COMPUTE_WH"), + )?; + let username = prompt_required_text("Enter the Snowflake username:", None)?; + let password = prompt_password("Enter the Snowflake password:")?; + // Role is optional for Snowflake - use standard Text prompt + let role = Text::new("Enter the Snowflake role (optional):").prompt()?; + let role_opt = if role.trim().is_empty() { + None + } else { + Some(role) + }; + let database = prompt_required_text("Enter the default Snowflake database name:", None)?; + let schema = prompt_required_text("Enter the default Snowflake schema:", None)?; // Show summary and confirm println!("\n{}", "Connection Summary:".bold()); @@ -1204,9 +959,7 @@ async fn setup_snowflake( println!("Role: {}", r.cyan()); } println!("Default Database: {}", database.cyan()); - if let Some(s) = &schema_opt { - println!("Default Schema: {}", s.cyan()); - } + println!("Default Schema: {}", schema.cyan()); let confirm = Confirm::new("Do you want to create this data source?") .with_default(true) @@ -1217,7 +970,7 @@ async fn setup_snowflake( return Ok(()); } - // Create API request + // Create credentials and request let snowflake_creds = SnowflakeCredentials { account_id, warehouse_id, @@ -1225,92 +978,99 @@ async fn setup_snowflake( password, role: role_opt.clone(), default_database: database.clone(), - default_schema: schema_opt.clone(), + default_schema: Some(schema.clone()), }; let credential = Credential::Snowflake(snowflake_creds); - let request = PostDataSourcesRequest { name: name.clone(), credential }; + let request = PostDataSourcesRequest { + name: name.clone(), + credential, + }; - // Send to API - let spinner = ProgressBar::new_spinner(); - spinner.set_style(ProgressStyle::default_spinner().tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ ").template("{spinner:.green} {msg}").unwrap()); - spinner.set_message("Sending credentials to Buster API..."); - spinner.enable_steady_tick(Duration::from_millis(100)); + // Send to API using helper let client = BusterClient::new(buster_url, buster_api_key)?; + create_data_source_with_progress(&client, request).await?; - match client.post_data_sources(request).await { - Ok(_) => { - spinner.finish_with_message("✓ Data source created successfully!".green().bold().to_string()); - println!("\nData source '{}' is now available for use with Buster.", name.cyan()); - create_buster_config_file(config_path, &name, &database, schema_opt.as_deref())?; - println!("You can now use this data source with other Buster commands."); - Ok(()) - } - Err(e) => { - spinner.finish_with_message("✗ Failed to create data source".red().bold().to_string()); - println!("\nError: {}", e); - Err(anyhow::anyhow!("Failed to create data source: {}", e)) + // If successful, proceed to create config file + println!( + "\nData source '{}' is now available for use with Buster.", + name.cyan() + ); + create_buster_config_file(config_path, &name, &database, Some(&schema))?; + println!("You can now use this data source with other Buster commands."); + + Ok(()) +} + +// --- Helper to suggest model paths from dbt_project.yml --- +fn suggest_model_paths_from_dbt(buster_config_dir: &Path) -> (Option, String) { + let mut suggested_model_paths_str = "".to_string(); + let mut log_message = "".to_string(); + + let dbt_project_path = buster_config_dir.join("dbt_project.yml"); + if dbt_project_path.exists() && dbt_project_path.is_file() { + match fs::read_to_string(&dbt_project_path) { + Ok(content) => match serde_yaml::from_str::(&content) { + Ok(dbt_config) => { + let paths_to_suggest = dbt_config + .model_paths + .unwrap_or_else(|| vec!["models".to_string()]); + if !paths_to_suggest.is_empty() { + suggested_model_paths_str = paths_to_suggest.join(","); + log_message = format!( + "Found dbt_project.yml, suggesting model paths: {}", + suggested_model_paths_str.cyan() + ); + } + } + Err(e) => { + log_message = format!( + "Warning: Failed to parse {}: {}. Proceeding without suggested model paths.", + dbt_project_path.display(), e + ); + } + }, + Err(e) => { + log_message = format!( + "Warning: Failed to read {}: {}. Proceeding without suggested model paths.", + dbt_project_path.display(), + e + ); + } } } + (Some(suggested_model_paths_str), log_message) } // Helper function to create buster.yml file fn create_buster_config_file( path: &Path, data_source_name: &str, - database: &str, - schema: Option<&str>, + database: &str, // This represents 'database' for most, 'project_id' for BQ, 'catalog' for Databricks + schema: Option<&str>, // Made optional again at signature level ) -> Result<()> { - // --- BEGIN DBT PROJECT DETECTION --- - let mut suggested_model_paths = "".to_string(); // Default to empty string + // --- Suggest model paths based on dbt_project.yml --- + let (suggested_paths_opt, suggestion_log) = if let Some(parent_dir) = path.parent() { + suggest_model_paths_from_dbt(parent_dir) + } else { + (None, "".to_string()) // Cannot determine parent dir + }; - // Construct path to dbt_project.yml relative to the buster.yml path - if let Some(parent_dir) = path.parent() { - let dbt_project_path = parent_dir.join("dbt_project.yml"); - if dbt_project_path.exists() && dbt_project_path.is_file() { - match fs::read_to_string(&dbt_project_path) { - Ok(content) => { - match serde_yaml::from_str::(&content) { - Ok(dbt_config) => { - // Use specified model-paths or default ["models"] if present in file - let paths_to_suggest = dbt_config.model_paths.unwrap_or_else(|| vec!["models".to_string()]); - if !paths_to_suggest.is_empty() { - suggested_model_paths = paths_to_suggest.join(","); - println!( - "{}", - format!("Found dbt_project.yml, suggesting model paths: {}", suggested_model_paths.cyan()).dimmed() - ); - } - }, - Err(e) => { - // Log error but don't fail the init process - eprintln!( - "{}", - format!("Warning: Failed to parse {}: {}. Proceeding without suggested model paths.", dbt_project_path.display(), e).yellow() - ); - } - } - }, - Err(e) => { - eprintln!( - "{}", - format!("Warning: Failed to read {}: {}. Proceeding without suggested model paths.", dbt_project_path.display(), e).yellow() - ); - } - } + if !suggestion_log.is_empty() { + if suggestion_log.starts_with("Warning") { + eprintln!("{}", suggestion_log.yellow()); + } else { + println!("{}", suggestion_log.dimmed()); } } - // --- END DBT PROJECT DETECTION --- - // Prompt for model paths (optional), now with potential initial input - let model_paths_input = Text::new( - "Enter paths to your SQL models (optional, comma-separated):", - ) - .with_default(&suggested_model_paths) // Use with_default instead - .with_help_message( - "Leave blank if none, or specify paths like './models,./analytics/models'", - ) - .prompt()?; + let model_paths_input = + Text::new("Enter paths to your SQL models (optional, comma-separated):") + .with_default(suggested_paths_opt.as_deref().unwrap_or("")) // Use with_default + .with_help_message( + "Leave blank if none, or specify paths like './models,./analytics/models'", + ) + .prompt()?; // Process the comma-separated input into a vector if not empty let model_paths_vec = if model_paths_input.trim().is_empty() { @@ -1340,15 +1100,15 @@ fn create_buster_config_file( // --- Top-level fields are None when generating new config --- data_source_name: None, schema: None, - database: None, + database: None, exclude_files: None, // Define top-level excludes if needed exclude_tags: None, model_paths: None, - // --- Populate the projects field --- + // --- Populate the projects field --- projects: Some(vec![main_context]), }; - let yaml = serde_yaml::to_string(&config)?; + let yaml = serde_yaml::to_string(&config)?; fs::write(path, yaml)?; println!(