mirror of https://github.com/buster-so/buster.git
feat: add support for custom environment variables in CLI commands
- Introduced `env_vars` option in the `start` command to allow users to set environment variables in KEY=VALUE format. - Implemented `update_arbitrary_env_vars` function to update the .env file with provided variables. - Adjusted `setup_persistent_app_environment` to handle custom environment variables. - Added `MAX_RECURSION` variable to .env.example and updated recursion limit logic in the agent implementation.
This commit is contained in:
parent
86c275e70d
commit
e9a269983b
|
@ -11,6 +11,7 @@ SUPABASE_URL="http://kong:8000"
|
|||
SUPABASE_SERVICE_ROLE_KEY="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.ey AgCiAgICAicm9sZSI6ICJzZXJ2aWNlX3JvbGUiLAogICAgImlzcyI6ICJzdXBhYmFzZS1kZW1vIiwKICAgICJpYXQiOiAxNjQxNzY5MjAwLAogICAgImV4cCI6IDE3OTk1MzU2MDAKfQ.DaYlNEoUrrEn2Ig7tqibS-PHK5vgusbcbo7X36XVt4Q"
|
||||
POSTHOG_TELEMETRY_KEY="phc_zZraCicSTfeXX5b9wWQv2rWG8QB4Z3xlotOT7gFtoNi"
|
||||
TELEMETRY_ENABLED="true"
|
||||
MAX_RECURSION="15"
|
||||
|
||||
# AI VARS
|
||||
RERANK_API_KEY="your_rerank_api_key"
|
||||
|
|
|
@ -554,8 +554,13 @@ impl Agent {
|
|||
*current = Some(thread_ref.clone());
|
||||
}
|
||||
|
||||
let max_recursion = std::env::var("MAX_RECURSION")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u32>().ok())
|
||||
.unwrap_or(15);
|
||||
|
||||
// Limit recursion to a maximum of 15 times
|
||||
if recursion_depth >= 15 {
|
||||
if recursion_depth >= max_recursion {
|
||||
let max_depth_msg = format!("Maximum recursion depth ({}) reached.", recursion_depth);
|
||||
warn!("{}", max_depth_msg);
|
||||
let message = AgentMessage::assistant(
|
||||
|
|
|
@ -2,9 +2,11 @@ use crate::error::BusterError;
|
|||
use dirs;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_yaml;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::fs;
|
||||
use std::io::{self, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::str::FromStr;
|
||||
use colored::*;
|
||||
use inquire::{Confirm, Password, Select, Text, PasswordDisplayMode, validator::Validation};
|
||||
|
||||
|
@ -179,6 +181,60 @@ pub fn update_env_file(
|
|||
})
|
||||
}
|
||||
|
||||
pub fn update_arbitrary_env_vars(
|
||||
target_dotenv_path: &Path,
|
||||
env_vars: &[(String, String)],
|
||||
) -> Result<(), BusterError> {
|
||||
let mut new_env_lines: Vec<String> = Vec::new();
|
||||
let mut updated_vars: HashSet<String> = HashSet::new();
|
||||
|
||||
// Read existing .env file if it exists
|
||||
if target_dotenv_path.exists() {
|
||||
let env_content = fs::read_to_string(target_dotenv_path).map_err(|e| {
|
||||
BusterError::CommandError(format!(
|
||||
"Failed to read .env file at {}: {}",
|
||||
target_dotenv_path.display(),
|
||||
e
|
||||
))
|
||||
})?;
|
||||
|
||||
for line in env_content.lines() {
|
||||
let mut line_replaced = false;
|
||||
|
||||
// Check if this line starts with any of our env vars
|
||||
for (key, value) in env_vars {
|
||||
if line.starts_with(&format!("{}=", key)) {
|
||||
new_env_lines.push(format!("{}=\"{}\"", key, value));
|
||||
updated_vars.insert(key.clone());
|
||||
line_replaced = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// If no replacement was made, keep the original line
|
||||
if !line_replaced {
|
||||
new_env_lines.push(line.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Add any environment variables that weren't found in the existing file
|
||||
for (key, value) in env_vars {
|
||||
if !updated_vars.contains(key) {
|
||||
new_env_lines.push(format!("{}=\"{}\"", key, value));
|
||||
}
|
||||
}
|
||||
|
||||
// Write the updated content back to the file
|
||||
fs::write(target_dotenv_path, new_env_lines.join("\n")).map_err(|e| {
|
||||
BusterError::CommandError(format!(
|
||||
"Failed to write updated .env file to {}: {}",
|
||||
target_dotenv_path.display(),
|
||||
e
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, Clone, Default)]
|
||||
pub struct ModelInfo {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
|
|
|
@ -23,7 +23,7 @@ use colored::*;
|
|||
#[exclude = "supabase/docker-compose.override.yml"]
|
||||
struct StaticAssets;
|
||||
|
||||
async fn setup_persistent_app_environment() -> Result<PathBuf, BusterError> {
|
||||
async fn setup_persistent_app_environment(env_vars: Option<&[(String, String)]>) -> Result<PathBuf, BusterError> {
|
||||
let app_base_dir = config_utils::get_app_base_dir().map_err(|e| {
|
||||
BusterError::CommandError(format!("Failed to get app base directory: {}", e))
|
||||
})?;
|
||||
|
@ -172,6 +172,25 @@ async fn setup_persistent_app_environment() -> Result<PathBuf, BusterError> {
|
|||
))
|
||||
})?;
|
||||
|
||||
// Update arbitrary environment variables if provided
|
||||
if let Some(env_vars) = env_vars {
|
||||
if !env_vars.is_empty() {
|
||||
println!("Updating custom environment variables...");
|
||||
config_utils::update_arbitrary_env_vars(&main_dot_env_path, env_vars)
|
||||
.map_err(|e| {
|
||||
BusterError::CommandError(format!(
|
||||
"Failed to update custom environment variables in {}: {}",
|
||||
main_dot_env_path.display(),
|
||||
e
|
||||
))
|
||||
})?;
|
||||
|
||||
for (key, value) in env_vars {
|
||||
println!("Set {}={}", key, if key.to_lowercase().contains("key") || key.to_lowercase().contains("password") || key.to_lowercase().contains("secret") { "***" } else { value });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
println!("--- Configuration Setup/Check Complete ---");
|
||||
// --- END Configuration Checks/Updates ---
|
||||
|
||||
|
@ -371,9 +390,10 @@ async fn run_docker_compose_command(
|
|||
}
|
||||
}
|
||||
|
||||
pub async fn start(no_track: bool) -> Result<(), BusterError> {
|
||||
pub async fn start(no_track: bool, env_vars: Vec<(String, String)>) -> Result<(), BusterError> {
|
||||
// First, run the setup/check which includes printing headers/prompts if needed
|
||||
let app_base_dir = setup_persistent_app_environment().await?;
|
||||
let env_vars_option = if env_vars.is_empty() { None } else { Some(env_vars.as_slice()) };
|
||||
let app_base_dir = setup_persistent_app_environment(env_vars_option).await?;
|
||||
// Then, run the docker command in that directory
|
||||
run_docker_compose_command(&app_base_dir, &["up", "-d"], "Starting", no_track).await
|
||||
}
|
||||
|
|
|
@ -3,10 +3,10 @@ mod error;
|
|||
mod types;
|
||||
mod utils;
|
||||
|
||||
use anyhow;
|
||||
use clap::{Parser, Subcommand};
|
||||
use commands::{auth::check_authentication, auth::AuthArgs, init, run};
|
||||
use utils::updater::check_for_updates;
|
||||
use anyhow;
|
||||
|
||||
pub const APP_NAME: &str = "buster";
|
||||
pub const VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
|
@ -95,6 +95,10 @@ pub enum Commands {
|
|||
/// Disable telemetry tracking
|
||||
#[arg(long, default_value_t = false)]
|
||||
no_track: bool,
|
||||
/// Set environment variables (can be used multiple times)
|
||||
/// Format: KEY=VALUE
|
||||
#[arg(long = "env", action = clap::ArgAction::Append)]
|
||||
env_vars: Vec<String>,
|
||||
},
|
||||
/// Stop the Buster services
|
||||
Stop,
|
||||
|
@ -161,7 +165,22 @@ async fn main() {
|
|||
} => commands::generate::generate_semantic_models_command(path, target_semantic_file).await,
|
||||
Commands::Parse { path } => commands::parse::parse_models_command(path).await,
|
||||
Commands::Config => commands::config::manage_settings_interactive().await.map_err(anyhow::Error::from),
|
||||
Commands::Start { no_track } => run::start(no_track).await.map_err(anyhow::Error::from),
|
||||
Commands::Start { no_track, env_vars } => {
|
||||
// Parse env vars from KEY=VALUE format
|
||||
let parsed_env_vars: Result<Vec<(String, String)>, _> = env_vars
|
||||
.iter()
|
||||
.map(|env_str| {
|
||||
env_str.split_once('=')
|
||||
.map(|(k, v)| (k.to_string(), v.to_string()))
|
||||
.ok_or_else(|| anyhow::anyhow!("Invalid env var format '{}'. Expected KEY=VALUE", env_str))
|
||||
})
|
||||
.collect();
|
||||
|
||||
match parsed_env_vars {
|
||||
Ok(env_pairs) => run::start(no_track, env_pairs).await.map_err(anyhow::Error::from),
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
Commands::Stop => run::stop().await.map_err(anyhow::Error::from),
|
||||
Commands::Reset => run::reset().await.map_err(anyhow::Error::from),
|
||||
};
|
||||
|
|
|
@ -0,0 +1,159 @@
|
|||
use anyhow::Result;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use tempfile::TempDir;
|
||||
|
||||
// Import from the buster CLI library
|
||||
use buster_cli::commands::config_utils;
|
||||
|
||||
/// Test that update_arbitrary_env_vars correctly updates environment variables
|
||||
#[tokio::test]
|
||||
async fn test_update_arbitrary_env_vars_new_file() -> Result<()> {
|
||||
let temp_dir = TempDir::new()?;
|
||||
let env_file_path = temp_dir.path().join(".env");
|
||||
|
||||
// Test data
|
||||
let env_vars = vec![
|
||||
("DATABASE_URL".to_string(), "postgres://localhost:5432/testdb".to_string()),
|
||||
("API_KEY".to_string(), "secret123".to_string()),
|
||||
("DEBUG".to_string(), "true".to_string()),
|
||||
];
|
||||
|
||||
// Update env vars (should create new file)
|
||||
config_utils::update_arbitrary_env_vars(&env_file_path, &env_vars)?;
|
||||
|
||||
// Verify file was created and contains correct content
|
||||
assert!(env_file_path.exists());
|
||||
let content = fs::read_to_string(&env_file_path)?;
|
||||
|
||||
assert!(content.contains("DATABASE_URL=\"postgres://localhost:5432/testdb\""));
|
||||
assert!(content.contains("API_KEY=\"secret123\""));
|
||||
assert!(content.contains("DEBUG=\"true\""));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_update_arbitrary_env_vars_existing_file() -> Result<()> {
|
||||
let temp_dir = TempDir::new()?;
|
||||
let env_file_path = temp_dir.path().join(".env");
|
||||
|
||||
// Create initial .env file
|
||||
let initial_content = r#"
|
||||
# Existing configuration
|
||||
EXISTING_VAR="keep_this"
|
||||
OLD_API_KEY="old_secret"
|
||||
DATABASE_URL="postgres://localhost:5432/olddb"
|
||||
"#;
|
||||
fs::write(&env_file_path, initial_content)?;
|
||||
|
||||
// Test data - some new, some updating existing
|
||||
let env_vars = vec![
|
||||
("DATABASE_URL".to_string(), "postgres://localhost:5432/newdb".to_string()),
|
||||
("NEW_VAR".to_string(), "new_value".to_string()),
|
||||
("API_SECRET".to_string(), "super_secret".to_string()),
|
||||
];
|
||||
|
||||
// Update env vars
|
||||
config_utils::update_arbitrary_env_vars(&env_file_path, &env_vars)?;
|
||||
|
||||
// Verify content
|
||||
let content = fs::read_to_string(&env_file_path)?;
|
||||
|
||||
// Should preserve existing vars not being updated
|
||||
assert!(content.contains("EXISTING_VAR=\"keep_this\""));
|
||||
assert!(content.contains("OLD_API_KEY=\"old_secret\""));
|
||||
|
||||
// Should update existing var
|
||||
assert!(content.contains("DATABASE_URL=\"postgres://localhost:5432/newdb\""));
|
||||
assert!(!content.contains("postgres://localhost:5432/olddb"));
|
||||
|
||||
// Should add new vars
|
||||
assert!(content.contains("NEW_VAR=\"new_value\""));
|
||||
assert!(content.contains("API_SECRET=\"super_secret\""));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_update_arbitrary_env_vars_preserves_comments() -> Result<()> {
|
||||
let temp_dir = TempDir::new()?;
|
||||
let env_file_path = temp_dir.path().join(".env");
|
||||
|
||||
// Create initial .env file with comments
|
||||
let initial_content = r#"# Database configuration
|
||||
DATABASE_URL="postgres://localhost:5432/olddb"
|
||||
|
||||
# API configuration
|
||||
API_KEY="old_key"
|
||||
# This is a comment that should be preserved
|
||||
KEEP_THIS="unchanged"
|
||||
"#;
|
||||
fs::write(&env_file_path, initial_content)?;
|
||||
|
||||
// Update only one variable
|
||||
let env_vars = vec![
|
||||
("API_KEY".to_string(), "new_key".to_string()),
|
||||
];
|
||||
|
||||
config_utils::update_arbitrary_env_vars(&env_file_path, &env_vars)?;
|
||||
|
||||
let content = fs::read_to_string(&env_file_path)?;
|
||||
|
||||
// Should preserve comments and other variables
|
||||
assert!(content.contains("# Database configuration"));
|
||||
assert!(content.contains("# API configuration"));
|
||||
assert!(content.contains("# This is a comment that should be preserved"));
|
||||
assert!(content.contains("DATABASE_URL=\"postgres://localhost:5432/olddb\""));
|
||||
assert!(content.contains("KEEP_THIS=\"unchanged\""));
|
||||
|
||||
// Should update the target variable
|
||||
assert!(content.contains("API_KEY=\"new_key\""));
|
||||
assert!(!content.contains("API_KEY=\"old_key\""));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_update_arbitrary_env_vars_empty_list() -> Result<()> {
|
||||
let temp_dir = TempDir::new()?;
|
||||
let env_file_path = temp_dir.path().join(".env");
|
||||
|
||||
// Create initial .env file
|
||||
let initial_content = "EXISTING_VAR=\"value\"\n";
|
||||
fs::write(&env_file_path, initial_content)?;
|
||||
|
||||
// Update with empty list
|
||||
let env_vars = vec![];
|
||||
config_utils::update_arbitrary_env_vars(&env_file_path, &env_vars)?;
|
||||
|
||||
// Should preserve existing content
|
||||
let content = fs::read_to_string(&env_file_path)?;
|
||||
assert_eq!(content, initial_content);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_update_arbitrary_env_vars_handles_special_characters() -> Result<()> {
|
||||
let temp_dir = TempDir::new()?;
|
||||
let env_file_path = temp_dir.path().join(".env");
|
||||
|
||||
// Test data with special characters
|
||||
let env_vars = vec![
|
||||
("URL_WITH_PARAMS".to_string(), "http://localhost:3000/api?key=value&other=123".to_string()),
|
||||
("PASSWORD_WITH_SYMBOLS".to_string(), "p@ssw0rd!#$%^&*()".to_string()),
|
||||
("JSON_CONFIG".to_string(), "{\"key\":\"value\",\"number\":42}".to_string()),
|
||||
];
|
||||
|
||||
config_utils::update_arbitrary_env_vars(&env_file_path, &env_vars)?;
|
||||
|
||||
let content = fs::read_to_string(&env_file_path)?;
|
||||
|
||||
// Verify special characters are preserved
|
||||
assert!(content.contains("URL_WITH_PARAMS=\"http://localhost:3000/api?key=value&other=123\""));
|
||||
assert!(content.contains("PASSWORD_WITH_SYMBOLS=\"p@ssw0rd!#$%^&*()\""));
|
||||
assert!(content.contains("JSON_CONFIG=\"{\"key\":\"value\",\"number\":42}\""));
|
||||
|
||||
Ok(())
|
||||
}
|
Loading…
Reference in New Issue