From bf58045a2841db7c46b44d8c3dbccb28a161697e Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 12 May 2026 17:31:54 +0800 Subject: [PATCH] Revert "feat(cli): add dbx-cli runtime and agent integration (#229)" This reverts commit 58c47adef5a893219259f5835cfd367b345bdb64, reversing changes made to 3bcec8ef7f0ace0cdea0b629415f78af7c13f5d8. --- .gitignore | 1 - Cargo.lock | 16 - Cargo.toml | 2 +- crates/dbx-cli/Cargo.toml | 21 - crates/dbx-cli/src/commands.rs | 1617 ----------------- crates/dbx-cli/src/main.rs | 10 - crates/dbx-cli/src/runtime_client.rs | 190 -- crates/dbx-core/src/cli.rs | 345 ---- crates/dbx-core/src/db/mysql.rs | 27 +- crates/dbx-core/src/db/postgres.rs | 42 +- crates/dbx-core/src/db/sqlite.rs | 15 +- crates/dbx-core/src/handoff.rs | 69 - crates/dbx-core/src/lib.rs | 4 - crates/dbx-core/src/query.rs | 233 +-- crates/dbx-core/src/schema_snapshot.rs | 164 -- crates/dbx-core/src/sql_safety.rs | 490 ----- crates/dbx-core/src/storage.rs | 295 --- crates/dbx-core/tests/schema_snapshot.rs | 140 -- .../plans/2026-05-10-dbx-cli-runtime.md | 1537 ---------------- src-tauri/Cargo.toml | 1 - src-tauri/src/commands/agent_runtime.rs | 889 --------- src-tauri/src/commands/mod.rs | 1 - src-tauri/src/lib.rs | 14 +- src/App.vue | 28 +- src/components/agent/AgentHandoffDialog.vue | 78 - src/components/grid/DataGrid.vue | 21 +- src/components/layout/AppDialogs.vue | 2 - src/lib/agentHandoff.ts | 48 - src/lib/agentRuntimeSnapshot.ts | 62 - src/lib/api.ts | 6 - src/lib/appStartup.ts | 13 - src/lib/http.ts | 16 - src/lib/tauri.ts | 18 - src/stores/agentRuntimeStore.ts | 109 -- src/stores/queryStore.ts | 35 +- tests/agentHandoff.test.ts | 74 - tests/agentRuntimeSnapshot.test.ts | 102 -- 37 files changed, 56 insertions(+), 6679 deletions(-) delete mode 100644 crates/dbx-cli/Cargo.toml delete mode 100644 crates/dbx-cli/src/commands.rs delete mode 100644 crates/dbx-cli/src/main.rs delete mode 100644 crates/dbx-cli/src/runtime_client.rs delete mode 100644 crates/dbx-core/src/cli.rs delete mode 100644 crates/dbx-core/src/handoff.rs delete mode 100644 crates/dbx-core/src/schema_snapshot.rs delete mode 100644 crates/dbx-core/src/sql_safety.rs delete mode 100644 crates/dbx-core/tests/schema_snapshot.rs delete mode 100644 docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md delete mode 100644 src-tauri/src/commands/agent_runtime.rs delete mode 100644 src/components/agent/AgentHandoffDialog.vue delete mode 100644 src/lib/agentHandoff.ts delete mode 100644 src/lib/agentRuntimeSnapshot.ts delete mode 100644 src/lib/appStartup.ts delete mode 100644 src/stores/agentRuntimeStore.ts delete mode 100644 tests/agentHandoff.test.ts delete mode 100644 tests/agentRuntimeSnapshot.test.ts diff --git a/.gitignore b/.gitignore index ef8705469..0cd799edb 100644 --- a/.gitignore +++ b/.gitignore @@ -40,7 +40,6 @@ Thumbs.db # Temp tmp/ -.worktrees/ # Logs *.log diff --git a/Cargo.lock b/Cargo.lock index ea6031c52..c8f06c8e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1644,7 +1644,6 @@ dependencies = [ "dbx-core", "duckdb", "futures", - "libc", "log", "mongodb", "percent-encoding", @@ -1674,21 +1673,6 @@ dependencies = [ "zip 4.6.1", ] -[[package]] -name = "dbx-cli" -version = "0.5.2" -dependencies = [ - "chrono", - "dbx-core", - "libc", - "reqwest 0.12.28", - "serde", - "serde_json", - "tempfile", - "tokio", - "uuid", -] - [[package]] name = "dbx-core" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 1b5c694ec..390ef8ef7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [workspace] resolver = "2" -members = ["src-tauri", "crates/dbx-core", "crates/dbx-cli", "src-web"] +members = ["src-tauri", "crates/dbx-core", "src-web"] [profile.release] panic = "abort" diff --git a/crates/dbx-cli/Cargo.toml b/crates/dbx-cli/Cargo.toml deleted file mode 100644 index acb34da20..000000000 --- a/crates/dbx-cli/Cargo.toml +++ /dev/null @@ -1,21 +0,0 @@ -[package] -name = "dbx-cli" -version = "0.5.2" -edition = "2021" - -[[bin]] -name = "dbx-cli" -path = "src/main.rs" - -[dependencies] -dbx-core = { path = "../dbx-core" } -chrono = { version = "0.4", features = ["serde"] } -serde = { version = "1.0", features = ["derive"] } -serde_json = "1.0" -tokio = { version = "1", features = ["full"] } -reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } -uuid = { version = "1", features = ["v4", "serde"] } -libc = "0.2" - -[dev-dependencies] -tempfile = "3" diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs deleted file mode 100644 index 8ecfd6e85..000000000 --- a/crates/dbx-cli/src/commands.rs +++ /dev/null @@ -1,1617 +0,0 @@ -use dbx_core::cli::{fail, fail_safe, ok, CliEnvelope, CliErrorCode, CliSource}; - -const DEFAULT_RESULT_LIMIT: u32 = 50; -const MAX_RESULT_LIMIT: u32 = 1000; - -pub(crate) async fn run(args: Vec) -> Result<(), CliEnvelope<()>> { - let output = dispatch(args).await; - println!("{}", serde_json::to_string_pretty(&output).unwrap()); - - if matches!(output, CliEnvelope::Failure { .. }) { - std::process::exit(1); - } - - Ok(()) -} - -pub(crate) async fn dispatch(args: Vec) -> CliEnvelope { - let parsed = match parse_args(args) { - Ok(parsed) => parsed, - Err(err) => return err, - }; - - match parsed.positionals.as_slice() { - [cmd, rest @ ..] if cmd == "context" => context(rest).await, - [cmd, sub, rest @ ..] if cmd == "conn" && sub == "list" => conn_list(rest).await, - [cmd, sub, name, rest @ ..] if cmd == "conn" && sub == "show" => conn_show(name, rest).await, - [cmd, sub, rest @ ..] if cmd == "schema" && sub == "snapshot" => schema_snapshot(rest).await, - [cmd, rest @ ..] if cmd == "safe-query" => safe_query(rest).await, - [cmd, rest @ ..] if cmd == "handoff" => handoff(rest).await, - [cmd, rest @ ..] if cmd == "selection" => selection(rest).await, - [cmd, sub, rest @ ..] if cmd == "result" && sub == "current" => result_current(rest).await, - _ => fail(CliSource::Headless, CliErrorCode::InternalError, "Unknown command", false), - } -} - -struct ParsedArgs { - positionals: Vec, -} - -#[derive(Clone, Copy, PartialEq, Eq)] -enum FlagKind { - Value, - Bool, -} - -#[derive(Clone, Copy)] -struct FlagSpec { - kind: FlagKind, - allow_dash_value: bool, -} - -fn parse_args(args: Vec) -> Result> { - let mut positionals = Vec::new(); - let mut index = 0; - - while index < args.len() { - let arg = &args[index]; - if arg == "--format" { - let Some(value) = args.get(index + 1) else { - return Err(invalid_args(format!("{arg} requires a value"))); - }; - if value != "json" { - return Err(invalid_args("Only --format json is supported")); - } - index += 2; - } else if arg.starts_with("--") { - let Some(spec) = - flag_spec(positionals.first().map(String::as_str), positionals.get(1).map(String::as_str), arg) - else { - return Err(invalid_args(format!("Unknown flag: {arg}"))); - }; - - if spec.kind == FlagKind::Value { - let Some(value) = args.get(index + 1) else { - return Err(invalid_args(format!("{arg} requires a value"))); - }; - if !spec.allow_dash_value && value.starts_with("--") { - return Err(invalid_args(format!("{arg} requires a value"))); - } - positionals.push(arg.clone()); - positionals.push(value.clone()); - index += 2; - } else { - positionals.push(arg.clone()); - index += 1; - } - } else { - positionals.push(arg.clone()); - index += 1; - } - } - - Ok(ParsedArgs { positionals }) -} - -fn flag_spec(command: Option<&str>, subcommand: Option<&str>, flag: &str) -> Option { - let value = FlagKind::Value; - let boolean = FlagKind::Bool; - let normal_value = FlagSpec { kind: value, allow_dash_value: false }; - let free_text_value = FlagSpec { kind: value, allow_dash_value: true }; - let boolean_flag = FlagSpec { kind: boolean, allow_dash_value: false }; - - match (command, subcommand, flag) { - (Some("conn"), Some("show"), "--redacted") => Some(boolean_flag), - (Some("schema"), Some("snapshot"), "--conn" | "--db") => Some(normal_value), - (Some("safe-query"), _, "--conn" | "--db" | "--limit") => Some(normal_value), - (Some("safe-query"), _, "--sql") => Some(free_text_value), - (Some("handoff"), _, "--conn" | "--sql-file") => Some(normal_value), - (Some("handoff"), _, "--title" | "--sql" | "--description") => Some(free_text_value), - (Some("result"), Some("current"), "--limit") => Some(normal_value), - _ => None, - } -} - -fn reject_unexpected_positionals( - args: &[String], - command: &str, - subcommand: Option<&str>, -) -> Result<(), CliEnvelope> { - let mut index = 0; - while index < args.len() { - let arg = &args[index]; - let Some(spec) = flag_spec(Some(command), subcommand, arg) else { - return Err(invalid_args(format!("Unexpected positional argument: {arg}"))); - }; - - index += match spec.kind { - FlagKind::Value => 2, - FlagKind::Bool => 1, - }; - } - Ok(()) -} - -async fn open_state() -> Result { - let app_dir = crate::runtime_client::app_data_dir(); - std::fs::create_dir_all(&app_dir).map_err(|err| err.to_string())?; - let storage = dbx_core::storage::Storage::open(&app_dir.join("dbx.db")).await?; - Ok(dbx_core::connection::AppState::new(storage)) -} - -fn redacted_config(config: &dbx_core::models::connection::ConnectionConfig) -> serde_json::Value { - let risk = dbx_core::sql_safety::risk_for_connection("SELECT 1", &config.name, config.color.as_deref()); - serde_json::json!({ - "id": config.id, - "name": config.name, - "databaseType": config.db_type, - "driverProfile": config.driver_profile, - "driverLabel": config.driver_label, - "defaultDatabase": config.database, - "color": config.color, - "sshEnabled": config.ssh_enabled, - "redactedUrl": redacted_url(config), - "risk": risk, - }) -} - -fn redacted_url(config: &dbx_core::models::connection::ConnectionConfig) -> String { - let redacted = if matches!( - config.db_type, - dbx_core::models::connection::DatabaseType::Sqlite | dbx_core::models::connection::DatabaseType::DuckDb - ) { - redact_embedded_path_url(&config.redacted_connection_url()) - } else { - redact_uri_credentials(&config.redacted_connection_url()) - }; - - redact_sensitive_query_params(&redacted) -} - -fn redact_embedded_path_url(value: &str) -> String { - let suffix_start = value.find(['?', '#']).unwrap_or(value.len()); - format!("{}", &value[suffix_start..]) -} - -fn redact_uri_credentials(value: &str) -> String { - let Some(scheme_index) = value.find("://") else { - return value.to_string(); - }; - let authority_start = scheme_index + 3; - let rest = &value[authority_start..]; - let authority_len = rest.find(['/', '?', '#']).unwrap_or(rest.len()); - let authority = &rest[..authority_len]; - let Some(at_index) = authority.rfind('@') else { - return value.to_string(); - }; - - format!("{}{}{}", &value[..authority_start], &authority[at_index + 1..], &rest[authority_len..]) -} - -fn redact_sensitive_query_params(value: &str) -> String { - let Some(query_start) = value.find('?') else { - return value.to_string(); - }; - let suffix = &value[query_start + 1..]; - let (query, fragment) = match suffix.find('#') { - Some(fragment_start) => (&suffix[..fragment_start], &suffix[fragment_start..]), - None => (suffix, ""), - }; - if query.is_empty() { - return value.to_string(); - } - - let query = query.split('&').map(redact_query_pair).collect::>().join("&"); - format!("{}?{}{}", &value[..query_start], query, fragment) -} - -fn redact_query_pair(pair: &str) -> String { - if pair.is_empty() { - return String::new(); - } - let (key, has_value) = match pair.split_once('=') { - Some((key, _)) => (key, true), - None => (pair, false), - }; - if is_sensitive_query_key(key) { - let separator = if has_value { "=" } else { "" }; - return format!("{key}{separator}"); - } - pair.to_string() -} - -fn is_sensitive_query_key(key: &str) -> bool { - let key = key.to_ascii_lowercase(); - ["password", "passwd", "token", "secret", "key", "cert", "path", "file", "sslkey", "passfile"] - .iter() - .any(|needle| key.contains(needle)) -} - -async fn load_headless_connections( -) -> Result, CliEnvelope> { - let state = - open_state().await.map_err(|err| fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false))?; - state - .storage - .load_connections() - .await - .map_err(|err| fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false)) -} - -async fn find_connection( - name: &str, -) -> Result> { - let configs = load_headless_connections().await?; - let id_matches: Vec<_> = configs.iter().filter(|config| config.id == name).collect(); - match id_matches.len() { - 1 => return Ok(id_matches[0].clone()), - 0 => {} - _ => { - return Err(fail( - CliSource::Headless, - CliErrorCode::AmbiguousConnection, - "Connection id is ambiguous", - true, - )); - } - } - - let matches: Vec<_> = configs.into_iter().filter(|config| config.name == name).collect(); - - match matches.len() { - 1 => Ok(matches.into_iter().next().unwrap()), - 0 => Err(fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "Connection not found", true)), - _ => Err(fail(CliSource::Headless, CliErrorCode::AmbiguousConnection, "Connection name is ambiguous", true)), - } -} - -async fn state_with_connection( - config: dbx_core::models::connection::ConnectionConfig, -) -> Result { - let state = open_state().await?; - state.configs.write().await.insert(config.id.clone(), config.clone()); - - match config.db_type { - dbx_core::models::connection::DatabaseType::Sqlite => { - let path = dbx_core::connection::expand_tilde(&config.host); - let pool = dbx_core::db::sqlite::connect_path(&path).await?; - state.connections.write().await.insert(config.id.clone(), dbx_core::connection::PoolKind::Sqlite(pool)); - } - dbx_core::models::connection::DatabaseType::DuckDb => { - let path = dbx_core::connection::expand_tilde(&config.host); - let pool = dbx_core::db::duckdb_driver::connect_path(&path)?; - state.connections.write().await.insert(config.id.clone(), dbx_core::connection::PoolKind::DuckDb(pool)); - } - _ => { - state.get_or_create_pool(&config.id, config.database.as_deref()).await?; - } - } - - Ok(state) -} - -async fn context(args: &[String]) -> CliEnvelope { - if let Err(err) = reject_unexpected_positionals(args, "context", None) { - return err; - } - - match crate::runtime_client::get_json("/context").await { - Ok(data) => ok(CliSource::GuiRuntime, data), - Err(_) => { - let configs = load_headless_connections().await.unwrap_or_default(); - ok( - CliSource::Headless, - serde_json::json!({ - "runtime": "headless", - "activeConnection": configs.first().map(redacted_config), - "configSource": "headless", - }), - ) - } - } -} - -async fn conn_list(args: &[String]) -> CliEnvelope { - if let Err(err) = reject_unexpected_positionals(args, "conn", Some("list")) { - return err; - } - - match load_headless_connections().await { - Ok(configs) => ok( - CliSource::Headless, - serde_json::json!({ - "connections": configs.iter().map(redacted_config).collect::>(), - }), - ), - Err(err) => err, - } -} - -async fn conn_show(name: &str, args: &[String]) -> CliEnvelope { - if let Err(err) = reject_unexpected_positionals(args, "conn", Some("show")) { - return err; - } - - match find_connection(name).await { - Ok(config) => ok(CliSource::Headless, redacted_config(&config)), - Err(err) => err, - } -} - -async fn schema_snapshot(args: &[String]) -> CliEnvelope { - if let Err(err) = reject_unexpected_positionals(args, "schema", Some("snapshot")) { - return err; - } - - let Some(conn) = option_value(args, "--conn") else { - return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); - }; - let config = match find_connection(conn).await { - Ok(config) => config, - Err(err) => return err, - }; - let state = match state_with_connection(config.clone()).await { - Ok(state) => state, - Err(err) => return fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false), - }; - - match dbx_core::schema_snapshot::snapshot(&state, &config.id, option_value(args, "--db"), None).await { - Ok(snapshot) => ok(CliSource::Headless, serde_json::to_value(snapshot).unwrap()), - Err(err) => fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false), - } -} - -async fn safe_query(args: &[String]) -> CliEnvelope { - if let Err(err) = reject_unexpected_positionals(args, "safe-query", None) { - return err; - } - - let Some(conn) = option_value(args, "--conn") else { - return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); - }; - let Some(sql) = option_value(args, "--sql") else { - return fail(CliSource::Headless, CliErrorCode::QueryClassificationFailed, "--sql is required", true); - }; - let limit = match parse_result_limit(args) { - Ok(limit) => limit as usize, - Err(err) => return err, - }; - let config = match find_connection(conn).await { - Ok(config) => config, - Err(err) => return err, - }; - let risk = dbx_core::sql_safety::risk_for_connection(sql, &config.name, config.color.as_deref()); - match risk.operation_class { - dbx_core::sql_safety::OperationClass::Read => {} - dbx_core::sql_safety::OperationClass::Write if risk.is_production => { - return blocked_query(CliErrorCode::ProductionWriteBlocked, &risk); - } - dbx_core::sql_safety::OperationClass::Write => { - return blocked_query(CliErrorCode::HandoffRequired, &risk); - } - dbx_core::sql_safety::OperationClass::Ddl => { - return blocked_query(CliErrorCode::DdlBlocked, &risk); - } - dbx_core::sql_safety::OperationClass::Unknown => { - return blocked_query(CliErrorCode::QueryClassificationFailed, &risk); - } - } - let state = match state_with_connection(config.clone()).await { - Ok(state) => state, - Err(err) => return fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false), - }; - let database = option_value(args, "--db").or(config.database.as_deref()).unwrap_or(""); - - match dbx_core::query::execute_sql_statement_with_row_limit(&state, &config.id, database, sql, None, None, limit) - .await - { - Ok(result) => ok( - CliSource::Headless, - serde_json::json!({ - "risk": risk, - "result": result, - }), - ), - Err(err) => fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false), - } -} - -fn blocked_query(code: CliErrorCode, risk: &dbx_core::sql_safety::RiskMetadata) -> CliEnvelope { - fail(CliSource::Headless, code, serde_json::to_string(risk).unwrap(), true) -} - -async fn handoff(args: &[String]) -> CliEnvelope { - if let Err(err) = reject_unexpected_positionals(args, "handoff", None) { - return err; - } - - let Some(conn) = required_option(args, "--conn") else { - return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); - }; - let Some(title) = required_option(args, "--title") else { - return invalid_args("--title is required"); - }; - let sql_inline = option_value(args, "--sql"); - let sql_file = option_value(args, "--sql-file"); - let sql = match (sql_inline, sql_file) { - (Some(_), Some(_)) => return invalid_args("Use exactly one of --sql or --sql-file"), - (Some(sql), None) => sql.to_string(), - (None, Some(path)) => match std::fs::read_to_string(path) { - Ok(sql) => sql, - Err(err) => return invalid_args(format!("Failed to read --sql-file: {err}")), - }, - (None, None) => return invalid_args("Use exactly one of --sql or --sql-file"), - }; - - if sql.trim().is_empty() { - return invalid_args("SQL must not be empty"); - } - - let config = match find_connection(conn).await { - Ok(config) => config, - Err(err) => return err, - }; - let risk = dbx_core::sql_safety::risk_for_connection(&sql, &config.name, config.color.as_deref()); - let item = dbx_core::handoff::HandoffItem::queued( - config.id, - config.name, - config.database, - title.to_string(), - option_value(args, "--description").map(str::to_string), - sql, - risk.operation_class, - risk.risk_level, - risk.is_production, - ); - - if let Ok(data) = crate::runtime_client::post_json("/handoff", serde_json::to_value(&item).unwrap()).await { - return ok(CliSource::GuiRuntime, data); - } - - match queue_handoff(&item).await { - Ok(()) => ok(CliSource::Headless, serde_json::json!({ "id": item.id, "status": "queued" })), - Err(err) => fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false), - } -} - -async fn selection(args: &[String]) -> CliEnvelope { - if let Err(err) = reject_unexpected_positionals(args, "selection", None) { - return err; - } - - match crate::runtime_client::get_json("/selection").await { - Ok(data) => ok(CliSource::GuiRuntime, data), - Err(_) => runtime_required("dbx selection requires DBX GUI runtime."), - } -} - -async fn result_current(args: &[String]) -> CliEnvelope { - if let Err(err) = reject_unexpected_positionals(args, "result", Some("current")) { - return err; - } - - let limit = match parse_result_limit(args) { - Ok(limit) => limit, - Err(err) => return err, - }; - - match crate::runtime_client::get_json_with_query("/result/current", &[("limit", limit.to_string())]).await { - Ok(data) => ok(CliSource::GuiRuntime, data), - Err(_) => runtime_required("dbx result current requires DBX GUI runtime."), - } -} - -fn runtime_required(message: &str) -> CliEnvelope { - fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, message, true) -} - -fn invalid_args(message: impl Into) -> CliEnvelope { - fail(CliSource::Headless, CliErrorCode::InternalError, message, true) -} - -fn option_value<'a>(args: &'a [String], key: &str) -> Option<&'a str> { - args.windows(2).find(|pair| pair[0] == key).map(|pair| pair[1].as_str()) -} - -fn required_option<'a>(args: &'a [String], key: &str) -> Option<&'a str> { - option_value(args, key).map(str::trim).filter(|value| !value.is_empty()) -} - -fn parse_result_limit(args: &[String]) -> Result> { - let default_limit = DEFAULT_RESULT_LIMIT.to_string(); - let raw = option_value(args, "--limit").unwrap_or(&default_limit); - let limit = raw - .parse::() - .map_err(|_| invalid_args(format!("--limit must be a positive integer between 1 and {MAX_RESULT_LIMIT}")))?; - if !(1..=MAX_RESULT_LIMIT).contains(&limit) { - return Err(invalid_args(format!("--limit must be a positive integer between 1 and {MAX_RESULT_LIMIT}"))); - } - Ok(limit) -} - -async fn queue_handoff(item: &dbx_core::handoff::HandoffItem) -> Result<(), String> { - let app_dir = crate::runtime_client::app_data_dir(); - std::fs::create_dir_all(&app_dir).map_err(|err| err.to_string())?; - let storage = dbx_core::storage::Storage::open(&app_dir.join("dbx.db")).await?; - storage.save_handoff(item).await -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::runtime_client::ENV_LOCK; - use dbx_core::cli::CliErrorCode; - use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; - use std::sync::{Arc, Mutex}; - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - use tokio::net::TcpListener; - - fn assert_failure_code(env: CliEnvelope, expected: CliErrorCode) { - match env { - CliEnvelope::Failure { error, .. } => assert_eq!(error.code, expected), - CliEnvelope::Success { .. } => panic!("expected failure envelope"), - } - } - - fn assert_failure_message_contains(env: CliEnvelope, expected: &str) { - match env { - CliEnvelope::Failure { error, .. } => assert!( - error.message.contains(expected), - "expected error message to contain {expected:?}, got {:?}", - error.message - ), - CliEnvelope::Success { .. } => panic!("expected failure envelope"), - } - } - - fn mysql_fixture(id: &str, name: &str, color: Option<&str>, password: &str) -> ConnectionConfig { - ConnectionConfig { - id: id.to_string(), - name: name.to_string(), - db_type: DatabaseType::Mysql, - driver_profile: None, - driver_label: Some("MySQL".to_string()), - url_params: None, - host: "127.0.0.1".to_string(), - port: 3306, - username: "root".to_string(), - password: password.to_string(), - database: Some("app".to_string()), - color: color.map(str::to_string), - ssh_enabled: true, - ssh_host: "bastion.internal".to_string(), - ssh_port: 22, - ssh_user: "deploy".to_string(), - ssh_password: "ssh-secret".to_string(), - ssh_key_path: String::new(), - ssh_key_passphrase: "key-secret".to_string(), - ssh_expose_lan: false, - ssh_connect_timeout_secs: dbx_core::models::connection::default_ssh_connect_timeout_secs(), - proxy_enabled: false, - proxy_type: dbx_core::models::connection::ProxyType::Socks5, - proxy_host: String::new(), - proxy_port: 1080, - proxy_username: String::new(), - proxy_password: String::new(), - ssl: false, - sysdba: false, - connection_string: Some(format!("mysql://root:{password}@127.0.0.1:3306/app")), - external_config: None, - jdbc_driver_class: None, - jdbc_driver_paths: Vec::new(), - } - } - - fn embedded_fixture(id: &str, name: &str, db_type: DatabaseType, path: &std::path::Path) -> ConnectionConfig { - let mut config = mysql_fixture(id, name, None, ""); - config.db_type = db_type; - config.host = path.display().to_string(); - config.port = 0; - config.username = String::new(); - config.password = String::new(); - config.database = None; - config.connection_string = None; - config - } - - async fn seed_connections(dir: &std::path::Path, configs: &[ConnectionConfig]) { - std::fs::create_dir_all(dir).unwrap(); - let storage = dbx_core::storage::Storage::open(&dir.join("dbx.db")).await.unwrap(); - storage.save_connections(configs).await.unwrap(); - } - - #[derive(Debug, Clone)] - struct RuntimeRequestFixture { - first_line: String, - authorization: Option, - body: String, - } - - #[cfg(unix)] - fn set_owner_only_mode(path: &std::path::Path) { - use std::os::unix::fs::PermissionsExt; - - let mut permissions = std::fs::metadata(path).unwrap().permissions(); - permissions.set_mode(0o600); - std::fs::set_permissions(path, permissions).unwrap(); - } - - fn write_runtime_discovery(dir: &std::path::Path, port: u16, token: &str) { - let path = dir.join("agent-runtime.json"); - std::fs::write(&path, serde_json::json!({ "port": port, "token": token }).to_string()).unwrap(); - #[cfg(unix)] - set_owner_only_mode(&path); - } - - async fn start_runtime_fixture( - dir: &std::path::Path, - expected_requests: usize, - ) -> Arc>> { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - write_runtime_discovery(dir, port, "runtime-token"); - let requests = Arc::new(Mutex::new(Vec::new())); - let captured = requests.clone(); - - tokio::spawn(async move { - for _ in 0..expected_requests { - let (mut stream, _) = listener.accept().await.unwrap(); - let request = read_fixture_request(&mut stream).await; - let body = fixture_response(&request); - captured.lock().unwrap().push(request); - let response = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", - body.len(), - body - ); - stream.write_all(response.as_bytes()).await.unwrap(); - } - }); - - requests - } - - async fn read_fixture_request(stream: &mut tokio::net::TcpStream) -> RuntimeRequestFixture { - let mut buf = Vec::new(); - let header_end = loop { - if let Some(index) = buf.windows(4).position(|window| window == b"\r\n\r\n") { - break index; - } - let mut chunk = [0u8; 1024]; - let n = stream.read(&mut chunk).await.unwrap(); - assert!(n > 0, "runtime client closed before sending headers"); - buf.extend_from_slice(&chunk[..n]); - }; - let header_text = String::from_utf8_lossy(&buf[..header_end]); - let mut lines = header_text.lines(); - let first_line = lines.next().unwrap_or_default().to_string(); - let headers: Vec<(String, String)> = lines - .filter_map(|line| { - let (name, value) = line.split_once(':')?; - Some((name.trim().to_string(), value.trim().to_string())) - }) - .collect(); - let content_length = headers - .iter() - .find(|(name, _)| name.eq_ignore_ascii_case("content-length")) - .and_then(|(_, value)| value.parse::().ok()) - .unwrap_or(0); - let body_start = header_end + 4; - let body_end = body_start + content_length; - - while buf.len() < body_end { - let mut chunk = [0u8; 1024]; - let n = stream.read(&mut chunk).await.unwrap(); - assert!(n > 0, "runtime client closed before sending body"); - buf.extend_from_slice(&chunk[..n]); - } - - RuntimeRequestFixture { - first_line, - authorization: headers - .iter() - .find(|(name, _)| name.eq_ignore_ascii_case("authorization")) - .map(|(_, value)| value.clone()), - body: String::from_utf8_lossy(&buf[body_start..body_end]).to_string(), - } - } - - fn fixture_response(request: &RuntimeRequestFixture) -> String { - assert_eq!(request.authorization.as_deref(), Some("Bearer runtime-token")); - - if request.first_line.starts_with("GET /context ") || request.first_line.starts_with("GET /context?") { - return serde_json::json!({ - "activeConnectionId": "runtime-conn", - "activeConnectionName": "Runtime Connection", - "database": "app", - "sql": "select * from users" - }) - .to_string(); - } - if request.first_line.starts_with("GET /selection ") || request.first_line.starts_with("GET /selection?") { - return serde_json::json!({ "type": "grid-cells", "cells": [[1, "ada@example.com"]] }).to_string(); - } - if request.first_line.starts_with("GET /result/current?") { - return serde_json::json!({ "columns": ["id"], "rows": [[1], [2]], "truncated": true }).to_string(); - } - if request.first_line.starts_with("POST /handoff ") || request.first_line.starts_with("POST /handoff?") { - let item: dbx_core::handoff::HandoffItem = serde_json::from_str(&request.body).unwrap(); - return serde_json::json!({ "id": item.id, "status": "shown" }).to_string(); - } - - panic!("unexpected runtime request: {}", request.first_line); - } - - async fn create_sqlite_fixture(path: &std::path::Path) { - std::fs::File::create(path).unwrap(); - let pool = dbx_core::db::sqlite::connect_path(&path.display().to_string()).await.unwrap(); - dbx_core::db::sqlite::execute_query(&pool, "CREATE TABLE teams (id INTEGER PRIMARY KEY, name TEXT NOT NULL)") - .await - .unwrap(); - dbx_core::db::sqlite::execute_query( - &pool, - "CREATE TABLE users (id INTEGER PRIMARY KEY, team_id INTEGER NOT NULL, email TEXT NOT NULL, FOREIGN KEY(team_id) REFERENCES teams(id))", - ) - .await - .unwrap(); - dbx_core::db::sqlite::execute_query(&pool, "INSERT INTO teams (id, name) VALUES (1, 'Core'), (2, 'Data')") - .await - .unwrap(); - dbx_core::db::sqlite::execute_query( - &pool, - "INSERT INTO users (id, team_id, email) VALUES (1, 1, 'ada@example.com'), (2, 2, 'grace@example.com')", - ) - .await - .unwrap(); - pool.close().await; - } - - async fn create_large_sqlite_fixture(path: &std::path::Path, row_count: usize) { - std::fs::File::create(path).unwrap(); - let pool = dbx_core::db::sqlite::connect_path(&path.display().to_string()).await.unwrap(); - dbx_core::db::sqlite::execute_query( - &pool, - "CREATE TABLE numbers (id INTEGER PRIMARY KEY, value TEXT NOT NULL)", - ) - .await - .unwrap(); - for id in 1..=row_count { - dbx_core::db::sqlite::execute_query( - &pool, - &format!("INSERT INTO numbers (id, value) VALUES ({id}, 'value-{id}')"), - ) - .await - .unwrap(); - } - pool.close().await; - } - - fn success_data(env: CliEnvelope) -> serde_json::Value { - match env { - CliEnvelope::Success { source, data, .. } => { - assert_eq!(source, CliSource::Headless); - data - } - CliEnvelope::Failure { error, .. } => panic!("expected success envelope, got {error:?}"), - } - } - - fn gui_success_data(env: CliEnvelope) -> serde_json::Value { - match env { - CliEnvelope::Success { source, data, .. } => { - assert_eq!(source, CliSource::GuiRuntime); - data - } - CliEnvelope::Failure { error, .. } => panic!("expected GUI runtime success envelope, got {error:?}"), - } - } - - #[test] - fn parser_allows_free_text_option_values_to_start_with_dashes() { - let parsed = parse_args(vec![ - "handoff".into(), - "--conn".into(), - "local".into(), - "--title".into(), - "-- review generated SQL".into(), - "--sql".into(), - "-- explain select 1".into(), - "--description".into(), - "-- optional note".into(), - ]) - .expect("free-text values beginning with -- should parse"); - - assert_eq!( - parsed.positionals, - vec![ - "handoff", - "--conn", - "local", - "--title", - "-- review generated SQL", - "--sql", - "-- explain select 1", - "--description", - "-- optional note", - ] - ); - } - - #[test] - fn parser_accepts_global_json_format_before_between_and_after_command_args() { - let cases = [ - vec!["--format", "json", "safe-query", "--conn", "local", "--sql", "select 1"], - vec!["safe-query", "--conn", "local", "--format", "json", "--sql", "select 1"], - vec!["safe-query", "--conn", "local", "--sql", "select 1", "--format", "json"], - ]; - - for args in cases { - let parsed = parse_args(args.iter().map(|value| value.to_string()).collect()) - .unwrap_or_else(|err| panic!("expected parse success for {args:?}, got {err:?}")); - assert_eq!(parsed.positionals, vec!["safe-query", "--conn", "local", "--sql", "select 1"]); - } - } - - #[tokio::test] - async fn rejects_extra_positionals_for_each_command() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - - let cases = [ - vec!["context", "extra"], - vec!["conn", "list", "extra"], - vec!["conn", "show", "local", "extra"], - vec!["schema", "snapshot", "--conn", "local", "extra"], - vec!["safe-query", "--conn", "local", "--sql", "select 1", "extra"], - vec!["handoff", "--conn", "local", "--title", "Review", "--sql", "select 1", "extra"], - vec!["selection", "extra"], - vec!["result", "current", "--limit", "50", "extra"], - ]; - - for args in cases { - let env = dispatch(args.iter().map(|value| value.to_string()).collect()).await; - assert_failure_message_contains(env, "Unexpected positional argument"); - } - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn gui_only_commands_return_runtime_required_without_runtime() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - - assert_failure_code( - dispatch(vec!["selection".into(), "--format".into(), "json".into()]).await, - CliErrorCode::GuiRuntimeRequired, - ); - assert_failure_code( - dispatch(vec![ - "result".into(), - "current".into(), - "--limit".into(), - "25".into(), - "--format".into(), - "json".into(), - ]) - .await, - CliErrorCode::GuiRuntimeRequired, - ); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn recognizes_all_eight_cli_commands_with_json_format() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - - let cases = [ - vec!["context", "--format", "json"], - vec!["conn", "list", "--format", "json"], - vec!["conn", "show", "__missing__", "--redacted", "--format", "json"], - vec!["schema", "snapshot", "--format", "json"], - vec!["safe-query", "--format", "json"], - vec!["handoff", "--format", "json"], - vec!["selection", "--format", "json"], - vec!["result", "current", "--limit", "50", "--format", "json"], - ]; - - for args in cases { - let env = dispatch(args.iter().map(|value| value.to_string()).collect()).await; - let json = serde_json::to_value(&env).unwrap(); - assert!(json.get("ok").is_some(), "missing ok for args: {args:?}"); - assert!(json.get("source").is_some(), "missing source for args: {args:?}"); - assert!(json.get("data").is_some() || json.get("error").is_some(), "missing data/error for args: {args:?}"); - } - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn context_reads_headless_storage_and_returns_redacted_active_connection() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections(dir.path(), &[mysql_fixture("prod-id", "prod-main", Some("#ef4444"), "super-secret")]).await; - - let data = success_data(dispatch(vec!["context".into(), "--format".into(), "json".into()]).await); - - assert_eq!(data["runtime"], "headless"); - assert_eq!(data["activeConnection"]["id"], "prod-id"); - assert_eq!(data["activeConnection"]["name"], "prod-main"); - assert_eq!(data["activeConnection"]["risk"]["isProduction"], true); - assert_eq!(data["configSource"], "headless"); - let json = serde_json::to_string(&data).unwrap(); - assert!(json.contains("configSource")); - assert!(!json.contains(&dir.path().join("dbx.db").display().to_string())); - assert!(!json.contains("super-secret")); - assert!(!json.contains("ssh-secret")); - assert!(!json.contains("key-secret")); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn context_returns_headless_minimal_context_for_empty_or_unopenable_app_data() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let empty_app_data = dir.path().join("empty-app-data"); - let unopenable_app_data = dir.path().join("app-data-file"); - std::fs::write(&unopenable_app_data, "not a directory").unwrap(); - - for app_data in [&empty_app_data, &unopenable_app_data] { - std::env::set_var("DBX_APP_DATA_DIR", app_data); - - let data = success_data(dispatch(vec!["context".into(), "--format".into(), "json".into()]).await); - - assert_eq!(data["runtime"], "headless"); - assert!(data["activeConnection"].is_null()); - assert_eq!(data["configSource"], "headless"); - } - - std::env::set_var("DBX_APP_DATA_DIR", &unopenable_app_data); - assert_failure_code( - dispatch(vec!["conn".into(), "list".into(), "--format".into(), "json".into()]).await, - CliErrorCode::InternalError, - ); - assert_failure_code( - dispatch(vec!["conn".into(), "show".into(), "missing".into(), "--format".into(), "json".into()]).await, - CliErrorCode::InternalError, - ); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn runtime_discovery_serves_context_selection_and_current_result_limit() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - let requests = start_runtime_fixture(dir.path(), 3).await; - - let context = gui_success_data(dispatch(vec!["context".into(), "--format".into(), "json".into()]).await); - let selection = gui_success_data(dispatch(vec!["selection".into(), "--format".into(), "json".into()]).await); - let result = gui_success_data( - dispatch(vec![ - "result".into(), - "current".into(), - "--limit".into(), - "7".into(), - "--format".into(), - "json".into(), - ]) - .await, - ); - - assert_eq!(context["activeConnectionId"], "runtime-conn"); - assert_eq!(selection["type"], "grid-cells"); - assert_eq!(result["rows"], serde_json::json!([[1], [2]])); - - let requests = requests.lock().unwrap(); - assert_eq!(requests[0].first_line, "GET /context? HTTP/1.1"); - assert_eq!(requests[1].first_line, "GET /selection? HTTP/1.1"); - assert_eq!(requests[2].first_line, "GET /result/current?limit=7 HTTP/1.1"); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn conn_list_reads_storage_and_returns_redacted_connection_dtos_with_risk_metadata() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections( - dir.path(), - &[ - mysql_fixture("prod-id", "prod-main", Some("#ef4444"), "prod-secret"), - mysql_fixture("dev-id", "dev-main", Some("#22c55e"), "dev-secret"), - ], - ) - .await; - - let data = success_data(dispatch(vec!["conn".into(), "list".into(), "--format".into(), "json".into()]).await); - let connections = data["connections"].as_array().expect("connections should be an array"); - - assert_eq!(connections.len(), 2); - assert_eq!(connections[0]["id"], "prod-id"); - assert_eq!(connections[0]["redactedUrl"], "mysql://127.0.0.1:3306/app?ssl-mode=preferred&charset=utf8mb4"); - assert_eq!(connections[0]["risk"]["isProduction"], true); - assert_eq!(connections[0]["risk"]["productionReason"], "red connection color"); - assert_eq!(connections[1]["risk"]["isProduction"], false); - let json = serde_json::to_string(&data).unwrap(); - assert!(!json.contains("prod-secret")); - assert!(!json.contains("dev-secret")); - assert!(!json.contains("root:")); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn conn_list_redacts_sqlite_and_duckdb_file_paths() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - let sqlite_path = dir.path().join("private").join("app.sqlite"); - let duckdb_path = dir.path().join("private").join("warehouse.duckdb"); - seed_connections( - dir.path(), - &[ - embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &sqlite_path), - embedded_fixture("duckdb-id", "local-duckdb", DatabaseType::DuckDb, &duckdb_path), - ], - ) - .await; - - let data = success_data(dispatch(vec!["conn".into(), "list".into(), "--format".into(), "json".into()]).await); - let connections = data["connections"].as_array().expect("connections should be an array"); - - assert_eq!(connections[0]["redactedUrl"], "?mode=rwc"); - assert_eq!(connections[1]["redactedUrl"], "?mode=rwc"); - let json = serde_json::to_string(&data).unwrap(); - assert!(!json.contains("app.sqlite")); - assert!(!json.contains("warehouse.duckdb")); - assert!(!json.contains(dir.path().to_str().unwrap())); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn conn_show_supports_id_lookup_and_returns_not_found_or_ambiguous_errors() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections( - dir.path(), - &[ - mysql_fixture("prod-id", "prod-main", Some("#ef4444"), "prod-secret"), - mysql_fixture("dup-1", "shared", None, "first-secret"), - mysql_fixture("dup-2", "shared", None, "second-secret"), - ], - ) - .await; - - let shown = success_data( - dispatch(vec![ - "conn".into(), - "show".into(), - "prod-id".into(), - "--redacted".into(), - "--format".into(), - "json".into(), - ]) - .await, - ); - assert_eq!(shown["id"], "prod-id"); - assert_eq!(shown["risk"]["isProduction"], true); - assert!(!serde_json::to_string(&shown).unwrap().contains("prod-secret")); - - assert_failure_code( - dispatch(vec!["conn".into(), "show".into(), "__missing__".into(), "--format".into(), "json".into()]).await, - CliErrorCode::ConnectionNotFound, - ); - assert_failure_code( - dispatch(vec!["conn".into(), "show".into(), "shared".into(), "--format".into(), "json".into()]).await, - CliErrorCode::AmbiguousConnection, - ); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn conn_list_and_show_redact_sensitive_query_params() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - let mut mysql = mysql_fixture("prod-id", "prod-main", Some("#ef4444"), "prod-secret"); - mysql.url_params = - Some("password=pw&sslkey=/tmp/client.key&sslcert=/tmp/client.crt&token=abc&charset=utf8mb4".to_string()); - seed_connections(dir.path(), &[mysql]).await; - - let list = success_data(dispatch(vec!["conn".into(), "list".into(), "--format".into(), "json".into()]).await); - let listed_url = list["connections"][0]["redactedUrl"].as_str().unwrap(); - assert!(listed_url.contains("password=")); - assert!(listed_url.contains("sslkey=")); - assert!(listed_url.contains("sslcert=")); - assert!(listed_url.contains("token=")); - assert!(listed_url.contains("charset=utf8mb4")); - - let shown = success_data( - dispatch(vec!["conn".into(), "show".into(), "prod-id".into(), "--format".into(), "json".into()]).await, - ); - let shown_json = serde_json::to_string(&shown).unwrap(); - assert!(!shown_json.contains("pw")); - assert!(!shown_json.contains("/tmp/client.key")); - assert!(!shown_json.contains("/tmp/client.crt")); - assert!(!shown_json.contains("abc")); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn conn_show_unique_id_match_takes_precedence_over_ambiguous_name() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections( - dir.path(), - &[ - mysql_fixture("shared", "id-target", Some("#22c55e"), "target-secret"), - mysql_fixture("dup-1", "shared", None, "first-secret"), - mysql_fixture("dup-2", "shared", None, "second-secret"), - ], - ) - .await; - - let shown = success_data( - dispatch(vec!["conn".into(), "show".into(), "shared".into(), "--format".into(), "json".into()]).await, - ); - - assert_eq!(shown["id"], "shared"); - assert_eq!(shown["name"], "id-target"); - let json = serde_json::to_string(&shown).unwrap(); - assert!(!json.contains("target-secret")); - assert!(!json.contains("first-secret")); - assert!(!json.contains("second-secret")); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn schema_snapshot_executes_headless_sqlite_snapshot_from_storage() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let data_path = dir.path().join("fixture.sqlite"); - create_sqlite_fixture(&data_path).await; - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections( - dir.path(), - &[embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &data_path)], - ) - .await; - - let data = success_data( - dispatch(vec![ - "schema".into(), - "snapshot".into(), - "--conn".into(), - "local-sqlite".into(), - "--format".into(), - "json".into(), - ]) - .await, - ); - - assert_eq!(data["connectionId"], "sqlite-id"); - assert_eq!(data["database"], "main"); - let tables = data["tables"].as_array().expect("tables should be an array"); - assert!(tables.iter().any(|table| table["name"] == "users")); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn safe_query_executes_read_sqlite_query_with_limit_and_risk() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let data_path = dir.path().join("fixture.sqlite"); - create_sqlite_fixture(&data_path).await; - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections( - dir.path(), - &[embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &data_path)], - ) - .await; - - let data = success_data( - dispatch(vec![ - "safe-query".into(), - "--conn".into(), - "sqlite-id".into(), - "--sql".into(), - "SELECT email FROM users ORDER BY id".into(), - "--limit".into(), - "1".into(), - "--format".into(), - "json".into(), - ]) - .await, - ); - - assert_eq!(data["risk"]["operationClass"], "read"); - assert_eq!(data["result"]["columns"], serde_json::json!(["email"])); - assert_eq!(data["result"]["rows"], serde_json::json!([["ada@example.com"]])); - assert_eq!(data["result"]["truncated"], true); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn safe_query_expands_tilde_for_sqlite_connection_path() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let home = dir.path().join("home"); - std::fs::create_dir_all(&home).unwrap(); - let data_path = home.join("fixture.sqlite"); - create_sqlite_fixture(&data_path).await; - let previous_home = std::env::var_os("HOME"); - std::env::set_var("HOME", &home); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - let mut config = embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &data_path); - config.host = "~/fixture.sqlite".to_string(); - seed_connections(dir.path(), &[config]).await; - - let data = success_data( - dispatch(vec![ - "safe-query".into(), - "--conn".into(), - "local-sqlite".into(), - "--sql".into(), - "SELECT COUNT(*) FROM users".into(), - "--format".into(), - "json".into(), - ]) - .await, - ); - - assert_eq!(data["result"]["rows"], serde_json::json!([[2]])); - - std::env::remove_var("DBX_APP_DATA_DIR"); - if let Some(previous_home) = previous_home { - std::env::set_var("HOME", previous_home); - } else { - std::env::remove_var("HOME"); - } - } - - #[tokio::test] - async fn safe_query_large_sqlite_result_set_respects_limit() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let data_path = dir.path().join("large.sqlite"); - create_large_sqlite_fixture(&data_path, 200).await; - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections( - dir.path(), - &[embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &data_path)], - ) - .await; - - let data = success_data( - dispatch(vec![ - "safe-query".into(), - "--conn".into(), - "sqlite-id".into(), - "--sql".into(), - "SELECT id, value FROM numbers ORDER BY id".into(), - "--limit".into(), - "7".into(), - "--format".into(), - "json".into(), - ]) - .await, - ); - - let rows = data["result"]["rows"].as_array().expect("rows should be an array"); - assert!(rows.len() <= 7, "safe-query returned {} rows for limit 7", rows.len()); - assert_eq!(rows.len(), 7); - assert_eq!(data["result"]["truncated"], true); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn schema_snapshot_and_safe_query_sanitize_internal_error_envelopes() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let private_path = dir.path().join("private").join("missing.sqlite"); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections( - dir.path(), - &[embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &private_path)], - ) - .await; - - let cases = [ - dispatch(vec![ - "schema".into(), - "snapshot".into(), - "--conn".into(), - "sqlite-id".into(), - "--format".into(), - "json".into(), - ]) - .await, - dispatch(vec![ - "safe-query".into(), - "--conn".into(), - "sqlite-id".into(), - "--sql".into(), - "SELECT 1".into(), - "--format".into(), - "json".into(), - ]) - .await, - ]; - - for env in cases { - match env { - CliEnvelope::Failure { error, .. } => { - let message = error.message; - assert!(!message.contains(dir.path().to_str().unwrap())); - assert!(!message.contains("missing.sqlite")); - assert!(!message.contains("Database file does not exist")); - assert!(!message.contains("SQLite connection failed")); - } - CliEnvelope::Success { .. } => panic!("expected sanitized failure envelope"), - } - } - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn safe_query_blocks_non_read_sql_with_structured_risk_errors() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let data_path = dir.path().join("fixture.sqlite"); - create_sqlite_fixture(&data_path).await; - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - let mut prod = embedded_fixture("prod-sqlite", "prod-sqlite", DatabaseType::Sqlite, &data_path); - prod.color = Some("#ef4444".to_string()); - let dev = embedded_fixture("dev-sqlite", "dev-sqlite", DatabaseType::Sqlite, &data_path); - seed_connections(dir.path(), &[prod, dev]).await; - - let cases = [ - ( - "prod-sqlite", - "UPDATE users SET email = 'x@example.com' WHERE id = 1", - CliErrorCode::ProductionWriteBlocked, - "write", - true, - ), - ( - "dev-sqlite", - "UPDATE users SET email = 'x@example.com' WHERE id = 1", - CliErrorCode::HandoffRequired, - "write", - false, - ), - ("prod-sqlite", "DROP TABLE users", CliErrorCode::DdlBlocked, "ddl", true), - ("prod-sqlite", "VACUUM", CliErrorCode::QueryClassificationFailed, "unknown", true), - ]; - - for (conn, sql, expected_code, expected_class, expected_production) in cases { - match dispatch(vec![ - "safe-query".into(), - "--conn".into(), - conn.into(), - "--sql".into(), - sql.into(), - "--format".into(), - "json".into(), - ]) - .await - { - CliEnvelope::Failure { error, .. } => { - assert_eq!(error.code, expected_code); - let risk: serde_json::Value = serde_json::from_str(&error.message).unwrap(); - assert_eq!(risk["operationClass"], expected_class); - assert_eq!(risk["isProduction"], expected_production); - } - CliEnvelope::Success { .. } => panic!("expected safe-query to block {sql}"), - } - } - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn unknown_command_returns_internal_error_envelope() { - assert_failure_code( - dispatch(vec!["not-a-command".into(), "--format".into(), "json".into()]).await, - CliErrorCode::InternalError, - ); - } - - #[tokio::test] - async fn rejects_non_json_format() { - assert_failure_code( - dispatch(vec!["context".into(), "--format".into(), "text".into()]).await, - CliErrorCode::InternalError, - ); - } - - #[tokio::test] - async fn rejects_unknown_flags_with_error_envelope() { - assert_failure_code( - dispatch(vec!["context".into(), "--unknown".into(), "value".into()]).await, - CliErrorCode::InternalError, - ); - } - - #[tokio::test] - async fn rejects_missing_option_values_with_error_envelope() { - assert_failure_code( - dispatch(vec!["safe-query".into(), "--conn".into(), "--sql".into(), "select 1".into()]).await, - CliErrorCode::InternalError, - ); - } - - #[tokio::test] - async fn validates_handoff_required_options_and_sql_source() { - assert_failure_code( - dispatch(vec!["handoff".into(), "--title".into(), "Review".into(), "--sql".into(), "select 1".into()]) - .await, - CliErrorCode::ConnectionNotFound, - ); - assert_failure_code( - dispatch(vec!["handoff".into(), "--conn".into(), "local".into(), "--sql".into(), "select 1".into()]).await, - CliErrorCode::InternalError, - ); - assert_failure_code( - dispatch(vec![ - "handoff".into(), - "--conn".into(), - "local".into(), - "--title".into(), - "Review".into(), - "--sql".into(), - "select 1".into(), - "--sql-file".into(), - "query.sql".into(), - ]) - .await, - CliErrorCode::InternalError, - ); - } - - #[tokio::test] - async fn handoff_queues_sql_file_when_runtime_is_unavailable() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let sql_file = dir.path().join("query.sql"); - std::fs::write(&sql_file, "select 1").unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections(dir.path(), &[mysql_fixture("local-id", "local", Some("#22c55e"), "local-secret")]).await; - - let env = dispatch(vec![ - "handoff".into(), - "--conn".into(), - "local".into(), - "--title".into(), - "Review".into(), - "--sql-file".into(), - sql_file.display().to_string(), - ]) - .await; - - match env { - CliEnvelope::Success { source, data, .. } => { - assert_eq!(source, CliSource::Headless); - assert_eq!(data["status"], "queued"); - } - CliEnvelope::Failure { error, .. } => panic!("expected queued handoff, got {error:?}"), - } - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn handoff_posts_to_discovered_runtime_instead_of_headless_queue() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections(dir.path(), &[mysql_fixture("local-id", "local", Some("#22c55e"), "local-secret")]).await; - let requests = start_runtime_fixture(dir.path(), 1).await; - - let data = gui_success_data( - dispatch(vec![ - "handoff".into(), - "--conn".into(), - "local".into(), - "--title".into(), - "Review generated SQL".into(), - "--description".into(), - "from agent".into(), - "--sql".into(), - "UPDATE users SET active = 0 WHERE id = 1".into(), - "--format".into(), - "json".into(), - ]) - .await, - ); - - assert_eq!(data["status"], "shown"); - let requests = requests.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert_eq!(requests[0].first_line, "POST /handoff? HTTP/1.1"); - let posted: serde_json::Value = serde_json::from_str(&requests[0].body).unwrap(); - assert_eq!(posted["connectionId"], "local-id"); - assert_eq!(posted["connectionName"], "local"); - assert_eq!(posted["title"], "Review generated SQL"); - assert_eq!(posted["description"], "from agent"); - assert_eq!(posted["status"], "queued"); - - let storage = dbx_core::storage::Storage::open(&dir.path().join("dbx.db")).await.unwrap(); - assert!(storage.load_pending_handoffs().await.unwrap().is_empty()); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn handoff_queues_with_resolved_connection_metadata_and_risk() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - seed_connections(dir.path(), &[mysql_fixture("prod-id", "prod-main", Some("#ef4444"), "prod-secret")]).await; - - let env = dispatch(vec![ - "handoff".into(), - "--conn".into(), - "prod-id".into(), - "--title".into(), - "Review write".into(), - "--description".into(), - "generated by agent".into(), - "--sql".into(), - "UPDATE users SET active = 0 WHERE id = 1".into(), - "--format".into(), - "json".into(), - ]) - .await; - - match env { - CliEnvelope::Success { source, data, .. } => { - assert_eq!(source, CliSource::Headless); - assert_eq!(data["status"], "queued"); - } - CliEnvelope::Failure { error, .. } => panic!("expected queued handoff, got {error:?}"), - } - - let storage = dbx_core::storage::Storage::open(&dir.path().join("dbx.db")).await.unwrap(); - let queued = storage.load_pending_handoffs().await.unwrap(); - - assert_eq!(queued.len(), 1); - assert_eq!(queued[0].connection_id, "prod-id"); - assert_eq!(queued[0].connection_name, "prod-main"); - assert_eq!(queued[0].database.as_deref(), Some("app")); - assert_eq!(queued[0].risk_level, dbx_core::sql_safety::RiskLevel::High); - assert!(queued[0].is_production); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[tokio::test] - async fn result_current_rejects_non_positive_and_over_limit_values() { - for limit in ["0", "-1", "1001", "abc"] { - assert_failure_code( - dispatch(vec!["result".into(), "current".into(), "--limit".into(), limit.into()]).await, - CliErrorCode::InternalError, - ); - } - } -} diff --git a/crates/dbx-cli/src/main.rs b/crates/dbx-cli/src/main.rs deleted file mode 100644 index 9a5a19356..000000000 --- a/crates/dbx-cli/src/main.rs +++ /dev/null @@ -1,10 +0,0 @@ -mod commands; -mod runtime_client; - -#[tokio::main] -async fn main() { - if let Err(err) = commands::run(std::env::args().skip(1).collect()).await { - println!("{}", serde_json::to_string_pretty(&err).unwrap_or_else(|_| "{\"ok\":false}".to_string())); - std::process::exit(1); - } -} diff --git a/crates/dbx-cli/src/runtime_client.rs b/crates/dbx-cli/src/runtime_client.rs deleted file mode 100644 index effa954a0..000000000 --- a/crates/dbx-cli/src/runtime_client.rs +++ /dev/null @@ -1,190 +0,0 @@ -use serde::{Deserialize, Serialize}; -use std::fs::Metadata; -use std::path::PathBuf; - -#[cfg(test)] -pub(crate) static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct RuntimeDiscovery { - pub port: u16, - pub token: String, -} - -pub fn app_data_dir() -> PathBuf { - if let Ok(path) = std::env::var("DBX_APP_DATA_DIR") { - return PathBuf::from(path); - } - - let home = std::env::var(if cfg!(windows) { "APPDATA" } else { "HOME" }).unwrap_or_else(|_| ".".to_string()); - - if cfg!(target_os = "macos") { - PathBuf::from(home).join("Library/Application Support/com.dbx.app") - } else if cfg!(windows) { - PathBuf::from(home).join("com.dbx.app") - } else { - PathBuf::from(home).join(".config/com.dbx.app") - } -} - -pub fn load_runtime() -> Option { - let path = app_data_dir().join("agent-runtime.json"); - let metadata = std::fs::symlink_metadata(&path).ok()?; - if !is_secure_runtime_file(&metadata) { - return None; - } - let json = std::fs::read_to_string(path).ok()?; - serde_json::from_str(&json).ok() -} - -pub async fn get_json(path: &str) -> Result { - let runtime = load_runtime().ok_or_else(|| "runtime unavailable".to_string())?; - let url = runtime_url(&runtime, path, &[])?; - - let response = - reqwest::Client::new().get(url).bearer_auth(runtime.token).send().await.map_err(|err| err.to_string())?; - - let status = response.status(); - if !status.is_success() { - return Err(format!("runtime request failed with status {status}")); - } - - response.json().await.map_err(|err| err.to_string()) -} - -pub async fn get_json_with_query(path: &str, query: &[(&str, String)]) -> Result { - let runtime = load_runtime().ok_or_else(|| "runtime unavailable".to_string())?; - let url = runtime_url(&runtime, path, query)?; - - let response = - reqwest::Client::new().get(url).bearer_auth(runtime.token).send().await.map_err(|err| err.to_string())?; - - let status = response.status(); - if !status.is_success() { - return Err(format!("runtime request failed with status {status}")); - } - - response.json().await.map_err(|err| err.to_string()) -} - -pub async fn post_json(path: &str, body: serde_json::Value) -> Result { - let runtime = load_runtime().ok_or_else(|| "runtime unavailable".to_string())?; - let url = runtime_url(&runtime, path, &[])?; - - let response = reqwest::Client::new() - .post(url) - .bearer_auth(runtime.token) - .json(&body) - .send() - .await - .map_err(|err| err.to_string())?; - - let status = response.status(); - if !status.is_success() { - return Err(format!("runtime request failed with status {status}")); - } - - response.json().await.map_err(|err| err.to_string()) -} - -fn runtime_url(runtime: &RuntimeDiscovery, path: &str, query: &[(&str, String)]) -> Result { - if path.contains('\r') || path.contains('\n') || path.contains("://") { - return Err("invalid runtime path".to_string()); - } - - let mut url = reqwest::Url::parse(&format!("http://127.0.0.1:{}/", runtime.port)).map_err(|err| err.to_string())?; - url.set_path(path.trim_start_matches('/')); - { - let mut pairs = url.query_pairs_mut(); - for (key, value) in query { - pairs.append_pair(key, value); - } - } - Ok(url) -} - -fn is_secure_runtime_file(metadata: &Metadata) -> bool { - if metadata.file_type().is_symlink() || !metadata.is_file() { - return false; - } - runtime_file_owner_and_mode_are_secure(metadata) -} - -#[cfg(unix)] -fn runtime_file_owner_and_mode_are_secure(metadata: &Metadata) -> bool { - use std::os::unix::fs::MetadataExt; - - let owner_only = metadata.mode() & 0o077 == 0; - let owned_by_effective_user = metadata.uid() == unsafe { libc::geteuid() }; - owner_only && owned_by_effective_user -} - -#[cfg(not(unix))] -fn runtime_file_owner_and_mode_are_secure(_metadata: &Metadata) -> bool { - true -} - -#[cfg(test)] -mod tests { - use super::*; - - #[cfg(unix)] - fn set_mode(path: &std::path::Path, mode: u32) { - use std::os::unix::fs::PermissionsExt; - - let mut permissions = std::fs::metadata(path).unwrap().permissions(); - permissions.set_mode(mode); - std::fs::set_permissions(path, permissions).unwrap(); - } - - #[test] - fn load_runtime_rejects_symlink_discovery_file() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let target = dir.path().join("target.json"); - let link = dir.path().join("agent-runtime.json"); - std::fs::write(&target, r#"{"port":4321,"token":"secret"}"#).unwrap(); - #[cfg(unix)] - std::os::unix::fs::symlink(&target, &link).unwrap(); - #[cfg(not(unix))] - std::fs::write(&link, r#"{"port":4321,"token":"secret"}"#).unwrap(); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - - assert!(load_runtime().is_none()); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[cfg(unix)] - #[test] - fn load_runtime_rejects_group_or_world_accessible_discovery_file() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let discovery = dir.path().join("agent-runtime.json"); - std::fs::write(&discovery, r#"{"port":4321,"token":"secret"}"#).unwrap(); - set_mode(&discovery, 0o644); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - - assert!(load_runtime().is_none()); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } - - #[cfg(unix)] - #[test] - fn load_runtime_accepts_owner_only_discovery_file() { - let _guard = ENV_LOCK.lock().unwrap(); - let dir = tempfile::tempdir().unwrap(); - let discovery = dir.path().join("agent-runtime.json"); - std::fs::write(&discovery, r#"{"port":4321,"token":"secret"}"#).unwrap(); - set_mode(&discovery, 0o600); - std::env::set_var("DBX_APP_DATA_DIR", dir.path()); - - let runtime = load_runtime().expect("secure runtime discovery should load"); - assert_eq!(runtime.port, 4321); - assert_eq!(runtime.token, "secret"); - - std::env::remove_var("DBX_APP_DATA_DIR"); - } -} diff --git a/crates/dbx-core/src/cli.rs b/crates/dbx-core/src/cli.rs deleted file mode 100644 index 6296e925a..000000000 --- a/crates/dbx-core/src/cli.rs +++ /dev/null @@ -1,345 +0,0 @@ -use serde::de::{self, IgnoredAny, MapAccess, Visitor}; -use serde::{Deserialize, Deserializer, Serialize}; -use std::fmt; -use std::marker::PhantomData; - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "kebab-case")] -pub enum CliSource { - GuiRuntime, - Headless, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "SCREAMING_SNAKE_CASE")] -pub enum CliErrorCode { - GuiRuntimeRequired, - ConnectionNotFound, - AmbiguousConnection, - SecretUnavailable, - SshTunnelFailed, - QueryClassificationFailed, - HandoffRequired, - DdlBlocked, - ProductionWriteBlocked, - UnsupportedDatabaseType, - Timeout, - InternalError, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct CliError { - pub code: CliErrorCode, - pub message: String, - pub recoverable: bool, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(untagged)] -pub enum CliEnvelope { - Success { ok: bool, source: CliSource, data: T }, - Failure { ok: bool, source: CliSource, error: CliError }, -} - -impl<'de, T> Deserialize<'de> for CliEnvelope -where - T: Deserialize<'de>, -{ - fn deserialize(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - deserializer.deserialize_map(CliEnvelopeVisitor { marker: PhantomData }) - } -} - -struct CliEnvelopeVisitor { - marker: PhantomData, -} - -impl<'de, T> Visitor<'de> for CliEnvelopeVisitor -where - T: Deserialize<'de>, -{ - type Value = CliEnvelope; - - fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { - formatter.write_str("a CLI envelope with consistent ok/data/error fields") - } - - fn visit_map(self, mut map: M) -> Result - where - M: MapAccess<'de>, - { - let mut ok = None; - let mut source = None; - let mut data = None; - let mut has_data = false; - let mut error = None; - let mut has_error = false; - - while let Some(key) = map.next_key::()? { - match key { - CliEnvelopeField::Ok => { - if ok.is_some() { - return Err(de::Error::duplicate_field("ok")); - } - ok = Some(map.next_value()?); - } - CliEnvelopeField::Source => { - if source.is_some() { - return Err(de::Error::duplicate_field("source")); - } - source = Some(map.next_value()?); - } - CliEnvelopeField::Data => { - if has_data { - return Err(de::Error::duplicate_field("data")); - } - has_data = true; - data = Some(map.next_value()?); - } - CliEnvelopeField::Error => { - if has_error { - return Err(de::Error::duplicate_field("error")); - } - has_error = true; - error = Some(map.next_value()?); - } - CliEnvelopeField::Ignore => { - let _ = map.next_value::()?; - } - } - } - - let ok = ok.ok_or_else(|| de::Error::missing_field("ok"))?; - let source = source.ok_or_else(|| de::Error::missing_field("source"))?; - - match (ok, has_data, has_error) { - (true, true, false) => { - Ok(CliEnvelope::Success { ok, source, data: data.expect("data presence was checked") }) - } - (false, false, true) => { - Ok(CliEnvelope::Failure { ok, source, error: error.expect("error presence was checked") }) - } - (true, _, _) => Err(de::Error::custom("ok=true requires data without error")), - (false, _, _) => Err(de::Error::custom("ok=false requires error without data")), - } - } -} - -enum CliEnvelopeField { - Ok, - Source, - Data, - Error, - Ignore, -} - -impl<'de> Deserialize<'de> for CliEnvelopeField { - fn deserialize(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - deserializer.deserialize_identifier(CliEnvelopeFieldVisitor) - } -} - -struct CliEnvelopeFieldVisitor; - -impl<'de> Visitor<'de> for CliEnvelopeFieldVisitor { - type Value = CliEnvelopeField; - - fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { - formatter.write_str("a CLI envelope field") - } - - fn visit_str(self, value: &str) -> Result - where - E: de::Error, - { - Ok(match value { - "ok" => CliEnvelopeField::Ok, - "source" => CliEnvelopeField::Source, - "data" => CliEnvelopeField::Data, - "error" => CliEnvelopeField::Error, - _ => CliEnvelopeField::Ignore, - }) - } -} - -pub fn ok(source: CliSource, data: T) -> CliEnvelope { - CliEnvelope::Success { ok: true, source, data } -} - -pub fn fail(source: CliSource, code: CliErrorCode, message: impl Into, recoverable: bool) -> CliEnvelope { - CliEnvelope::Failure { ok: false, source, error: CliError { code, message: message.into(), recoverable } } -} - -pub fn fail_safe( - source: CliSource, - fallback_code: CliErrorCode, - message: impl AsRef, - recoverable: bool, -) -> CliEnvelope { - let (code, message) = map_safe_error(fallback_code, message.as_ref()); - fail(source, code, message, recoverable) -} - -pub fn map_safe_error(fallback_code: CliErrorCode, message: &str) -> (CliErrorCode, String) { - let lower = message.to_ascii_lowercase(); - if lower.contains("timed out") || lower.contains("timeout") { - return (CliErrorCode::Timeout, "Operation timed out.".to_string()); - } - if lower.contains("connection config not found") || lower == "connection not found" { - return (CliErrorCode::ConnectionNotFound, "Connection not found.".to_string()); - } - if lower.contains("unsupported database") || lower.contains("unsupported database type") { - return (CliErrorCode::UnsupportedDatabaseType, "Unsupported database type.".to_string()); - } - - (fallback_code, "Operation failed. See DBX logs for details.".to_string()) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn serializes_success_source_as_kebab_case() { - let env = ok(CliSource::GuiRuntime, serde_json::json!({"value": 1})); - let json = serde_json::to_string(&env).unwrap(); - assert!(json.contains("\"ok\":true")); - assert!(json.contains("\"source\":\"gui-runtime\"")); - } - - #[test] - fn serializes_error_code_as_screaming_snake_case() { - let env: CliEnvelope<()> = fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, "runtime needed", true); - let json = serde_json::to_string(&env).unwrap(); - assert!(json.contains("\"GUI_RUNTIME_REQUIRED\"")); - } - - #[test] - fn deserializes_success_when_ok_true_and_data_present() { - let env: CliEnvelope = serde_json::from_value(serde_json::json!({ - "ok": true, - "source": "gui-runtime", - "data": { "value": 1 } - })) - .unwrap(); - - match env { - CliEnvelope::Success { ok, source, data } => { - assert!(ok); - assert_eq!(source, CliSource::GuiRuntime); - assert_eq!(data, serde_json::json!({ "value": 1 })); - } - CliEnvelope::Failure { .. } => panic!("expected success envelope"), - } - } - - #[test] - fn deserializes_failure_when_ok_false_and_error_present() { - let env: CliEnvelope = serde_json::from_value(serde_json::json!({ - "ok": false, - "source": "headless", - "error": { - "code": "GUI_RUNTIME_REQUIRED", - "message": "runtime needed", - "recoverable": true - } - })) - .unwrap(); - - match env { - CliEnvelope::Failure { ok, source, error } => { - assert!(!ok); - assert_eq!(source, CliSource::Headless); - assert_eq!(error.code, CliErrorCode::GuiRuntimeRequired); - assert_eq!(error.message, "runtime needed"); - assert!(error.recoverable); - } - CliEnvelope::Success { .. } => panic!("expected failure envelope"), - } - } - - #[test] - fn rejects_success_shape_when_ok_is_false() { - let err = serde_json::from_value::>(serde_json::json!({ - "ok": false, - "source": "headless", - "data": { "runtime": "headless" } - })) - .unwrap_err(); - - assert!(err.to_string().contains("ok=false requires error without data")); - } - - #[test] - fn rejects_failure_shape_when_ok_is_true() { - let err = serde_json::from_value::>(serde_json::json!({ - "ok": true, - "source": "headless", - "error": { - "code": "GUI_RUNTIME_REQUIRED", - "message": "runtime needed", - "recoverable": true - } - })) - .unwrap_err(); - - assert!(err.to_string().contains("ok=true requires data without error")); - } - - #[test] - fn rejects_envelope_with_both_data_and_error() { - let err = serde_json::from_value::>(serde_json::json!({ - "ok": true, - "source": "gui-runtime", - "data": { "value": 1 }, - "error": { - "code": "INTERNAL_ERROR", - "message": "unexpected", - "recoverable": false - } - })) - .unwrap_err(); - - assert!(err.to_string().contains("ok=true requires data without error")); - } - - #[test] - fn maps_timeout_errors_without_leaking_internal_details() { - let (code, message) = map_safe_error( - CliErrorCode::InternalError, - "Query timed out after 30 seconds while reading /Users/alice/private.sqlite", - ); - - assert_eq!(code, CliErrorCode::Timeout); - assert_eq!(message, "Operation timed out."); - assert!(!message.contains("/Users/alice")); - } - - #[test] - fn sanitizes_paths_uri_credentials_and_driver_details() { - let (code, message) = map_safe_error( - CliErrorCode::InternalError, - "SQLite connection failed: Database file does not exist: /Users/alice/private.sqlite; url=mysql://root:secret@localhost/db", - ); - - assert_eq!(code, CliErrorCode::InternalError); - assert_eq!(message, "Operation failed. See DBX logs for details."); - assert!(!message.contains("/Users/alice")); - assert!(!message.contains("root:secret")); - assert!(!message.contains("SQLite connection failed")); - } - - #[test] - fn maps_connection_not_found_to_public_error_code() { - let (code, message) = map_safe_error(CliErrorCode::InternalError, "Connection config not found"); - - assert_eq!(code, CliErrorCode::ConnectionNotFound); - assert_eq!(message, "Connection not found."); - } -} diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 3fa858591..5946f0436 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -249,17 +249,7 @@ pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Resul } pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result { - execute_query_with_row_limit(pool, sql, bare, crate::query::MAX_ROWS).await -} - -pub async fn execute_query_with_row_limit( - pool: &MySqlPool, - sql: &str, - bare: bool, - row_limit: usize, -) -> Result { let start = Instant::now(); - let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN"]) { if bare { @@ -279,14 +269,14 @@ pub async fn execute_query_with_row_limit( .map(|i| mysql_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) .collect(), ); - if result_rows.len() > row_limit { + if result_rows.len() > crate::query::MAX_ROWS { break; } } - let truncated = result_rows.len() > row_limit; + let truncated = result_rows.len() > crate::query::MAX_ROWS; if truncated { - result_rows.truncate(row_limit); + result_rows.truncate(crate::query::MAX_ROWS); } Ok(QueryResult { @@ -311,14 +301,14 @@ pub async fn execute_query_with_row_limit( .map(|i| mysql_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) .collect(), ); - if result_rows.len() > row_limit { + if result_rows.len() > crate::query::MAX_ROWS { break; } } - let truncated = result_rows.len() > row_limit; + let truncated = result_rows.len() > crate::query::MAX_ROWS; if truncated { - result_rows.truncate(row_limit); + result_rows.truncate(crate::query::MAX_ROWS); } Ok(QueryResult { @@ -423,11 +413,6 @@ pub async fn list_triggers(pool: &MySqlPool, database: &str, table: &str) -> Res mod tests { use super::*; - #[allow(dead_code)] - fn mysql_limit_aware_execute_query_api_compiles(pool: &MySqlPool, sql: &str, bare: bool) { - let _ = execute_query_with_row_limit(pool, sql, bare, 7); - } - #[test] fn numeric_metadata_accepts_unsigned_information_schema_values() { assert_eq!(numeric_metadata_u64_to_i32(Some(65)), Some(65)); diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index c7891fe16..bec71c4fd 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -281,12 +281,7 @@ pub async fn get_columns(pool: &PgPool, schema: &str, table: &str) -> Result Result { - execute_query_with_row_limit(pool, sql, crate::query::MAX_ROWS).await -} - -pub async fn execute_query_with_row_limit(pool: &PgPool, sql: &str, row_limit: usize) -> Result { let start = Instant::now(); - let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { let mut stream = sqlx::query(sql).persistent(false).fetch(pool); @@ -306,7 +301,7 @@ pub async fn execute_query_with_row_limit(pool: &PgPool, sql: &str, row_limit: u .map(|i| pg_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) .collect(), ); - if result_rows.len() > row_limit { + if result_rows.len() > crate::query::MAX_ROWS { break; } } @@ -316,9 +311,9 @@ pub async fn execute_query_with_row_limit(pool: &PgPool, sql: &str, row_limit: u columns = desc.columns().iter().map(|c| c.name().to_string()).collect(); } - let truncated = result_rows.len() > row_limit; + let truncated = result_rows.len() > crate::query::MAX_ROWS; if truncated { - result_rows.truncate(row_limit); + result_rows.truncate(crate::query::MAX_ROWS); } Ok(QueryResult { @@ -342,21 +337,11 @@ pub async fn execute_query_with_row_limit(pool: &PgPool, sql: &str, row_limit: u } pub async fn execute_query_with_schema(pool: &PgPool, schema: &str, sql: &str) -> Result { - execute_query_with_schema_and_row_limit(pool, schema, sql, crate::query::MAX_ROWS).await -} - -pub async fn execute_query_with_schema_and_row_limit( - pool: &PgPool, - schema: &str, - sql: &str, - row_limit: usize, -) -> Result { let mut conn = pool.acquire().await.map_err(|e| e.to_string())?; let set_path = format!("SET search_path TO \"{}\", public", schema); sqlx::query(&set_path).execute(&mut *conn).await.map_err(|e| e.to_string())?; let start = Instant::now(); - let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { let mut stream = sqlx::query(sql).persistent(false).fetch(&mut *conn); @@ -376,7 +361,7 @@ pub async fn execute_query_with_schema_and_row_limit( .map(|i| pg_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) .collect(), ); - if result_rows.len() > row_limit { + if result_rows.len() > crate::query::MAX_ROWS { break; } } @@ -387,9 +372,9 @@ pub async fn execute_query_with_schema_and_row_limit( columns = desc.columns().iter().map(|c| c.name().to_string()).collect(); } - let truncated = result_rows.len() > row_limit; + let truncated = result_rows.len() > crate::query::MAX_ROWS; if truncated { - result_rows.truncate(row_limit); + result_rows.truncate(crate::query::MAX_ROWS); } Ok(QueryResult { @@ -461,21 +446,6 @@ pub async fn list_indexes(pool: &PgPool, schema: &str, table: &str) -> Result Result, String> { let rows: Vec = sqlx::query( "SELECT kcu.constraint_name, kcu.column_name, \ diff --git a/crates/dbx-core/src/db/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index f2d668820..a5b9ef618 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -162,16 +162,7 @@ pub async fn list_triggers(pool: &SqlitePool, _schema: &str, table: &str) -> Res } pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result { - execute_query_with_row_limit(pool, sql, crate::query::MAX_ROWS).await -} - -pub async fn execute_query_with_row_limit( - pool: &SqlitePool, - sql: &str, - row_limit: usize, -) -> Result { let start = Instant::now(); - let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "PRAGMA", "EXPLAIN", "WITH"]) { let desc = pool.describe(sql).await.map_err(|e| e.to_string())?; @@ -200,14 +191,14 @@ pub async fn execute_query_with_row_limit( }) .collect(), ); - if result_rows.len() > row_limit { + if result_rows.len() > crate::query::MAX_ROWS { break; } } - let truncated = result_rows.len() > row_limit; + let truncated = result_rows.len() > crate::query::MAX_ROWS; if truncated { - result_rows.truncate(row_limit); + result_rows.truncate(crate::query::MAX_ROWS); } Ok(QueryResult { diff --git a/crates/dbx-core/src/handoff.rs b/crates/dbx-core/src/handoff.rs deleted file mode 100644 index b239647f2..000000000 --- a/crates/dbx-core/src/handoff.rs +++ /dev/null @@ -1,69 +0,0 @@ -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; - -use crate::sql_safety::{OperationClass, RiskLevel}; - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "lowercase")] -pub enum HandoffStatus { - Queued, - Shown, - Approved, - Rejected, - Executed, - Failed, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct HandoffItem { - pub id: String, - pub created_at: DateTime, - pub created_by: String, - #[serde(default)] - pub connection_id: String, - pub connection_name: String, - pub database: Option, - pub title: String, - pub description: Option, - pub sql: String, - pub operation_class: OperationClass, - pub risk_level: RiskLevel, - pub is_production: bool, - pub status: HandoffStatus, - pub result_summary: Option, - pub error: Option, -} - -impl HandoffItem { - pub fn queued( - connection_id: String, - connection_name: String, - database: Option, - title: String, - description: Option, - sql: String, - operation_class: OperationClass, - risk_level: RiskLevel, - is_production: bool, - ) -> Self { - Self { - id: Uuid::new_v4().to_string(), - created_at: Utc::now(), - created_by: "dbx-cli".to_string(), - connection_id, - connection_name, - database, - title, - description, - sql, - operation_class, - risk_level, - is_production, - status: HandoffStatus::Queued, - result_summary: None, - error: None, - } - } -} diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index 78dd76b02..cd34fdf92 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -1,10 +1,8 @@ pub mod ai; -pub mod cli; pub mod connection; pub mod connection_secrets; pub mod db; pub mod external; -pub mod handoff; pub mod history; pub mod models; pub mod mongo_ops; @@ -14,9 +12,7 @@ pub mod query_cancel; pub mod redis_ops; pub mod saved_sql; pub mod schema; -pub mod schema_snapshot; pub mod sql; -pub mod sql_safety; pub mod storage; pub mod table_import; pub mod transfer; diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index d007e9403..56f9c9a28 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -12,16 +12,7 @@ pub const MAX_ROWS: usize = 10000; pub const QUERY_CANCELED: &str = "Query canceled"; pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result { - duckdb_execute_with_row_limit(con, sql, MAX_ROWS) -} - -pub fn duckdb_execute_with_row_limit( - con: &duckdb::Connection, - sql: &str, - row_limit: usize, -) -> Result { let start = std::time::Instant::now(); - let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH", "PRAGMA"]) { let mut stmt = con.prepare(sql).map_err(|e| e.to_string())?; @@ -34,6 +25,9 @@ pub fn duckdb_execute_with_row_limit( let mut result_rows = Vec::new(); while let Some(row) = rows.next().map_err(|e| e.to_string())? { + if result_rows.len() >= MAX_ROWS { + break; + } let vals: Vec = (0..col_count) .map(|i| { row.get::<_, String>(i) @@ -51,15 +45,9 @@ pub fn duckdb_execute_with_row_limit( }) .collect(); result_rows.push(vals); - if result_rows.len() > row_limit { - break; - } } - let truncated = result_rows.len() > row_limit; - if truncated { - result_rows.truncate(row_limit); - } + let truncated = result_rows.len() >= MAX_ROWS; Ok(db::QueryResult { columns, rows: result_rows, @@ -79,14 +67,9 @@ pub fn duckdb_execute_with_row_limit( } } -pub fn truncate_result(result: db::QueryResult) -> db::QueryResult { - truncate_result_with_row_limit(result, MAX_ROWS) -} - -pub fn truncate_result_with_row_limit(mut result: db::QueryResult, row_limit: usize) -> db::QueryResult { - let row_limit = row_limit.max(1); - if result.rows.len() > row_limit { - result.rows.truncate(row_limit); +pub fn truncate_result(mut result: db::QueryResult) -> db::QueryResult { + if result.rows.len() > MAX_ROWS { + result.rows.truncate(MAX_ROWS); result.truncated = true; } result @@ -159,18 +142,6 @@ pub async fn do_execute( schema: Option<&str>, cancel_token: Option, ) -> Result { - do_execute_with_row_limit(state, pool_key, sql, schema, cancel_token, MAX_ROWS).await -} - -pub async fn do_execute_with_row_limit( - state: &AppState, - pool_key: &str, - sql: &str, - schema: Option<&str>, - cancel_token: Option, - row_limit: usize, -) -> Result { - let row_limit = row_limit.max(1); let connections = state.connections.read().await; let pool = connections.get(pool_key).ok_or("Connection not found")?; @@ -182,7 +153,7 @@ pub async fn do_execute_with_row_limit( wait_for_query(cancel_token, async move { let task = tokio::task::spawn_blocking(move || { let con = con.lock().map_err(|e| e.to_string())?; - duckdb_execute_with_row_limit(&con, &sql, row_limit) + duckdb_execute(&con, &sql) }); task.await.map_err(|e| e.to_string())? }) @@ -192,26 +163,22 @@ pub async fn do_execute_with_row_limit( let p = p.clone(); let bare = *mode == crate::connection::MysqlMode::Bare; drop(connections); - wait_for_query(cancel_token, db::mysql::execute_query_with_row_limit(&p, sql, bare, row_limit)).await + wait_for_query(cancel_token, db::mysql::execute_query(&p, sql, bare)).await } PoolKind::Postgres(p) => { let p = p.clone(); let schema = schema.map(|s| s.to_string()); drop(connections); if let Some(schema) = schema { - wait_for_query( - cancel_token, - db::postgres::execute_query_with_schema_and_row_limit(&p, &schema, sql, row_limit), - ) - .await + wait_for_query(cancel_token, db::postgres::execute_query_with_schema(&p, &schema, sql)).await } else { - wait_for_query(cancel_token, db::postgres::execute_query_with_row_limit(&p, sql, row_limit)).await + wait_for_query(cancel_token, db::postgres::execute_query(&p, sql)).await } } PoolKind::Sqlite(p) => { let p = p.clone(); drop(connections); - wait_for_query(cancel_token, db::sqlite::execute_query_with_row_limit(&p, sql, row_limit)).await + wait_for_query(cancel_token, db::sqlite::execute_query(&p, sql)).await } PoolKind::ClickHouse(client) => { let client = client.clone(); @@ -219,7 +186,7 @@ pub async fn do_execute_with_row_limit( drop(connections); wait_for_query(cancel_token, db::clickhouse_driver::execute_query(&client, &database, sql)) .await - .map(|result| truncate_result_with_row_limit(result, row_limit)) + .map(truncate_result) } PoolKind::SqlServer(client) => { let client = client.clone(); @@ -232,9 +199,7 @@ pub async fn do_execute_with_row_limit( }, None => client.lock().await, }; - wait_for_query(cancel_token, db::sqlserver::execute_query(&mut client, sql)) - .await - .map(|result| truncate_result_with_row_limit(result, row_limit)) + wait_for_query(cancel_token, db::sqlserver::execute_query(&mut client, sql)).await.map(truncate_result) } PoolKind::Oracle(pool) => { let client = pool.client(); @@ -253,11 +218,9 @@ pub async fn do_execute_with_row_limit( if let Some(schema) = schema { wait_for_query(cancel_token, db::oracle_driver::execute_query_with_schema(&*client, &schema, sql)) .await - .map(|result| truncate_result_with_row_limit(result, row_limit)) + .map(truncate_result) } else { - wait_for_query(cancel_token, db::oracle_driver::execute_query(&*client, sql)) - .await - .map(|result| truncate_result_with_row_limit(result, row_limit)) + wait_for_query(cancel_token, db::oracle_driver::execute_query(&*client, sql)).await.map(truncate_result) } } PoolKind::Elasticsearch(client) => { @@ -266,7 +229,7 @@ pub async fn do_execute_with_row_limit( drop(connections); wait_for_query(cancel_token, db::elasticsearch_driver::execute_rest_query(&client, &sql)) .await - .map(|result| truncate_result_with_row_limit(result, row_limit)) + .map(truncate_result) } PoolKind::Redis(_) => Err("Use Redis-specific commands".to_string()), PoolKind::MongoDb(_) => Err("Use MongoDB-specific commands".to_string()), @@ -287,7 +250,7 @@ pub async fn do_execute_with_row_limit( task.await.map_err(|e| e.to_string())? }) .await - .map(|result| truncate_result_with_row_limit(result, row_limit)) + .map(truncate_result) } PoolKind::Gaussdb(client) => { let client = client.clone(); @@ -298,7 +261,7 @@ pub async fn do_execute_with_row_limit( db::gaussdb_driver::execute_query(&mut client, &sql).await }) .await - .map(|result| truncate_result_with_row_limit(result, row_limit)) + .map(truncate_result) } PoolKind::ExternalTabular(ext_pool) => { if !starts_with_executable_sql_keyword(sql, &["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA"]) { @@ -310,7 +273,7 @@ pub async fn do_execute_with_row_limit( wait_for_query(cancel_token, async move { let task = tokio::task::spawn_blocking(move || { let con = con.lock().map_err(|e| e.to_string())?; - duckdb_execute_with_row_limit(&con, &sql, row_limit) + duckdb_execute(&con, &sql) }); task.await.map_err(|e| e.to_string())? }) @@ -331,7 +294,7 @@ pub async fn do_execute_with_row_limit( session.invoke::("executeQuery", params).await }) .await - .map(|result| truncate_result_with_row_limit(result, row_limit)) + .map(truncate_result) } } } @@ -344,19 +307,6 @@ pub async fn execute_sql_statement( schema: Option<&str>, cancel_token: Option, ) -> Result { - execute_sql_statement_with_row_limit(state, connection_id, database, sql, schema, cancel_token, MAX_ROWS).await -} - -pub async fn execute_sql_statement_with_row_limit( - state: &AppState, - connection_id: &str, - database: &str, - sql: &str, - schema: Option<&str>, - cancel_token: Option, - row_limit: usize, -) -> Result { - let row_limit = row_limit.max(1); let pool_key = if database.is_empty() { connection_id.to_string() } else { @@ -367,18 +317,16 @@ pub async fn execute_sql_statement_with_row_limit( return Err(canceled_error()); } - let result = do_execute_with_row_limit(state, &pool_key, sql, schema, cancel_token.clone(), row_limit).await; + let result = do_execute(state, &pool_key, sql, schema, cancel_token.clone()).await; - let result = match &result { + match &result { Err(e) if is_connection_error(e) && !is_canceled(&cancel_token) => { let db_opt = if database.is_empty() { None } else { Some(database) }; let new_key = state.reconnect_pool(connection_id, db_opt).await?; - do_execute_with_row_limit(state, &new_key, sql, schema, cancel_token, row_limit).await + do_execute(state, &new_key, sql, schema, cancel_token).await } _ => result, - }; - - result.map(|result| truncate_result_with_row_limit(result, row_limit)) + } } pub async fn execute_multi_core( @@ -798,54 +746,6 @@ async fn exec_tx_none_inner( #[cfg(test)] mod tests { use super::*; - use crate::connection::AppState; - use crate::models::connection::{ConnectionConfig, DatabaseType}; - use crate::storage::Storage; - - fn query_result_with_rows(row_count: usize, truncated: bool) -> db::QueryResult { - db::QueryResult { - columns: vec!["id".to_string()], - rows: (1..=row_count).map(|id| vec![serde_json::Value::Number(id.into())]).collect(), - affected_rows: 0, - execution_time_ms: 0, - truncated, - } - } - - #[test] - fn truncate_result_with_row_limit_truncates_rows_and_marks_result() { - let result = truncate_result_with_row_limit(query_result_with_rows(5, false), 3); - - assert_eq!(result.rows.len(), 3); - assert!(result.truncated); - assert_eq!(result.rows[0][0], serde_json::Value::Number(1.into())); - assert_eq!(result.rows[2][0], serde_json::Value::Number(3.into())); - } - - #[test] - fn truncate_result_with_row_limit_keeps_untruncated_result_when_within_limit() { - let result = truncate_result_with_row_limit(query_result_with_rows(3, false), 3); - - assert_eq!(result.rows.len(), 3); - assert!(!result.truncated); - } - - #[test] - fn truncate_result_with_row_limit_treats_zero_limit_as_one() { - let result = truncate_result_with_row_limit(query_result_with_rows(2, false), 0); - - assert_eq!(result.rows.len(), 1); - assert!(result.truncated); - } - - #[test] - fn mysql_and_postgres_pool_branches_call_limit_aware_drivers() { - let source = include_str!("query.rs"); - - assert!(source.contains("db::mysql::execute_query_with_row_limit(&p, sql, bare, row_limit)")); - assert!(source.contains("db::postgres::execute_query_with_schema_and_row_limit(&p, &schema, sql, row_limit)")); - assert!(source.contains("db::postgres::execute_query_with_row_limit(&p, sql, row_limit)")); - } #[tokio::test] async fn wait_for_query_returns_cancelled_when_token_is_cancelled() { @@ -916,87 +816,4 @@ mod tests { assert!(!is_connection_error("syntax error at position 5")); assert!(!is_connection_error("os error 13")); } - - fn sqlite_config(path: &std::path::Path) -> ConnectionConfig { - ConnectionConfig { - id: "sqlite-id".to_string(), - name: "local-sqlite".to_string(), - db_type: DatabaseType::Sqlite, - driver_profile: None, - driver_label: None, - url_params: None, - host: path.display().to_string(), - port: 0, - username: String::new(), - password: String::new(), - database: None, - color: None, - ssh_enabled: false, - ssh_host: String::new(), - ssh_port: 22, - ssh_user: String::new(), - ssh_password: String::new(), - ssh_key_path: String::new(), - ssh_key_passphrase: String::new(), - ssh_expose_lan: false, - ssh_connect_timeout_secs: crate::models::connection::default_ssh_connect_timeout_secs(), - proxy_enabled: false, - proxy_type: crate::models::connection::ProxyType::Socks5, - proxy_host: String::new(), - proxy_port: 1080, - proxy_username: String::new(), - proxy_password: String::new(), - ssl: false, - sysdba: false, - connection_string: None, - external_config: None, - jdbc_driver_class: None, - jdbc_driver_paths: Vec::new(), - } - } - - async fn sqlite_state_with_rows(row_count: usize) -> (AppState, std::path::PathBuf, std::path::PathBuf) { - let unique = uuid::Uuid::new_v4(); - let data_path = std::env::temp_dir().join(format!("dbx-query-data-{unique}.sqlite")); - let storage_path = std::env::temp_dir().join(format!("dbx-query-storage-{unique}.sqlite")); - std::fs::File::create(&data_path).unwrap(); - let pool = db::sqlite::connect_path(&data_path.display().to_string()).await.unwrap(); - db::sqlite::execute_query(&pool, "CREATE TABLE numbers (id INTEGER PRIMARY KEY)").await.unwrap(); - for id in 1..=row_count { - sqlx::query("INSERT INTO numbers (id) VALUES (?)").bind(id as i64).execute(&pool).await.unwrap(); - } - pool.close().await; - - let storage = Storage::open(&storage_path).await.unwrap(); - let state = AppState::new(storage); - let config = sqlite_config(&data_path); - state.configs.write().await.insert(config.id.clone(), config); - let pool = db::sqlite::connect_path(&data_path.display().to_string()).await.unwrap(); - state.connections.write().await.insert("sqlite-id".to_string(), PoolKind::Sqlite(pool)); - (state, data_path, storage_path) - } - - #[tokio::test] - async fn execute_sql_statement_with_row_limit_caps_sqlite_rows() { - let (state, data_path, storage_path) = sqlite_state_with_rows(200).await; - - let result = execute_sql_statement_with_row_limit( - &state, - "sqlite-id", - "", - "SELECT id FROM numbers ORDER BY id", - None, - None, - 7, - ) - .await - .unwrap(); - - assert!(result.rows.len() <= 7); - assert_eq!(result.rows.len(), 7); - assert!(result.truncated); - - let _ = std::fs::remove_file(data_path); - let _ = std::fs::remove_file(storage_path); - } } diff --git a/crates/dbx-core/src/schema_snapshot.rs b/crates/dbx-core/src/schema_snapshot.rs deleted file mode 100644 index 66e233a4c..000000000 --- a/crates/dbx-core/src/schema_snapshot.rs +++ /dev/null @@ -1,164 +0,0 @@ -use chrono::{DateTime, Utc}; -use futures::{stream, StreamExt, TryStreamExt}; -use serde::{Deserialize, Serialize}; - -use crate::connection::AppState; -use crate::models::connection::{ConnectionConfig, DatabaseType}; -use crate::{schema, types}; - -const TABLE_METADATA_CONCURRENCY: usize = 4; - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct TableSnapshot { - pub name: String, - pub table_type: String, - pub comment: Option, - pub columns: Vec, - pub indexes: Vec, - pub foreign_keys: Vec, - pub triggers: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SchemaSnapshot { - pub connection_id: String, - pub connection_name: String, - pub database: Option, - pub database_type: DatabaseType, - pub driver_profile: Option, - pub captured_at: DateTime, - pub databases: Vec, - pub schemas: Vec, - pub tables: Vec, -} - -pub async fn snapshot( - state: &AppState, - connection_id: &str, - database: Option<&str>, - schema_name: Option<&str>, -) -> Result { - let config = { - let configs = state.configs.read().await; - configs.get(connection_id).cloned().ok_or("Connection config not found")? - }; - - let db = snapshot_database(&config, database)?; - let databases = schema::list_databases_core(state, connection_id) - .await - .map_err(|err| format!("Failed to list databases: {err}"))?; - let schemas = if db.is_empty() { - Vec::new() - } else { - schema::list_schemas_core(state, connection_id, &db) - .await - .map_err(|err| format!("Failed to list schemas for database '{db}': {err}"))? - }; - let effective_schema = schema_name.or_else(|| schemas.first().map(String::as_str)).unwrap_or(""); - let table_infos = if db.is_empty() { - Vec::new() - } else { - schema::list_tables_core(state, connection_id, &db, effective_schema) - .await - .map_err(|err| format!("Failed to list tables for database '{db}' schema '{effective_schema}': {err}"))? - }; - - let tables = collect_table_snapshots(state, connection_id, &db, effective_schema, table_infos).await?; - - Ok(SchemaSnapshot { - connection_id: config.id, - connection_name: config.name, - database: (!db.is_empty()).then_some(db), - database_type: config.db_type, - driver_profile: config.driver_profile, - captured_at: Utc::now(), - databases, - schemas, - tables, - }) -} - -async fn collect_table_snapshots( - state: &AppState, - connection_id: &str, - database: &str, - schema_name: &str, - table_infos: Vec, -) -> Result, String> { - stream::iter(table_infos) - .map(|table| async move { table_snapshot(state, connection_id, database, schema_name, table).await }) - .buffered(TABLE_METADATA_CONCURRENCY) - .try_collect() - .await -} - -async fn table_snapshot( - state: &AppState, - connection_id: &str, - database: &str, - schema_name: &str, - table: types::TableInfo, -) -> Result { - let table_name = table.name.clone(); - let columns = schema::get_columns_core(state, connection_id, database, schema_name, &table_name) - .await - .map_err(|err| format!("Failed to list columns for table '{table_name}': {err}"))?; - let indexes = schema::list_indexes_core(state, connection_id, database, schema_name, &table_name) - .await - .map_err(|err| format!("Failed to list indexes for table '{table_name}': {err}"))?; - let foreign_keys = schema::list_foreign_keys_core(state, connection_id, database, schema_name, &table_name) - .await - .map_err(|err| format!("Failed to list foreign keys for table '{table_name}': {err}"))?; - let triggers = schema::list_triggers_core(state, connection_id, database, schema_name, &table_name) - .await - .map_err(|err| format!("Failed to list triggers for table '{table_name}': {err}"))?; - - Ok(TableSnapshot { - name: table.name, - table_type: table.table_type, - comment: table.comment, - columns, - indexes, - foreign_keys, - triggers, - }) -} - -fn snapshot_database(config: &ConnectionConfig, requested_database: Option<&str>) -> Result { - let database = requested_database - .map(str::trim) - .filter(|database| !database.is_empty()) - .or_else(|| config.effective_database()) - .or_else(|| embedded_default_database(&config.db_type)) - .map(str::to_string); - - match database { - Some(database) => Ok(database), - None if requires_database(&config.db_type) => Err(format!( - "Database is required for schema snapshot for connection '{}' ({:?})", - config.id, config.db_type - )), - None => Ok(String::new()), - } -} - -fn embedded_default_database(db_type: &DatabaseType) -> Option<&'static str> { - match db_type { - DatabaseType::Sqlite | DatabaseType::DuckDb => Some("main"), - _ => None, - } -} - -fn requires_database(db_type: &DatabaseType) -> bool { - matches!( - db_type, - DatabaseType::Mysql - | DatabaseType::Doris - | DatabaseType::StarRocks - | DatabaseType::ClickHouse - | DatabaseType::MongoDb - | DatabaseType::Jdbc - ) -} diff --git a/crates/dbx-core/src/sql_safety.rs b/crates/dbx-core/src/sql_safety.rs deleted file mode 100644 index 09cf4f68a..000000000 --- a/crates/dbx-core/src/sql_safety.rs +++ /dev/null @@ -1,490 +0,0 @@ -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "lowercase")] -pub enum OperationClass { - Read, - Write, - Ddl, - Unknown, -} - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "lowercase")] -pub enum RiskLevel { - Low, - Medium, - High, - Critical, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct RiskMetadata { - pub operation_class: OperationClass, - pub risk_level: RiskLevel, - pub is_production: bool, - pub production_reason: Option, - pub first_token: Option, -} - -#[derive(Debug, Clone, Copy)] -pub struct RiskContext<'a> { - pub connection_name: &'a str, - pub color: Option<&'a str>, - pub environment_label: Option<&'a str>, -} - -impl<'a> RiskContext<'a> { - pub fn new(connection_name: &'a str) -> Self { - Self { connection_name, color: None, environment_label: None } - } - - pub fn with_color(mut self, color: Option<&'a str>) -> Self { - self.color = color; - self - } - - pub fn with_environment_label(mut self, environment_label: Option<&'a str>) -> Self { - self.environment_label = environment_label; - self - } -} - -pub fn classify_sql(sql: &str) -> OperationClass { - let tokens = executable_tokens(sql); - classify_tokens(&tokens) -} - -fn classify_tokens(tokens: &[String]) -> OperationClass { - if tokens.iter().any(|token| is_ddl_token(token)) { - return OperationClass::Ddl; - } - if tokens.iter().any(|token| is_write_token(token)) { - return OperationClass::Write; - } - - match tokens.first().map(String::as_str) { - Some("SELECT" | "SHOW" | "DESCRIBE" | "EXPLAIN" | "WITH") => OperationClass::Read, - _ => OperationClass::Unknown, - } -} - -pub fn risk_for(sql: &str, context: RiskContext<'_>) -> RiskMetadata { - let operation_class = classify_sql(sql); - let (is_production, production_reason) = production_signal(context); - let risk_level = match (operation_class, is_production) { - (OperationClass::Read, _) => RiskLevel::Low, - (OperationClass::Write, _) if has_unfiltered_destructive_write(sql) => RiskLevel::Critical, - (OperationClass::Write, false) => RiskLevel::Medium, - (OperationClass::Write, true) => RiskLevel::High, - (OperationClass::Ddl, _) => RiskLevel::Critical, - (OperationClass::Unknown, _) => RiskLevel::High, - }; - - RiskMetadata { - operation_class, - risk_level, - is_production, - production_reason, - first_token: first_executable_token(sql), - } -} - -pub fn risk_for_connection(sql: &str, connection_name: &str, color: Option<&str>) -> RiskMetadata { - risk_for(sql, RiskContext::new(connection_name).with_color(color)) -} - -fn production_signal(context: RiskContext<'_>) -> (bool, Option) { - if let Some(environment_label) = context.environment_label { - if contains_non_production_signal(environment_label) { - return (false, None); - } - if contains_production_signal(environment_label) { - return (true, Some("environment label".to_string())); - } - } - - if matches!(context.color, Some("#ef4444")) { - return (true, Some("red connection color".to_string())); - } - - if contains_production_signal(context.connection_name) { - return (true, Some("connection name fallback".to_string())); - } - - (false, None) -} - -fn contains_production_signal(value: &str) -> bool { - let value = value.to_ascii_lowercase(); - ["prod", "production", "live"].iter().any(|needle| value.contains(needle)) -} - -fn contains_non_production_signal(value: &str) -> bool { - let value = value.to_ascii_lowercase(); - [ - "dev", - "development", - "test", - "testing", - "qa", - "stage", - "staging", - "local", - "sandbox", - "non-prod", - "non-production", - "non production", - "nonprod", - ] - .iter() - .any(|needle| value.contains(needle)) -} - -fn is_write_token(token: &str) -> bool { - matches!(token, "INSERT" | "UPDATE" | "DELETE" | "MERGE" | "REPLACE") -} - -fn is_ddl_token(token: &str) -> bool { - matches!(token, "CREATE" | "ALTER" | "DROP" | "TRUNCATE" | "RENAME") -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct SqlToken { - text: String, - depth: usize, -} - -fn has_unfiltered_destructive_write(sql: &str) -> bool { - scanned_executable_statements(sql).into_iter().any(|statement| { - statement.iter().enumerate().any(|(index, token)| { - if !matches!(token.text.as_str(), "DELETE" | "UPDATE") { - return false; - } - - !has_same_fragment_boundary(&statement, index, token.depth) - }) - }) -} - -fn has_same_fragment_boundary(statement: &[SqlToken], destructive_write_index: usize, depth: usize) -> bool { - for boundary in &statement[destructive_write_index + 1..] { - if boundary.depth < depth { - break; - } - if boundary.depth == depth && matches!(boundary.text.as_str(), "WHERE" | "LIMIT") { - return true; - } - } - - false -} - -fn executable_tokens(sql: &str) -> Vec { - executable_statements(sql).into_iter().flatten().collect() -} - -fn executable_statements(sql: &str) -> Vec> { - scanned_executable_statements(sql) - .into_iter() - .map(|statement| statement.into_iter().map(|token| token.text).collect()) - .collect() -} - -fn scanned_executable_statements(sql: &str) -> Vec> { - let mut statements = Vec::new(); - let mut current = Vec::new(); - let mut depth = 0; - scan_executable_tokens(sql, &mut current, &mut statements, &mut depth); - push_statement(&mut current, &mut statements); - statements -} - -fn scan_executable_tokens( - sql: &str, - current: &mut Vec, - statements: &mut Vec>, - depth: &mut usize, -) { - let bytes = sql.as_bytes(); - let mut i = 0; - - while i < bytes.len() { - if bytes[i].is_ascii_whitespace() { - i += 1; - continue; - } - - if bytes[i] == b';' { - if *depth == 0 { - push_statement(current, statements); - } - i += 1; - continue; - } - - if bytes[i] == b'(' { - *depth += 1; - i += 1; - continue; - } - - if bytes[i] == b')' { - *depth = depth.saturating_sub(1); - i += 1; - continue; - } - - if i + 1 < bytes.len() && bytes[i] == b'-' && bytes[i + 1] == b'-' { - i += 2; - while i < bytes.len() && bytes[i] != b'\n' { - i += 1; - } - continue; - } - - if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' { - if i + 2 < bytes.len() && bytes[i + 2] == b'!' { - let content_start = i + 3; - let content_end = block_comment_end(bytes, content_start); - scan_executable_tokens(&sql[content_start..content_end], current, statements, depth); - i = (content_end + 2).min(bytes.len()); - } else { - i += 2; - while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { - i += 1; - } - i = (i + 2).min(bytes.len()); - } - continue; - } - - if let Some(delimiter_len) = dollar_quote_delimiter_len(bytes, i) { - let delimiter = &sql[i..i + delimiter_len]; - i += delimiter_len; - if let Some(end) = sql[i..].find(delimiter) { - i += end + delimiter_len; - } else { - i = bytes.len(); - } - continue; - } - - if matches!(bytes[i], b'\'' | b'"' | b'`') { - let quote = bytes[i]; - i += 1; - while i < bytes.len() { - if bytes[i] == quote { - if i + 1 < bytes.len() && bytes[i + 1] == quote { - i += 2; - continue; - } - i += 1; - break; - } - i += 1; - } - continue; - } - - if bytes[i].is_ascii_alphabetic() || bytes[i] == b'_' { - let start = i; - i += 1; - while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') { - i += 1; - } - current.push(SqlToken { text: sql[start..i].to_ascii_uppercase(), depth: *depth }); - continue; - } - - i += 1; - } -} - -fn push_statement(current: &mut Vec, statements: &mut Vec>) { - if !current.is_empty() { - statements.push(std::mem::take(current)); - } -} - -fn block_comment_end(bytes: &[u8], mut i: usize) -> usize { - while i + 1 < bytes.len() { - if bytes[i] == b'*' && bytes[i + 1] == b'/' { - return i; - } - i += 1; - } - bytes.len() -} - -fn dollar_quote_delimiter_len(bytes: &[u8], start: usize) -> Option { - if bytes.get(start) != Some(&b'$') { - return None; - } - - let mut i = start + 1; - if bytes.get(i) == Some(&b'$') { - return Some(2); - } - - if !bytes.get(i).is_some_and(|byte| byte.is_ascii_alphabetic() || *byte == b'_') { - return None; - } - - i += 1; - while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') { - i += 1; - } - - (bytes.get(i) == Some(&b'$')).then_some(i - start + 1) -} - -fn first_executable_token(sql: &str) -> Option { - executable_tokens(sql).into_iter().next() -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn comments_do_not_hide_read_token() { - assert_eq!(classify_sql("-- comment\nSELECT 1"), OperationClass::Read); - assert_eq!(classify_sql("/* DROP TABLE x */ SELECT 1"), OperationClass::Read); - } - - #[test] - fn classifies_write_and_ddl() { - assert_eq!(classify_sql("update users set name = 'a'"), OperationClass::Write); - assert_eq!(classify_sql("DROP TABLE users"), OperationClass::Ddl); - } - - #[test] - fn with_does_not_hide_write_or_ddl() { - assert_eq!( - classify_sql("WITH moved AS (DELETE FROM orders RETURNING *) SELECT * FROM moved"), - OperationClass::Write - ); - assert_eq!(classify_sql("WITH dropped AS (DROP TABLE old_orders) SELECT 1"), OperationClass::Ddl); - } - - #[test] - fn explain_analyze_write_is_write() { - assert_eq!(classify_sql("EXPLAIN ANALYZE UPDATE users SET name = 'a'"), OperationClass::Write); - } - - #[test] - fn dangerous_statement_in_multi_statement_sql_is_not_read() { - assert_eq!(classify_sql("SELECT * FROM users; DELETE FROM users WHERE id = 1"), OperationClass::Write); - assert_eq!(classify_sql("SHOW TABLES; DROP TABLE users"), OperationClass::Ddl); - } - - #[test] - fn red_color_marks_production() { - let risk = risk_for_connection("SELECT * FROM orders", "prod-main", Some("#ef4444")); - assert!(risk.is_production); - assert_eq!(risk.risk_level, RiskLevel::Low); - } - - #[test] - fn environment_label_marks_production() { - let risk = risk_for( - "UPDATE orders SET status = 'done' WHERE id = 1", - RiskContext { connection_name: "analytics", color: None, environment_label: Some("Production") }, - ); - assert!(risk.is_production); - assert_eq!(risk.production_reason.as_deref(), Some("environment label")); - assert_eq!(risk.risk_level, RiskLevel::High); - } - - #[test] - fn environment_label_overrides_color_and_name_fallback() { - let non_prod_label = risk_for( - "SELECT * FROM orders", - RiskContext { connection_name: "prod-main", color: Some("#ef4444"), environment_label: Some("Staging") }, - ); - assert!(!non_prod_label.is_production); - assert_eq!(non_prod_label.production_reason, None); - - let prod_label = risk_for( - "SELECT * FROM orders", - RiskContext { connection_name: "analytics", color: Some("#22c55e"), environment_label: Some("Production") }, - ); - assert!(prod_label.is_production); - assert_eq!(prod_label.production_reason.as_deref(), Some("environment label")); - } - - #[test] - fn destructive_writes_without_where_or_limit_are_critical() { - assert_eq!(risk_for("DELETE FROM users", RiskContext::new("dev")).risk_level, RiskLevel::Critical); - assert_eq!( - risk_for("UPDATE users SET active = false", RiskContext::new("dev")).risk_level, - RiskLevel::Critical - ); - assert_eq!(risk_for("DELETE FROM users WHERE id = 1", RiskContext::new("dev")).risk_level, RiskLevel::Medium); - assert_eq!(risk_for("TRUNCATE TABLE users", RiskContext::new("dev")).risk_level, RiskLevel::Critical); - } - - #[test] - fn destructive_writes_only_count_top_level_where_or_limit_as_boundaries() { - assert_eq!( - risk_for( - "DELETE FROM users USING (SELECT id FROM archived WHERE stale = true) old", - RiskContext::new("dev") - ) - .risk_level, - RiskLevel::Critical - ); - assert_eq!( - risk_for( - "UPDATE users SET active = false FROM (SELECT id FROM flags LIMIT 10) flags", - RiskContext::new("dev") - ) - .risk_level, - RiskLevel::Critical - ); - assert_eq!( - risk_for( - "DELETE FROM users WHERE id IN (SELECT user_id FROM archived WHERE stale = true)", - RiskContext::new("dev") - ) - .risk_level, - RiskLevel::Medium - ); - } - - #[test] - fn cte_destructive_writes_do_not_use_sibling_cte_boundaries() { - assert_eq!( - risk_for( - "WITH deleted AS (DELETE FROM users RETURNING id), scoped AS (SELECT id FROM audit WHERE id = 1) SELECT * FROM scoped", - RiskContext::new("dev") - ) - .risk_level, - RiskLevel::Critical - ); - assert_eq!( - risk_for( - "WITH updated AS (UPDATE users SET active = false RETURNING id), scoped AS (SELECT id FROM audit LIMIT 1) SELECT * FROM scoped", - RiskContext::new("dev") - ) - .risk_level, - RiskLevel::Critical - ); - } - - #[test] - fn postgresql_dollar_quotes_do_not_contribute_tokens() { - assert_eq!(classify_sql("SELECT $$ DELETE FROM users $$"), OperationClass::Read); - assert_eq!(classify_sql("SELECT $tag$ DROP TABLE users $tag$"), OperationClass::Read); - } - - #[test] - fn mysql_executable_comment_contributes_tokens() { - assert_eq!(classify_sql("/*!50000 DELETE FROM users */ SELECT 1"), OperationClass::Write); - let risk = risk_for("/*! UPDATE users SET active = false */", RiskContext::new("dev")); - assert_eq!(risk.operation_class, OperationClass::Write); - assert_eq!(risk.first_token.as_deref(), Some("UPDATE")); - } -} diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index dddf77d36..27d588737 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -4,7 +4,6 @@ use std::str::FromStr; use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions}; use crate::ai::{AiChatMessage, AiConfig, AiConversation}; -use crate::handoff::{HandoffItem, HandoffStatus}; use crate::history::HistoryEntry; use crate::models::connection::ConnectionConfig; use crate::saved_sql::{SavedSqlFile, SavedSqlFolder, SavedSqlLibrary}; @@ -85,13 +84,6 @@ const SCHEMA_STATEMENTS: &[&str] = &[ created_at TEXT NOT NULL DEFAULT '', updated_at TEXT NOT NULL DEFAULT '' )", - "CREATE TABLE IF NOT EXISTS handoffs ( - seq INTEGER PRIMARY KEY AUTOINCREMENT, - id TEXT NOT NULL UNIQUE, - payload_json TEXT NOT NULL, - status TEXT NOT NULL, - created_at TEXT NOT NULL - )", ]; // --------------------------------------------------------------------------- @@ -109,7 +101,6 @@ impl Storage { sqlx::query(statement).execute(&pool).await.map_err(|e| e.to_string())?; } ensure_history_columns(&pool).await?; - ensure_handoffs_sequence(&pool).await?; Ok(Self { db: pool }) } @@ -143,47 +134,6 @@ async fn ensure_history_columns(pool: &SqlitePool) -> Result<(), String> { Ok(()) } -async fn ensure_handoffs_sequence(pool: &SqlitePool) -> Result<(), String> { - let (seq_columns,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM pragma_table_info('handoffs') WHERE name = 'seq'") - .fetch_one(pool) - .await - .map_err(|e| e.to_string())?; - - if seq_columns > 0 { - return Ok(()); - } - - let mut tx = pool.begin().await.map_err(|e| e.to_string())?; - sqlx::query( - "CREATE TABLE handoffs_migration ( - seq INTEGER PRIMARY KEY AUTOINCREMENT, - id TEXT NOT NULL UNIQUE, - payload_json TEXT NOT NULL, - status TEXT NOT NULL, - created_at TEXT NOT NULL - )", - ) - .execute(&mut *tx) - .await - .map_err(|e| e.to_string())?; - sqlx::query( - "INSERT INTO handoffs_migration (id, payload_json, status, created_at) - SELECT id, payload_json, status, created_at - FROM handoffs - ORDER BY created_at ASC, rowid ASC", - ) - .execute(&mut *tx) - .await - .map_err(|e| e.to_string())?; - sqlx::query("DROP TABLE handoffs").execute(&mut *tx).await.map_err(|e| e.to_string())?; - sqlx::query("ALTER TABLE handoffs_migration RENAME TO handoffs") - .execute(&mut *tx) - .await - .map_err(|e| e.to_string())?; - tx.commit().await.map_err(|e| e.to_string())?; - Ok(()) -} - // --------------------------------------------------------------------------- // History // --------------------------------------------------------------------------- @@ -292,89 +242,6 @@ impl Storage { } } -// --------------------------------------------------------------------------- -// Handoffs -// --------------------------------------------------------------------------- - -impl Storage { - pub async fn save_handoff(&self, item: &HandoffItem) -> Result<(), String> { - let json = serde_json::to_string(item).map_err(|e| e.to_string())?; - let status = handoff_status_value(&item.status)?; - - sqlx::query( - "INSERT INTO handoffs (id, payload_json, status, created_at) \ - VALUES (?, ?, ?, ?) \ - ON CONFLICT(id) DO UPDATE SET \ - payload_json = excluded.payload_json, \ - status = excluded.status, \ - created_at = excluded.created_at", - ) - .bind(&item.id) - .bind(json) - .bind(status) - .bind(item.created_at.to_rfc3339()) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; - - Ok(()) - } - - pub async fn update_handoff_status(&self, id: &str, status: HandoffStatus) -> Result { - let from_statuses = allowed_handoff_status_transitions(&status); - if from_statuses.is_empty() { - return Ok(false); - } - - let status_value = handoff_status_value(&status)?; - - let result = sqlx::query( - "UPDATE handoffs SET payload_json = json_set(payload_json, '$.status', ?), status = ? \ - WHERE id = ? AND status IN (SELECT value FROM json_each(?))", - ) - .bind(status_value.clone()) - .bind(status_value) - .bind(id) - .bind(serde_json::to_string(&from_statuses).map_err(|e| e.to_string())?) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; - - Ok(result.rows_affected() > 0) - } - - pub async fn load_pending_handoffs(&self) -> Result, String> { - let rows: Vec<(String,)> = sqlx::query_as( - "SELECT payload_json FROM handoffs \ - WHERE status IN ('queued', 'shown') \ - ORDER BY created_at ASC, seq ASC", - ) - .fetch_all(&self.db) - .await - .map_err(|e| e.to_string())?; - - rows.into_iter().map(|(json,)| serde_json::from_str(&json).map_err(|e| e.to_string())).collect() - } -} - -fn handoff_status_value(status: &HandoffStatus) -> Result { - serde_json::to_value(status) - .ok() - .and_then(|value| value.as_str().map(str::to_string)) - .ok_or_else(|| "Failed to serialize handoff status".to_string()) -} - -fn allowed_handoff_status_transitions(status: &HandoffStatus) -> &'static [&'static str] { - match status { - HandoffStatus::Queued => &[], - HandoffStatus::Shown => &["queued", "shown"], - HandoffStatus::Approved => &["queued", "shown", "approved"], - HandoffStatus::Rejected => &["queued", "shown"], - HandoffStatus::Executed => &["approved", "executed"], - HandoffStatus::Failed => &["approved", "executed", "failed"], - } -} - // --------------------------------------------------------------------------- // AI Config // --------------------------------------------------------------------------- @@ -1031,168 +898,6 @@ impl Storage { } } -#[cfg(test)] -mod handoff_tests { - use super::*; - use crate::handoff::{HandoffItem, HandoffStatus}; - use crate::sql_safety::{OperationClass, RiskLevel}; - - async fn open_temp_storage() -> Storage { - let path = std::env::temp_dir().join(format!("dbx-handoff-test-{}.db", uuid::Uuid::new_v4())); - Storage::open(&path).await.unwrap() - } - - fn queued_handoff(title: &str) -> HandoffItem { - HandoffItem::queued( - "prod-main-id".to_string(), - "prod-main".to_string(), - Some("app".to_string()), - title.to_string(), - Some("review write".to_string()), - "UPDATE users SET active = 0".to_string(), - OperationClass::Write, - RiskLevel::High, - true, - ) - } - - #[tokio::test] - async fn save_handoff_loads_pending_records_in_created_order() { - let storage = open_temp_storage().await; - let first = queued_handoff("first"); - let mut second = queued_handoff("second"); - second.created_at = first.created_at + chrono::Duration::seconds(1); - second.status = HandoffStatus::Shown; - - storage.save_handoff(&second).await.unwrap(); - storage.save_handoff(&first).await.unwrap(); - - let loaded = storage.load_pending_handoffs().await.unwrap(); - - assert_eq!(loaded.iter().map(|item| item.title.as_str()).collect::>(), vec!["first", "second"]); - assert_eq!(loaded[0].status, HandoffStatus::Queued); - assert_eq!(loaded[1].status, HandoffStatus::Shown); - assert_eq!(loaded[0].operation_class, OperationClass::Write); - assert!(loaded[0].is_production); - } - - #[tokio::test] - async fn load_pending_handoffs_keeps_fifo_order_for_matching_created_at() { - let storage = open_temp_storage().await; - let first = queued_handoff("first"); - let mut second = queued_handoff("second"); - second.created_at = first.created_at; - - storage.save_handoff(&first).await.unwrap(); - storage.save_handoff(&second).await.unwrap(); - - let loaded = storage.load_pending_handoffs().await.unwrap(); - - assert_eq!(loaded.iter().map(|item| item.title.as_str()).collect::>(), vec!["first", "second"]); - } - - #[tokio::test] - async fn handoffs_table_has_stable_autoincrement_sequence() { - let storage = open_temp_storage().await; - - let columns: Vec<(String,)> = sqlx::query_as("SELECT name FROM pragma_table_info('handoffs')") - .fetch_all(&storage.db) - .await - .unwrap(); - - assert!(columns.iter().any(|(name,)| name == "seq")); - } - - #[tokio::test] - async fn handoff_item_serializes_connection_id_and_display_name() { - let item = queued_handoff("serialize"); - - let value = serde_json::to_value(&item).unwrap(); - - assert_eq!(item.connection_id, "prod-main-id"); - assert_eq!(value["connectionId"], "prod-main-id"); - assert_eq!(value["connectionName"], "prod-main"); - } - - #[tokio::test] - async fn load_pending_handoffs_excludes_terminal_statuses() { - let storage = open_temp_storage().await; - let queued = queued_handoff("queued"); - let mut executed = queued_handoff("executed"); - executed.status = HandoffStatus::Executed; - - storage.save_handoff(&queued).await.unwrap(); - storage.save_handoff(&executed).await.unwrap(); - - let loaded = storage.load_pending_handoffs().await.unwrap(); - - assert_eq!(loaded.len(), 1); - assert_eq!(loaded[0].id, queued.id); - } - - #[tokio::test] - async fn update_handoff_status_updates_payload_and_pending_visibility() { - let storage = open_temp_storage().await; - let item = queued_handoff("review me"); - let queued_reject = queued_handoff("reject from queued"); - let missing = uuid::Uuid::new_v4().to_string(); - - storage.save_handoff(&item).await.unwrap(); - storage.save_handoff(&queued_reject).await.unwrap(); - - assert!(storage.update_handoff_status(&item.id, HandoffStatus::Shown).await.unwrap()); - assert!(!storage.update_handoff_status(&missing, HandoffStatus::Rejected).await.unwrap()); - assert!(storage.update_handoff_status(&queued_reject.id, HandoffStatus::Rejected).await.unwrap()); - - let shown = storage.load_pending_handoffs().await.unwrap(); - assert_eq!(shown.len(), 1); - assert_eq!(shown[0].status, HandoffStatus::Shown); - - assert!(storage.update_handoff_status(&item.id, HandoffStatus::Rejected).await.unwrap()); - assert!(storage.load_pending_handoffs().await.unwrap().is_empty()); - } - - #[tokio::test] - async fn update_handoff_status_does_not_let_shown_overwrite_rejected() { - let storage = open_temp_storage().await; - let item = queued_handoff("reject wins"); - - storage.save_handoff(&item).await.unwrap(); - - assert!(storage.update_handoff_status(&item.id, HandoffStatus::Rejected).await.unwrap()); - assert!(!storage.update_handoff_status(&item.id, HandoffStatus::Shown).await.unwrap()); - - assert!(storage.load_pending_handoffs().await.unwrap().is_empty()); - } - - #[tokio::test] - async fn update_handoff_status_only_marks_shown_from_queued_or_shown() { - let storage = open_temp_storage().await; - let queued = queued_handoff("queued"); - let mut shown = queued_handoff("shown"); - shown.status = HandoffStatus::Shown; - let mut approved = queued_handoff("approved"); - approved.status = HandoffStatus::Approved; - let mut executed = queued_handoff("executed"); - executed.status = HandoffStatus::Executed; - let mut failed = queued_handoff("failed"); - failed.status = HandoffStatus::Failed; - - for item in [&queued, &shown, &approved, &executed, &failed] { - storage.save_handoff(item).await.unwrap(); - } - - assert!(storage.update_handoff_status(&queued.id, HandoffStatus::Shown).await.unwrap()); - assert!(storage.update_handoff_status(&shown.id, HandoffStatus::Shown).await.unwrap()); - assert!(!storage.update_handoff_status(&approved.id, HandoffStatus::Shown).await.unwrap()); - assert!(!storage.update_handoff_status(&executed.id, HandoffStatus::Shown).await.unwrap()); - assert!(!storage.update_handoff_status(&failed.id, HandoffStatus::Shown).await.unwrap()); - - let loaded = storage.load_pending_handoffs().await.unwrap(); - assert_eq!(loaded.iter().map(|item| item.id.as_str()).collect::>(), vec![queued.id, shown.id]); - } -} - // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- diff --git a/crates/dbx-core/tests/schema_snapshot.rs b/crates/dbx-core/tests/schema_snapshot.rs deleted file mode 100644 index bed6036bf..000000000 --- a/crates/dbx-core/tests/schema_snapshot.rs +++ /dev/null @@ -1,140 +0,0 @@ -use std::str::FromStr; - -use dbx_core::connection::{AppState, PoolKind}; -use dbx_core::models::connection::{default_ssh_connect_timeout_secs, ConnectionConfig, DatabaseType}; -use dbx_core::schema_snapshot::snapshot; -use dbx_core::storage::Storage; -use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; - -fn sqlite_config(path: &std::path::Path) -> ConnectionConfig { - ConnectionConfig { - id: "sqlite-fixture".to_string(), - name: "SQLite Fixture".to_string(), - db_type: DatabaseType::Sqlite, - driver_profile: Some("builtin-sqlite".to_string()), - driver_label: None, - url_params: None, - host: path.display().to_string(), - port: 0, - username: String::new(), - password: String::new(), - database: None, - color: None, - ssh_enabled: false, - ssh_host: String::new(), - ssh_port: 22, - ssh_user: String::new(), - ssh_password: String::new(), - ssh_key_path: String::new(), - ssh_key_passphrase: String::new(), - ssh_expose_lan: false, - ssh_connect_timeout_secs: default_ssh_connect_timeout_secs(), - proxy_enabled: false, - proxy_type: dbx_core::models::connection::ProxyType::Socks5, - proxy_host: String::new(), - proxy_port: 1080, - proxy_username: String::new(), - proxy_password: String::new(), - ssl: false, - sysdba: false, - connection_string: None, - external_config: None, - jdbc_driver_class: None, - jdbc_driver_paths: Vec::new(), - } -} - -async fn create_sqlite_fixture(path: &std::path::Path) { - let url = format!("sqlite:{}?mode=rwc", path.display()); - let options = SqliteConnectOptions::from_str(&url).unwrap().create_if_missing(true); - let pool = SqlitePoolOptions::new().max_connections(1).connect_with(options).await.unwrap(); - - sqlx::query("PRAGMA foreign_keys = ON").execute(&pool).await.unwrap(); - sqlx::query("CREATE TABLE teams (id INTEGER PRIMARY KEY, name TEXT NOT NULL)").execute(&pool).await.unwrap(); - sqlx::query( - "CREATE TABLE users (id INTEGER PRIMARY KEY, team_id INTEGER NOT NULL, email TEXT NOT NULL UNIQUE, FOREIGN KEY(team_id) REFERENCES teams(id))", - ) - .execute(&pool) - .await - .unwrap(); - sqlx::query("CREATE INDEX idx_users_team_id ON users(team_id)").execute(&pool).await.unwrap(); - sqlx::query("CREATE TRIGGER trg_users_ai AFTER INSERT ON users BEGIN SELECT 1; END").execute(&pool).await.unwrap(); - sqlx::query("CREATE VIEW active_users AS SELECT id, email FROM users").execute(&pool).await.unwrap(); - - pool.close().await; -} - -async fn open_state() -> AppState { - let storage_path = std::env::temp_dir().join(format!("dbx-schema-snapshot-storage-{}.db", uuid::Uuid::new_v4())); - AppState::new(Storage::open(&storage_path).await.unwrap()) -} - -#[tokio::test] -async fn snapshot_standardizes_sqlite_tables_views_and_metadata() { - let data_path = std::env::temp_dir().join(format!("dbx-schema-snapshot-data-{}.db", uuid::Uuid::new_v4())); - create_sqlite_fixture(&data_path).await; - - let state = open_state().await; - let config = sqlite_config(&data_path); - let pool = dbx_core::db::sqlite::connect_path(&data_path.display().to_string()).await.unwrap(); - state.configs.write().await.insert(config.id.clone(), config.clone()); - state.connections.write().await.insert(config.id.clone(), PoolKind::Sqlite(pool)); - - let snapshot = snapshot(&state, &config.id, None, None).await.unwrap(); - - assert_eq!(snapshot.connection_id, "sqlite-fixture"); - assert_eq!(snapshot.connection_name, "SQLite Fixture"); - assert_eq!(snapshot.database.as_deref(), Some("main")); - assert_eq!(snapshot.database_type, DatabaseType::Sqlite); - assert_eq!(snapshot.driver_profile.as_deref(), Some("builtin-sqlite")); - let now = chrono::Utc::now(); - assert!(snapshot.captured_at <= now); - assert!(snapshot.captured_at > now - chrono::Duration::seconds(5)); - assert_eq!(snapshot.databases.iter().map(|db| db.name.as_str()).collect::>(), vec!["main"]); - - let table_names = snapshot.tables.iter().map(|table| table.name.as_str()).collect::>(); - assert_eq!(table_names, vec!["active_users", "teams", "users"]); - assert!(serde_json::to_value(&snapshot).unwrap().get("views").is_none()); - - let users = snapshot.tables.iter().find(|table| table.name == "users").unwrap(); - assert_eq!(users.table_type, "BASE TABLE"); - assert!(users.columns.iter().any(|column| column.name == "email" && !column.is_nullable)); - assert!(users.indexes.iter().any(|index| index.name == "idx_users_team_id" && index.columns == vec!["team_id"])); - assert!(users.foreign_keys.iter().any(|fk| fk.column == "team_id" && fk.ref_table == "teams")); - assert!(users.triggers.iter().any(|trigger| trigger.name == "trg_users_ai" && trigger.event == "INSERT")); - - let active_users = snapshot.tables.iter().find(|table| table.name == "active_users").unwrap(); - assert_eq!(active_users.table_type, "VIEW"); - assert!(active_users.columns.iter().any(|column| column.name == "email")); -} - -#[tokio::test] -async fn snapshot_propagates_schema_core_errors() { - let data_path = std::env::temp_dir().join(format!("dbx-schema-snapshot-missing-{}.db", uuid::Uuid::new_v4())); - let state = open_state().await; - let config = sqlite_config(&data_path); - state.configs.write().await.insert(config.id.clone(), config.clone()); - - let err = snapshot(&state, &config.id, None, None).await.unwrap_err(); - - assert!(err.contains("Failed to list databases"), "{err}"); - assert!(err.contains("Connection not found"), "{err}"); -} - -#[tokio::test] -async fn snapshot_requires_database_for_database_scoped_connections_without_default() { - let state = open_state().await; - let mut config = sqlite_config(std::path::Path::new("unused")); - config.id = "mysql-no-db".to_string(); - config.name = "MySQL without DB".to_string(); - config.db_type = DatabaseType::Mysql; - config.driver_profile = None; - config.host = "127.0.0.1".to_string(); - config.port = 3306; - state.configs.write().await.insert(config.id.clone(), config.clone()); - - let err = snapshot(&state, &config.id, None, None).await.unwrap_err(); - - assert!(err.contains("Database is required"), "{err}"); - assert!(err.contains("mysql-no-db"), "{err}"); -} diff --git a/docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md b/docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md deleted file mode 100644 index e3cd3baf7..000000000 --- a/docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md +++ /dev/null @@ -1,1537 +0,0 @@ -# DBX CLI Runtime Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Add a full DBX CLI with independent GUI Runtime support for the 8 specified commands while preserving the existing MCP package and behavior. - -**Architecture:** Add a new `dbx-cli` Rust binary crate that first attempts GUI Runtime discovery, then falls back to headless mode. Shared CLI data models, JSON envelope, SQL safety classification, schema snapshot, risk metadata, and handoff queue live in `dbx-core`. The desktop app starts a separate authenticated localhost runtime and the Vue UI synchronizes active context, selection, result samples, and handoff review state into that runtime. - -**Tech Stack:** Rust workspace, `dbx-core`, Tauri v2, Vue 3/Pinia, SQLite storage via `sqlx`, existing DBX schema/query/SSH core, Node test runner for frontend utilities, Cargo tests for Rust. - ---- - -## Non-Negotiable Constraints - -- Preserve all existing `mcp/` files, package behavior, tool names, and bridge endpoints. -- Do not remove or rename current Tauri commands. -- Do not add CLI commands that create, edit, or delete DBX connections. -- All CLI agent-facing commands return the envelope shape `{ ok, source, data }` or `{ ok, source, error }`. -- `safe-query` directly executes only READ SQL; WRITE and DDL return structured blocking errors and point to `handoff`. -- `selection` and `result current` return real GUI state when runtime is available, not a placeholder. -- `handoff` sends to GUI runtime when running and writes queued pending records when headless. - -## File Structure - -- Create `crates/dbx-cli/Cargo.toml`: binary crate configuration. -- Create `crates/dbx-cli/src/main.rs`: argument parsing, command dispatch, runtime/headless selection. -- Create `crates/dbx-cli/src/commands.rs`: CLI command handlers. -- Create `crates/dbx-cli/src/runtime_client.rs`: discovery file loading, token-authenticated HTTP calls. -- Create `crates/dbx-core/src/cli.rs`: envelope, error codes, CLI DTOs, app data path helper. -- Create `crates/dbx-core/src/sql_safety.rs`: SQL classification and risk metadata. -- Create `crates/dbx-core/src/schema_snapshot.rs`: normalized schema snapshot orchestration. -- Create `crates/dbx-core/src/handoff.rs`: handoff model and queued storage helpers. -- Modify `crates/dbx-core/src/lib.rs`: export new modules. -- Modify `crates/dbx-core/src/storage.rs`: add pending handoff table migration and load/save helpers. -- Modify root `Cargo.toml`: add `crates/dbx-cli` workspace member. -- Modify `src-tauri/src/lib.rs`: start independent agent runtime and register runtime commands. -- Create `src-tauri/src/commands/agent_runtime.rs`: authenticated runtime server and shared runtime state. -- Modify `src-tauri/src/commands/mod.rs`: expose `agent_runtime`. -- Modify `src/lib/api.ts` and `src/lib/tauri.ts`: add runtime state sync and handoff API wrappers. -- Create `src/stores/agentRuntimeStore.ts`: collect current UI context and push snapshots. -- Create `src/components/agent/AgentHandoffDialog.vue`: review queued/runtime handoffs. -- Modify `src/components/layout/AppDialogs.vue`: mount handoff dialog. -- Modify `src/stores/queryStore.ts`: notify runtime store on active tab/result changes. -- Modify `src/components/grid/DataGrid.vue`: notify runtime store on grid selection changes. -- Modify `src/composables/useTauriEvents.ts`: keep existing MCP listeners unchanged. -- Add Rust tests under `crates/dbx-core/src/sql_safety.rs` and `crates/dbx-core/src/cli.rs`. -- Add frontend tests under `tests/agentRuntimeStore.test.ts` for payload shaping. - ---- - -### Task 1: Workspace and CLI Crate Skeleton - -**Files:** -- Modify: `/Users/bytedance/open_source_poj/dbx/Cargo.toml` -- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/Cargo.toml` -- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/main.rs` - -- [ ] **Step 1: Add workspace member** - -Change the workspace members in `/Users/bytedance/open_source_poj/dbx/Cargo.toml` to include the CLI crate: - -```toml -[workspace] -resolver = "2" -members = ["src-tauri", "crates/dbx-core", "crates/dbx-cli", "src-web"] -``` - -- [ ] **Step 2: Create CLI crate manifest** - -Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/Cargo.toml`: - -```toml -[package] -name = "dbx-cli" -version = "0.5.2" -edition = "2021" - -[[bin]] -name = "dbx-cli" -path = "src/main.rs" - -[dependencies] -dbx-core = { path = "../dbx-core" } -chrono = { version = "0.4", features = ["serde"] } -serde = { version = "1.0", features = ["derive"] } -serde_json = "1.0" -tokio = { version = "1", features = ["full"] } -reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } -uuid = { version = "1", features = ["v4", "serde"] } -``` - -- [ ] **Step 3: Create placeholder CLI entry** - -Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/main.rs`: - -```rust -mod commands; -mod runtime_client; - -#[tokio::main] -async fn main() { - if let Err(err) = commands::run(std::env::args().skip(1).collect()).await { - println!("{}", serde_json::to_string_pretty(&err).unwrap_or_else(|_| "{\"ok\":false}".to_string())); - std::process::exit(1); - } -} -``` - -- [ ] **Step 4: Verify workspace recognizes the binary** - -Run: - -```bash -cargo metadata --format-version 1 --no-deps -``` - -Expected: output contains `"name":"dbx-cli"` for the package and `"name":"dbx-cli"` for the binary target. - ---- - -### Task 2: Core CLI Envelope and Error DTOs - -**Files:** -- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/cli.rs` -- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs` - -- [ ] **Step 1: Add envelope and error types** - -Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/cli.rs`: - -```rust -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "kebab-case")] -pub enum CliSource { - GuiRuntime, - Headless, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "SCREAMING_SNAKE_CASE")] -pub enum CliErrorCode { - GuiRuntimeRequired, - ConnectionNotFound, - AmbiguousConnection, - SecretUnavailable, - SshTunnelFailed, - QueryClassificationFailed, - HandoffRequired, - DdlBlocked, - ProductionWriteBlocked, - UnsupportedDatabaseType, - Timeout, - InternalError, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct CliError { - pub code: CliErrorCode, - pub message: String, - pub recoverable: bool, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(untagged)] -pub enum CliEnvelope { - Success { ok: bool, source: CliSource, data: T }, - Failure { ok: bool, source: CliSource, error: CliError }, -} - -pub fn ok(source: CliSource, data: T) -> CliEnvelope { - CliEnvelope::Success { ok: true, source, data } -} - -pub fn fail(source: CliSource, code: CliErrorCode, message: impl Into, recoverable: bool) -> CliEnvelope { - CliEnvelope::Failure { ok: false, source, error: CliError { code, message: message.into(), recoverable } } -} -``` - -- [ ] **Step 2: Export module** - -Add this line to `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs`: - -```rust -pub mod cli; -``` - -- [ ] **Step 3: Add envelope unit test** - -Append tests in `cli.rs`: - -```rust -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn serializes_success_source_as_kebab_case() { - let env = ok(CliSource::GuiRuntime, serde_json::json!({"value": 1})); - let json = serde_json::to_string(&env).unwrap(); - assert!(json.contains("\"ok\":true")); - assert!(json.contains("\"source\":\"gui-runtime\"")); - } - - #[test] - fn serializes_error_code_as_screaming_snake_case() { - let env: CliEnvelope<()> = fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, "runtime needed", true); - let json = serde_json::to_string(&env).unwrap(); - assert!(json.contains("\"GUI_RUNTIME_REQUIRED\"")); - } -} -``` - -- [ ] **Step 4: Run targeted test** - -Run: - -```bash -cargo test -p dbx-core cli::tests -``` - -Expected: both tests pass. - ---- - -### Task 3: SQL Safety and Risk Classification - -**Files:** -- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/sql_safety.rs` -- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs` - -- [ ] **Step 1: Implement classifier** - -Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/sql_safety.rs`: - -```rust -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "lowercase")] -pub enum OperationClass { - Read, - Write, - Ddl, - Unknown, -} - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "lowercase")] -pub enum RiskLevel { - Low, - Medium, - High, - Critical, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct RiskMetadata { - pub operation_class: OperationClass, - pub risk_level: RiskLevel, - pub is_production: bool, - pub production_reason: Option, - pub first_token: Option, -} - -pub fn classify_sql(sql: &str) -> OperationClass { - let token = first_executable_token(sql).map(|s| s.to_ascii_uppercase()); - match token.as_deref() { - Some("SELECT" | "SHOW" | "DESCRIBE" | "EXPLAIN" | "WITH") => OperationClass::Read, - Some("INSERT" | "UPDATE" | "DELETE" | "MERGE" | "REPLACE") => OperationClass::Write, - Some("CREATE" | "ALTER" | "DROP" | "TRUNCATE" | "RENAME") => OperationClass::Ddl, - _ => OperationClass::Unknown, - } -} - -pub fn risk_for(sql: &str, connection_name: &str, color: Option<&str>) -> RiskMetadata { - let operation_class = classify_sql(sql); - let (is_production, production_reason) = production_signal(connection_name, color); - let risk_level = match (operation_class, is_production) { - (OperationClass::Read, _) => RiskLevel::Low, - (OperationClass::Write, false) => RiskLevel::Medium, - (OperationClass::Write, true) => RiskLevel::High, - (OperationClass::Ddl, _) => RiskLevel::Critical, - (OperationClass::Unknown, _) => RiskLevel::High, - }; - RiskMetadata { - operation_class, - risk_level, - is_production, - production_reason, - first_token: first_executable_token(sql).map(str::to_string), - } -} - -fn production_signal(connection_name: &str, color: Option<&str>) -> (bool, Option) { - if matches!(color, Some("#ef4444")) { - return (true, Some("red connection color".to_string())); - } - let name = connection_name.to_ascii_lowercase(); - if ["prod", "production", "live"].iter().any(|needle| name.contains(needle)) { - return (true, Some("connection name fallback".to_string())); - } - (false, None) -} - -fn first_executable_token(sql: &str) -> Option<&str> { - let bytes = sql.as_bytes(); - let mut i = 0; - while i < bytes.len() { - while i < bytes.len() && bytes[i].is_ascii_whitespace() { - i += 1; - } - if i + 1 < bytes.len() && bytes[i] == b'-' && bytes[i + 1] == b'-' { - i += 2; - while i < bytes.len() && bytes[i] != b'\n' { - i += 1; - } - continue; - } - if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' { - i += 2; - while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { - i += 1; - } - i = (i + 2).min(bytes.len()); - continue; - } - break; - } - let start = i; - while i < bytes.len() && (bytes[i].is_ascii_alphabetic() || bytes[i] == b'_') { - i += 1; - } - (i > start).then_some(&sql[start..i]) -} -``` - -- [ ] **Step 2: Export module** - -Add this line to `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs`: - -```rust -pub mod sql_safety; -``` - -- [ ] **Step 3: Add classifier tests** - -Append tests to `sql_safety.rs`: - -```rust -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn comments_do_not_hide_read_token() { - assert_eq!(classify_sql("-- comment\nSELECT 1"), OperationClass::Read); - assert_eq!(classify_sql("/* DROP TABLE x */ SELECT 1"), OperationClass::Read); - } - - #[test] - fn classifies_write_and_ddl() { - assert_eq!(classify_sql("update users set name = 'a'"), OperationClass::Write); - assert_eq!(classify_sql("DROP TABLE users"), OperationClass::Ddl); - } - - #[test] - fn red_color_marks_production() { - let risk = risk_for("orders", "prod-main", Some("#ef4444")); - assert!(risk.is_production); - assert_eq!(risk.risk_level, RiskLevel::Low); - } -} -``` - -- [ ] **Step 4: Run targeted tests** - -Run: - -```bash -cargo test -p dbx-core sql_safety::tests -``` - -Expected: all classifier tests pass. - ---- - -### Task 4: Handoff Storage Model - -**Files:** -- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/handoff.rs` -- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs` -- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/storage.rs` - -- [ ] **Step 1: Add handoff types** - -Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/handoff.rs`: - -```rust -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; - -use crate::sql_safety::{OperationClass, RiskLevel}; - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "lowercase")] -pub enum HandoffStatus { - Queued, - Shown, - Approved, - Rejected, - Executed, - Failed, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct HandoffItem { - pub id: String, - pub created_at: DateTime, - pub created_by: String, - pub connection_name: String, - pub database: Option, - pub title: String, - pub description: Option, - pub sql: String, - pub operation_class: OperationClass, - pub risk_level: RiskLevel, - pub is_production: bool, - pub status: HandoffStatus, - pub result_summary: Option, - pub error: Option, -} - -impl HandoffItem { - pub fn queued( - connection_name: String, - database: Option, - title: String, - description: Option, - sql: String, - operation_class: OperationClass, - risk_level: RiskLevel, - is_production: bool, - ) -> Self { - Self { - id: Uuid::new_v4().to_string(), - created_at: Utc::now(), - created_by: "dbx-cli".to_string(), - connection_name, - database, - title, - description, - sql, - operation_class, - risk_level, - is_production, - status: HandoffStatus::Queued, - result_summary: None, - error: None, - } - } -} -``` - -- [ ] **Step 2: Export module** - -Add this line to `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs`: - -```rust -pub mod handoff; -``` - -- [ ] **Step 3: Add storage table** - -Add this SQL statement to `SCHEMA_STATEMENTS` in `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/storage.rs`: - -```rust -"CREATE TABLE IF NOT EXISTS handoffs ( - id TEXT PRIMARY KEY, - payload_json TEXT NOT NULL, - status TEXT NOT NULL, - created_at TEXT NOT NULL -)", -``` - -- [ ] **Step 4: Add storage methods** - -Append this impl block near the storage sections in `storage.rs`: - -```rust -impl Storage { - pub async fn save_handoff(&self, item: &crate::handoff::HandoffItem) -> Result<(), String> { - let json = serde_json::to_string(item).map_err(|e| e.to_string())?; - sqlx::query( - "INSERT OR REPLACE INTO handoffs (id, payload_json, status, created_at) VALUES (?, ?, ?, ?)", - ) - .bind(&item.id) - .bind(json) - .bind(format!("{:?}", item.status).to_lowercase()) - .bind(item.created_at.to_rfc3339()) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; - Ok(()) - } - - pub async fn load_pending_handoffs(&self) -> Result, String> { - let rows: Vec<(String,)> = sqlx::query_as( - "SELECT payload_json FROM handoffs WHERE status IN ('queued', 'shown') ORDER BY created_at ASC", - ) - .fetch_all(&self.db) - .await - .map_err(|e| e.to_string())?; - rows.into_iter() - .map(|(json,)| serde_json::from_str(&json).map_err(|e| e.to_string())) - .collect() - } -} -``` - -- [ ] **Step 5: Run core tests** - -Run: - -```bash -cargo test -p dbx-core -``` - -Expected: existing tests pass and storage compiles with new methods. - ---- - -### Task 5: Schema Snapshot Orchestration - -**Files:** -- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/schema_snapshot.rs` -- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs` - -- [ ] **Step 1: Add snapshot DTOs and function** - -Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/schema_snapshot.rs`: - -```rust -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; - -use crate::connection::AppState; -use crate::models::connection::DatabaseType; -use crate::{schema, types}; - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct TableSnapshot { - pub name: String, - pub table_type: String, - pub columns: Vec, - pub indexes: Vec, - pub foreign_keys: Vec, - pub triggers: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SchemaSnapshot { - pub connection_id: String, - pub connection_name: String, - pub database: Option, - pub database_type: DatabaseType, - pub driver_profile: Option, - pub captured_at: DateTime, - pub databases: Vec, - pub schemas: Vec, - pub tables: Vec, -} - -pub async fn snapshot( - state: &AppState, - connection_id: &str, - database: Option<&str>, - schema_name: Option<&str>, -) -> Result { - let config = { - let configs = state.configs.lock().await; - configs.get(connection_id).cloned().ok_or("Connection config not found")? - }; - let db = database.or(config.database.as_deref()).unwrap_or_default(); - let schemas = if db.is_empty() { - Vec::new() - } else { - schema::list_schemas_core(state, connection_id, db).await.unwrap_or_default() - }; - let effective_schema = schema_name.or_else(|| schemas.first().map(String::as_str)).unwrap_or(""); - let databases = schema::list_databases_core(state, connection_id).await.unwrap_or_default(); - let table_infos = if db.is_empty() { - Vec::new() - } else { - schema::list_tables_core(state, connection_id, db, effective_schema).await.unwrap_or_default() - }; - let mut tables = Vec::new(); - for table in table_infos { - let columns = schema::get_columns_core(state, connection_id, db, effective_schema, &table.name).await.unwrap_or_default(); - let indexes = schema::list_indexes_core(state, connection_id, db, effective_schema, &table.name).await.unwrap_or_default(); - let foreign_keys = schema::list_foreign_keys_core(state, connection_id, db, effective_schema, &table.name).await.unwrap_or_default(); - let triggers = schema::list_triggers_core(state, connection_id, db, effective_schema, &table.name).await.unwrap_or_default(); - tables.push(TableSnapshot { name: table.name, table_type: table.table_type, columns, indexes, foreign_keys, triggers }); - } - Ok(SchemaSnapshot { - connection_id: config.id, - connection_name: config.name, - database: (!db.is_empty()).then(|| db.to_string()), - database_type: config.db_type, - driver_profile: config.driver_profile, - captured_at: Utc::now(), - databases, - schemas, - tables, - }) -} -``` - -- [ ] **Step 2: Export module** - -Add this line to `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs`: - -```rust -pub mod schema_snapshot; -``` - -- [ ] **Step 3: Compile core** - -Run: - -```bash -cargo check -p dbx-core -``` - -Expected: `dbx-core` compiles. - ---- - -### Task 6: Runtime Discovery and CLI Command Dispatch - -**Files:** -- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/runtime_client.rs` -- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` - -- [ ] **Step 1: Add runtime client** - -Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/runtime_client.rs`: - -```rust -use serde::{Deserialize, Serialize}; -use std::path::PathBuf; - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct RuntimeDiscovery { - pub port: u16, - pub token: String, -} - -pub fn app_data_dir() -> PathBuf { - let home = std::env::var(if cfg!(windows) { "APPDATA" } else { "HOME" }).unwrap_or_else(|_| ".".to_string()); - if cfg!(target_os = "macos") { - PathBuf::from(home).join("Library/Application Support/com.dbx.app") - } else if cfg!(windows) { - PathBuf::from(home).join("com.dbx.app") - } else { - PathBuf::from(home).join(".config/com.dbx.app") - } -} - -pub fn load_runtime() -> Option { - let path = app_data_dir().join("agent-runtime.json"); - let json = std::fs::read_to_string(path).ok()?; - serde_json::from_str(&json).ok() -} - -pub async fn get_json(path: &str) -> Result { - let runtime = load_runtime().ok_or("runtime unavailable")?; - let url = format!("http://127.0.0.1:{}{}", runtime.port, path); - reqwest::Client::new() - .get(url) - .bearer_auth(runtime.token) - .send() - .await - .map_err(|e| e.to_string())? - .json() - .await - .map_err(|e| e.to_string()) -} - -pub async fn post_json(path: &str, body: serde_json::Value) -> Result { - let runtime = load_runtime().ok_or("runtime unavailable")?; - let url = format!("http://127.0.0.1:{}{}", runtime.port, path); - reqwest::Client::new() - .post(url) - .bearer_auth(runtime.token) - .json(&body) - .send() - .await - .map_err(|e| e.to_string())? - .json() - .await - .map_err(|e| e.to_string()) -} -``` - -- [ ] **Step 2: Add command dispatcher skeleton** - -Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` with handlers for all 8 commands: - -```rust -use dbx_core::cli::{fail, ok, CliEnvelope, CliErrorCode, CliSource}; - -pub async fn run(args: Vec) -> Result<(), CliEnvelope<()>> { - let output = match args.as_slice() { - [cmd, rest @ ..] if cmd == "context" => context(rest).await, - [cmd, sub, rest @ ..] if cmd == "conn" && sub == "list" => conn_list(rest).await, - [cmd, sub, name, rest @ ..] if cmd == "conn" && sub == "show" => conn_show(name, rest).await, - [cmd, sub, rest @ ..] if cmd == "schema" && sub == "snapshot" => schema_snapshot(rest).await, - [cmd, rest @ ..] if cmd == "safe-query" => safe_query(rest).await, - [cmd, rest @ ..] if cmd == "handoff" => handoff(rest).await, - [cmd, rest @ ..] if cmd == "selection" => selection(rest).await, - [cmd, sub, rest @ ..] if cmd == "result" && sub == "current" => result_current(rest).await, - _ => fail(CliSource::Headless, CliErrorCode::InternalError, "Unknown command", false), - }; - println!("{}", serde_json::to_string_pretty(&output).unwrap()); - if matches!(output, CliEnvelope::Failure { .. }) { - std::process::exit(1); - } - Ok(()) -} - -async fn context(_args: &[String]) -> CliEnvelope { - match crate::runtime_client::get_json("/context").await { - Ok(data) => ok(CliSource::GuiRuntime, data), - Err(_) => ok(CliSource::Headless, serde_json::json!({ "runtime": "headless" })), - } -} - -async fn conn_list(_args: &[String]) -> CliEnvelope { - ok(CliSource::Headless, serde_json::json!({ "connections": [] })) -} - -async fn conn_show(_name: &str, _args: &[String]) -> CliEnvelope { - ok(CliSource::Headless, serde_json::json!({})) -} - -async fn schema_snapshot(_args: &[String]) -> CliEnvelope { - fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "Connection is required", true) -} - -async fn safe_query(_args: &[String]) -> CliEnvelope { - fail(CliSource::Headless, CliErrorCode::QueryClassificationFailed, "SQL is required", true) -} - -async fn handoff(_args: &[String]) -> CliEnvelope { - fail(CliSource::Headless, CliErrorCode::InternalError, "Handoff payload is required", true) -} - -async fn selection(_args: &[String]) -> CliEnvelope { - match crate::runtime_client::get_json("/selection").await { - Ok(data) => ok(CliSource::GuiRuntime, data), - Err(_) => fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, "dbx selection requires DBX GUI runtime.", true), - } -} - -async fn result_current(args: &[String]) -> CliEnvelope { - let limit = option_value(args, "--limit").unwrap_or("50"); - match crate::runtime_client::get_json(&format!("/result/current?limit={limit}")).await { - Ok(data) => ok(CliSource::GuiRuntime, data), - Err(_) => fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, "dbx result current requires DBX GUI runtime.", true), - } -} - -fn option_value<'a>(args: &'a [String], key: &str) -> Option<&'a str> { - args.windows(2).find(|pair| pair[0] == key).map(|pair| pair[1].as_str()) -} -``` - -- [ ] **Step 3: Verify CLI compiles** - -Run: - -```bash -cargo check -p dbx-cli -``` - -Expected: CLI crate compiles. - ---- - -### Task 7: Headless Connection Commands - -**Files:** -- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` - -- [ ] **Step 1: Add headless state helper and redaction** - -Add helper functions to `commands.rs`: - -```rust -async fn open_state() -> Result { - let db_path = crate::runtime_client::app_data_dir().join("dbx.db"); - let storage = dbx_core::storage::Storage::open(&db_path).await?; - Ok(dbx_core::connection::AppState::new(storage)) -} - -fn redacted_config(config: &dbx_core::models::connection::ConnectionConfig) -> serde_json::Value { - serde_json::json!({ - "id": config.id, - "name": config.name, - "databaseType": config.db_type, - "driverProfile": config.driver_profile, - "driverLabel": config.driver_label, - "defaultDatabase": config.database, - "color": config.color, - "sshEnabled": config.ssh_enabled, - "redactedUrl": config.redacted_connection_url(), - }) -} - -async fn find_connection(name: &str) -> Result> { - let state = open_state().await.map_err(|e| fail(CliSource::Headless, CliErrorCode::InternalError, e, false))?; - let configs = state.storage.load_connections().await.map_err(|e| fail(CliSource::Headless, CliErrorCode::InternalError, e, false))?; - let matches: Vec<_> = configs.into_iter().filter(|c| c.name == name || c.id == name).collect(); - match matches.len() { - 1 => Ok(matches.into_iter().next().unwrap()), - 0 => Err(fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "Connection not found", true)), - _ => Err(fail(CliSource::Headless, CliErrorCode::AmbiguousConnection, "Connection name is ambiguous", true)), - } -} -``` - -- [ ] **Step 2: Replace `conn_list` implementation** - -Use this implementation: - -```rust -async fn conn_list(_args: &[String]) -> CliEnvelope { - let state = match open_state().await { - Ok(state) => state, - Err(e) => return fail(CliSource::Headless, CliErrorCode::InternalError, e, false), - }; - match state.storage.load_connections().await { - Ok(configs) => ok(CliSource::Headless, serde_json::json!({ - "connections": configs.iter().map(redacted_config).collect::>() - })), - Err(e) => fail(CliSource::Headless, CliErrorCode::InternalError, e, false), - } -} -``` - -- [ ] **Step 3: Replace `conn_show` implementation** - -Use this implementation: - -```rust -async fn conn_show(name: &str, _args: &[String]) -> CliEnvelope { - match find_connection(name).await { - Ok(config) => ok(CliSource::Headless, redacted_config(&config)), - Err(err) => err, - } -} -``` - -- [ ] **Step 4: Verify commands** - -Run: - -```bash -cargo run -p dbx-cli --bin dbx-cli -- conn list --format json -cargo run -p dbx-cli --bin dbx-cli -- conn show __missing__ --redacted --format json -``` - -Expected: first command returns an envelope; second returns `CONNECTION_NOT_FOUND` without panicking. - ---- - -### Task 8: Headless Schema Snapshot and Safe Query - -**Files:** -- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` - -- [ ] **Step 1: Add connection registration helper** - -Add this helper: - -```rust -async fn state_with_connection(config: dbx_core::models::connection::ConnectionConfig) -> Result { - let state = open_state().await?; - state.configs.lock().await.insert(config.id.clone(), config.clone()); - state.get_or_create_pool(&config.id, config.database.as_deref()).await?; - Ok(state) -} -``` - -- [ ] **Step 2: Replace `schema_snapshot`** - -Use this implementation: - -```rust -async fn schema_snapshot(args: &[String]) -> CliEnvelope { - let Some(conn_name) = option_value(args, "--conn") else { - return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); - }; - let db = option_value(args, "--db"); - let config = match find_connection(conn_name).await { - Ok(config) => config, - Err(err) => return err, - }; - let state = match state_with_connection(config.clone()).await { - Ok(state) => state, - Err(e) => return fail(CliSource::Headless, CliErrorCode::InternalError, e, false), - }; - match dbx_core::schema_snapshot::snapshot(&state, &config.id, db, None).await { - Ok(snapshot) => ok(CliSource::Headless, serde_json::to_value(snapshot).unwrap()), - Err(e) => fail(CliSource::Headless, CliErrorCode::InternalError, e, false), - } -} -``` - -- [ ] **Step 3: Replace `safe_query`** - -Use this implementation: - -```rust -async fn safe_query(args: &[String]) -> CliEnvelope { - let Some(conn_name) = option_value(args, "--conn") else { - return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); - }; - let Some(sql) = option_value(args, "--sql") else { - return fail(CliSource::Headless, CliErrorCode::QueryClassificationFailed, "--sql is required", true); - }; - let config = match find_connection(conn_name).await { - Ok(config) => config, - Err(err) => return err, - }; - let risk = dbx_core::sql_safety::risk_for(sql, &config.name, config.color.as_deref()); - match risk.operation_class { - dbx_core::sql_safety::OperationClass::Read => {} - dbx_core::sql_safety::OperationClass::Write if risk.is_production => { - return fail(CliSource::Headless, CliErrorCode::ProductionWriteBlocked, serde_json::to_string(&risk).unwrap(), true); - } - dbx_core::sql_safety::OperationClass::Write => { - return fail(CliSource::Headless, CliErrorCode::HandoffRequired, serde_json::to_string(&risk).unwrap(), true); - } - dbx_core::sql_safety::OperationClass::Ddl => { - return fail(CliSource::Headless, CliErrorCode::DdlBlocked, serde_json::to_string(&risk).unwrap(), true); - } - dbx_core::sql_safety::OperationClass::Unknown => { - return fail(CliSource::Headless, CliErrorCode::QueryClassificationFailed, serde_json::to_string(&risk).unwrap(), true); - } - } - let state = match state_with_connection(config.clone()).await { - Ok(state) => state, - Err(e) => return fail(CliSource::Headless, CliErrorCode::InternalError, e, false), - }; - let database = option_value(args, "--db").or(config.database.as_deref()).unwrap_or(""); - match dbx_core::query::execute_sql_statement(&state, &config.id, database, sql, None, None).await { - Ok(result) => ok(CliSource::Headless, serde_json::json!({ "risk": risk, "result": result })), - Err(e) => fail(CliSource::Headless, CliErrorCode::InternalError, e, false), - } -} -``` - -- [ ] **Step 4: Verify SQL blocking** - -Run: - -```bash -cargo run -p dbx-cli --bin dbx-cli -- safe-query --conn __missing__ --sql "DROP TABLE users" --format json -``` - -Expected: returns `CONNECTION_NOT_FOUND`. With a real DBX connection, `DROP TABLE users` returns `DDL_BLOCKED`. - ---- - -### Task 9: GUI Runtime Server - -**Files:** -- Create: `/Users/bytedance/open_source_poj/dbx/src-tauri/src/commands/agent_runtime.rs` -- Modify: `/Users/bytedance/open_source_poj/dbx/src-tauri/src/commands/mod.rs` -- Modify: `/Users/bytedance/open_source_poj/dbx/src-tauri/src/lib.rs` - -- [ ] **Step 1: Add runtime command module** - -Create `/Users/bytedance/open_source_poj/dbx/src-tauri/src/commands/agent_runtime.rs` with: - -```rust -use serde::{Deserialize, Serialize}; -use std::sync::Arc; -use tauri::{AppHandle, Manager}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::TcpListener; -use tokio::sync::RwLock; - -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AgentRuntimeSnapshot { - pub active_connection_id: Option, - pub active_connection_name: Option, - pub database: Option, - pub schema: Option, - pub active_tab_id: Option, - pub active_tab_title: Option, - pub sql: Option, - pub selected_sql: Option, - pub selection: Option, - pub result: Option, -} - -#[derive(Clone)] -pub struct AgentRuntimeState { - pub token: String, - pub snapshot: Arc>, - pub handoffs: Arc>>, -} - -#[tauri::command] -pub async fn agent_runtime_update_snapshot( - state: tauri::State<'_, AgentRuntimeState>, - snapshot: AgentRuntimeSnapshot, -) -> Result<(), String> { - *state.snapshot.write().await = snapshot; - Ok(()) -} - -#[tauri::command] -pub async fn agent_runtime_load_handoffs( - app_state: tauri::State<'_, Arc>, - runtime: tauri::State<'_, AgentRuntimeState>, -) -> Result, String> { - let mut items = app_state.storage.load_pending_handoffs().await?; - items.extend(runtime.handoffs.read().await.iter().cloned()); - Ok(items) -} - -pub fn start(app: AppHandle) -> AgentRuntimeState { - let token = uuid::Uuid::new_v4().to_string(); - let state = AgentRuntimeState { - token: token.clone(), - snapshot: Arc::new(RwLock::new(AgentRuntimeSnapshot::default())), - handoffs: Arc::new(RwLock::new(Vec::new())), - }; - let server_state = state.clone(); - tauri::async_runtime::spawn(async move { - let listener = match TcpListener::bind("127.0.0.1:0").await { - Ok(listener) => listener, - Err(err) => { - log::warn!("Agent runtime bind failed: {err}"); - return; - } - }; - let port = listener.local_addr().map(|addr| addr.port()).unwrap_or(0); - if let Ok(dir) = app.path().app_data_dir() { - let payload = serde_json::json!({ "port": port, "token": token }); - let _ = std::fs::write(dir.join("agent-runtime.json"), serde_json::to_string(&payload).unwrap()); - } - loop { - let Ok((stream, _)) = listener.accept().await else { continue }; - let st = server_state.clone(); - tauri::async_runtime::spawn(async move { - handle_connection(stream, st).await; - }); - } - }); - state -} - -async fn handle_connection(mut stream: tokio::net::TcpStream, state: AgentRuntimeState) { - let mut buf = vec![0u8; 65536]; - let Ok(n) = stream.read(&mut buf).await else { return }; - if n == 0 { - return; - } - let request = String::from_utf8_lossy(&buf[..n]); - if !request.contains(&format!("Authorization: Bearer {}", state.token)) { - respond_json(&mut stream, "401 Unauthorized", serde_json::json!({"error":"unauthorized"})).await; - return; - } - let first = request.lines().next().unwrap_or(""); - if first.starts_with("GET /context") { - respond_json(&mut stream, "200 OK", serde_json::to_value(&*state.snapshot.read().await).unwrap()).await; - } else if first.starts_with("GET /selection") { - let snapshot = state.snapshot.read().await; - respond_json(&mut stream, "200 OK", snapshot.selection.clone().unwrap_or_else(|| serde_json::json!({"type":"none"}))).await; - } else if first.starts_with("GET /result/current") { - let snapshot = state.snapshot.read().await; - respond_json(&mut stream, "200 OK", snapshot.result.clone().unwrap_or_else(|| serde_json::json!({"columns":[],"rows":[]}))).await; - } else if first.starts_with("POST /handoff") { - let body = request.split("\r\n\r\n").nth(1).unwrap_or(""); - if let Ok(item) = serde_json::from_str::(body) { - state.handoffs.write().await.push(item.clone()); - respond_json(&mut stream, "200 OK", serde_json::json!({"id": item.id, "status": "shown"})).await; - } else { - respond_json(&mut stream, "400 Bad Request", serde_json::json!({"error":"invalid handoff"})).await; - } - } else { - respond_json(&mut stream, "404 Not Found", serde_json::json!({"error":"not found"})).await; - } -} - -async fn respond_json(stream: &mut tokio::net::TcpStream, status: &str, body: serde_json::Value) { - let body = serde_json::to_string(&body).unwrap(); - let resp = format!("HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}", body.len()); - let _ = stream.write_all(resp.as_bytes()).await; -} -``` - -- [ ] **Step 2: Expose module** - -Add to `/Users/bytedance/open_source_poj/dbx/src-tauri/src/commands/mod.rs`: - -```rust -pub mod agent_runtime; -``` - -- [ ] **Step 3: Start runtime in Tauri setup** - -In `/Users/bytedance/open_source_poj/dbx/src-tauri/src/lib.rs`, after `commands::mcp_bridge::start(app_handle, state);`, add: - -```rust -let runtime_state = commands::agent_runtime::start(app.handle().clone()); -app.manage(runtime_state); -``` - -Add command to `invoke_handler`: - -```rust -commands::agent_runtime::agent_runtime_update_snapshot, -commands::agent_runtime::agent_runtime_load_handoffs, -``` - -- [ ] **Step 4: Compile Tauri crate** - -Run: - -```bash -cargo check -p dbx -``` - -Expected: desktop crate compiles. - ---- - -### Task 10: Frontend Runtime State Sync - -**Files:** -- Modify: `/Users/bytedance/open_source_poj/dbx/src/lib/tauri.ts` -- Modify: `/Users/bytedance/open_source_poj/dbx/src/lib/api.ts` -- Create: `/Users/bytedance/open_source_poj/dbx/src/stores/agentRuntimeStore.ts` -- Modify: `/Users/bytedance/open_source_poj/dbx/src/stores/queryStore.ts` -- Modify: `/Users/bytedance/open_source_poj/dbx/src/components/grid/DataGrid.vue` - -- [ ] **Step 1: Add Tauri wrappers** - -Add to `/Users/bytedance/open_source_poj/dbx/src/lib/tauri.ts`: - -```ts -export async function agentRuntimeUpdateSnapshot(snapshot: unknown): Promise { - return invoke("agent_runtime_update_snapshot", { snapshot }); -} - -export async function agentRuntimeLoadHandoffs(): Promise { - return invoke("agent_runtime_load_handoffs"); -} -``` - -- [ ] **Step 2: Export API forwards** - -Add to `/Users/bytedance/open_source_poj/dbx/src/lib/api.ts`: - -```ts -export const agentRuntimeUpdateSnapshot = forward("agentRuntimeUpdateSnapshot"); -export const agentRuntimeLoadHandoffs = forward("agentRuntimeLoadHandoffs"); -``` - -- [ ] **Step 3: Create runtime store** - -Create `/Users/bytedance/open_source_poj/dbx/src/stores/agentRuntimeStore.ts`: - -```ts -import { defineStore } from "pinia"; -import { ref } from "vue"; -import * as api from "@/lib/api"; -import { useConnectionStore } from "@/stores/connectionStore"; -import { useQueryStore } from "@/stores/queryStore"; - -export const useAgentRuntimeStore = defineStore("agentRuntime", () => { - const selection = ref({ type: "none" }); - let timer: ReturnType | null = null; - - function setSelection(value: unknown) { - selection.value = value; - scheduleSync(); - } - - function scheduleSync() { - if (timer) clearTimeout(timer); - timer = setTimeout(() => void syncNow(), 100); - } - - async function syncNow() { - const connectionStore = useConnectionStore(); - const queryStore = useQueryStore(); - const tab = queryStore.tabs.find((item) => item.id === queryStore.activeTabId); - const conn = tab ? connectionStore.getConfig(tab.connectionId) : undefined; - await api.agentRuntimeUpdateSnapshot({ - activeConnectionId: tab?.connectionId, - activeConnectionName: conn?.name, - database: tab?.database, - schema: tab?.schema, - activeTabId: tab?.id, - activeTabTitle: tab?.title, - sql: tab?.sql, - selectedSql: undefined, - selection: selection.value, - result: tab?.result - ? { - columns: tab.result.columns, - rows: tab.result.rows.slice(0, 50), - truncated: tab.result.rows.length > 50 || tab.result.truncated, - executionTimeMs: tab.result.execution_time_ms, - } - : undefined, - }); - } - - return { selection, setSelection, scheduleSync, syncNow }; -}); -``` - -- [ ] **Step 4: Notify sync from query store** - -In `/Users/bytedance/open_source_poj/dbx/src/stores/queryStore.ts`, import: - -```ts -import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore"; -``` - -After result assignment in `executeTabSql`, call: - -```ts -useAgentRuntimeStore().scheduleSync(); -``` - -Also call `useAgentRuntimeStore().scheduleSync();` after active tab changes in `createTab`, `openSavedSql`, `closeTab`, and `setActiveResultIndex`. - -- [ ] **Step 5: Notify grid selection** - -In `/Users/bytedance/open_source_poj/dbx/src/components/grid/DataGrid.vue`, import: - -```ts -import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore"; -``` - -Create a store instance: - -```ts -const agentRuntimeStore = useAgentRuntimeStore(); -``` - -Where selection range changes, call: - -```ts -agentRuntimeStore.setSelection({ - type: "grid-cells", - data: extractSelection(props.result.columns, props.result.rows, selectionRange.value), -}); -``` - -- [ ] **Step 6: Run frontend typecheck** - -Run: - -```bash -pnpm build -``` - -Expected: Vue typecheck and Vite build pass. - ---- - -### Task 11: Handoff CLI and GUI Review UI - -**Files:** -- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` -- Create: `/Users/bytedance/open_source_poj/dbx/src/components/agent/AgentHandoffDialog.vue` -- Modify: `/Users/bytedance/open_source_poj/dbx/src/components/layout/AppDialogs.vue` - -- [ ] **Step 1: Implement CLI handoff** - -Replace `handoff` in `commands.rs` with: - -```rust -async fn handoff(args: &[String]) -> CliEnvelope { - let Some(conn_name) = option_value(args, "--conn") else { - return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); - }; - let Some(title) = option_value(args, "--title") else { - return fail(CliSource::Headless, CliErrorCode::InternalError, "--title is required", true); - }; - let sql = if let Some(sql_file) = option_value(args, "--sql-file") { - match std::fs::read_to_string(sql_file) { - Ok(sql) => sql, - Err(e) => return fail(CliSource::Headless, CliErrorCode::InternalError, e.to_string(), true), - } - } else if let Some(sql) = option_value(args, "--sql") { - sql.to_string() - } else { - return fail(CliSource::Headless, CliErrorCode::InternalError, "--sql-file or --sql is required", true); - }; - let config = match find_connection(conn_name).await { - Ok(config) => config, - Err(err) => return err, - }; - let risk = dbx_core::sql_safety::risk_for(&sql, &config.name, config.color.as_deref()); - let item = dbx_core::handoff::HandoffItem::queued( - config.name, - config.database, - title.to_string(), - option_value(args, "--description").map(str::to_string), - sql, - risk.operation_class, - risk.risk_level, - risk.is_production, - ); - if let Ok(data) = crate::runtime_client::post_json("/handoff", serde_json::to_value(&item).unwrap()).await { - return ok(CliSource::GuiRuntime, data); - } - match open_state().await.and_then(|state| futures::executor::block_on(state.storage.save_handoff(&item))) { - Ok(()) => ok(CliSource::Headless, serde_json::json!({ "id": item.id, "status": "queued" })), - Err(e) => fail(CliSource::Headless, CliErrorCode::InternalError, e, false), - } -} -``` - -- [ ] **Step 2: Create handoff review dialog** - -Create `/Users/bytedance/open_source_poj/dbx/src/components/agent/AgentHandoffDialog.vue`: - -```vue - - - -``` - -- [ ] **Step 3: Mount dialog** - -In `/Users/bytedance/open_source_poj/dbx/src/components/layout/AppDialogs.vue`, import and add: - -```vue - -``` - -- [ ] **Step 4: Verify GUI build** - -Run: - -```bash -pnpm build -``` - -Expected: build passes and dialog component compiles. - ---- - -### Task 12: Runtime Context, Selection, and Result CLI Verification - -**Files:** -- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` -- Modify: `/Users/bytedance/open_source_poj/dbx/src/stores/agentRuntimeStore.ts` - -- [ ] **Step 1: Improve context headless fallback** - -Replace headless branch of `context` with: - -```rust -Err(_) => { - let state = match open_state().await { - Ok(state) => state, - Err(e) => return fail(CliSource::Headless, CliErrorCode::InternalError, e, false), - }; - let configs = state.storage.load_connections().await.unwrap_or_default(); - ok(CliSource::Headless, serde_json::json!({ - "runtime": "headless", - "activeConnection": configs.first().map(redacted_config), - "configSource": crate::runtime_client::app_data_dir().join("dbx.db").display().to_string() - })) -} -``` - -- [ ] **Step 2: Ensure result limit is applied** - -In `/Users/bytedance/open_source_poj/dbx/src/stores/agentRuntimeStore.ts`, keep `rows.slice(0, 50)` and ensure `truncated` is set when original rows exceed 50. - -- [ ] **Step 3: Manual runtime verification** - -Run DBX desktop: - -```bash -pnpm tauri dev -``` - -In a second terminal, run: - -```bash -cargo run -p dbx-cli --bin dbx-cli -- context --format json -cargo run -p dbx-cli --bin dbx-cli -- selection --format json -cargo run -p dbx-cli --bin dbx-cli -- result current --limit 50 --format json -``` - -Expected: all three commands return `source: "gui-runtime"` when desktop is running. - ---- - -### Task 13: Final Verification and Compatibility Check - -**Files:** -- No new files. - -- [ ] **Step 1: Verify Rust workspace** - -Run: - -```bash -cargo check --workspace -cargo test -p dbx-core -``` - -Expected: workspace checks and core tests pass. - -- [ ] **Step 2: Verify frontend** - -Run: - -```bash -pnpm build -``` - -Expected: Vue typecheck and Vite build pass. - -- [ ] **Step 3: Verify MCP untouched** - -Run: - -```bash -git diff -- mcp -``` - -Expected: no output. - -- [ ] **Step 4: Verify CLI command surface** - -Run: - -```bash -cargo run -p dbx-cli --bin dbx-cli -- context --format json -cargo run -p dbx-cli --bin dbx-cli -- conn list --format json -cargo run -p dbx-cli --bin dbx-cli -- conn show __missing__ --redacted --format json -cargo run -p dbx-cli --bin dbx-cli -- selection --format json -cargo run -p dbx-cli --bin dbx-cli -- result current --limit 50 --format json -``` - -Expected: all commands return valid JSON envelopes. GUI-only commands return `GUI_RUNTIME_REQUIRED` if DBX desktop is not running. - -## Self-Review - -- Spec coverage: The plan covers all 8 CLI commands, runtime discovery/token, GUI context, selection, current result, safe-query classification, schema snapshot, and handoff queue/display. -- Compatibility: The plan explicitly avoids changing `mcp/` and verifies this with `git diff -- mcp`. -- Type consistency: Rust DTOs use camelCase where sent to frontend/CLI. Error codes use screaming snake case. Source uses kebab case. -- Known implementation detail to watch: Task 11 uses a synchronous block inside async CLI code for queued handoff save; if it causes compile friction, replace it with direct async `state.storage.save_handoff(&item).await` inside a `match` after opening state. diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index d4ec4bf0e..08b00aa33 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -48,5 +48,4 @@ russh = "0.60" csv = "1.4.0" calamine = "0.30.1" zip = "4.6.1" -libc = "0.2" dbx-core = { path = "../crates/dbx-core" } diff --git a/src-tauri/src/commands/agent_runtime.rs b/src-tauri/src/commands/agent_runtime.rs deleted file mode 100644 index f878862dd..000000000 --- a/src-tauri/src/commands/agent_runtime.rs +++ /dev/null @@ -1,889 +0,0 @@ -use serde::{Deserialize, Serialize}; -use std::fs::OpenOptions; -use std::io::Write; -use std::path::{Path, PathBuf}; -use std::sync::Arc; -use tauri::{AppHandle, Manager}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::{oneshot, RwLock}; - -use super::connection::AppState; -use dbx_core::handoff::{HandoffItem, HandoffStatus}; -use dbx_core::models::connection::ConnectionConfig; -use dbx_core::sql_safety::{classify_sql, risk_for, risk_for_connection, OperationClass, RiskContext, RiskLevel}; - -const BIND_ADDR: &str = "127.0.0.1:0"; -const DISCOVERY_FILE: &str = "agent-runtime.json"; -const MAX_HEADER_BYTES: usize = 16 * 1024; -const MAX_BODY_BYTES: usize = 1024 * 1024; - -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AgentRuntimeSnapshot { - pub active_connection_id: Option, - pub active_connection_name: Option, - pub database: Option, - pub schema: Option, - pub active_tab_id: Option, - pub active_tab_title: Option, - pub sql: Option, - pub selected_sql: Option, - pub selection: Option, - pub result: Option, -} - -#[derive(Clone)] -pub struct AgentRuntimeState { - pub token: String, - pub snapshot: Arc>, - pub handoffs: Arc>>, - pub app_state: Option>, -} - -pub struct AgentRuntimeServer { - state: AgentRuntimeState, - discovery_path: PathBuf, - shutdown: std::sync::Mutex>>, -} - -impl AgentRuntimeServer { - pub fn state(&self) -> &AgentRuntimeState { - &self.state - } - - pub fn cleanup(&self) { - if let Ok(mut shutdown) = self.shutdown.lock() { - if let Some(tx) = shutdown.take() { - let _ = tx.send(()); - } - } - cleanup_discovery_file(&self.discovery_path); - } -} - -impl Drop for AgentRuntimeServer { - fn drop(&mut self) { - if let Ok(mut shutdown) = self.shutdown.lock() { - if let Some(tx) = shutdown.take() { - let _ = tx.send(()); - } - } - cleanup_discovery_file(&self.discovery_path); - } -} - -#[derive(Debug, PartialEq, Eq)] -struct RuntimeResponse { - status: &'static str, - body: serde_json::Value, -} - -#[derive(Debug)] -struct RuntimeRequest { - first_line: String, - headers: Vec<(String, String)>, - body: String, -} - -#[tauri::command] -pub async fn agent_runtime_update_snapshot( - runtime: tauri::State<'_, AgentRuntimeServer>, - snapshot: AgentRuntimeSnapshot, -) -> Result<(), String> { - *runtime.state().snapshot.write().await = snapshot; - Ok(()) -} - -#[tauri::command] -pub async fn agent_runtime_load_handoffs( - app_state: tauri::State<'_, Arc>, - runtime: tauri::State<'_, AgentRuntimeServer>, -) -> Result, String> { - let mut items = app_state.storage.load_pending_handoffs().await?; - items.extend(pending_runtime_handoffs(runtime.state()).await); - Ok(items) -} - -#[tauri::command] -pub async fn agent_runtime_mark_handoff_shown( - app_state: tauri::State<'_, Arc>, - runtime: tauri::State<'_, AgentRuntimeServer>, - id: String, -) -> Result { - update_handoff_status(app_state.inner().as_ref(), runtime.state(), &id, HandoffStatus::Shown).await -} - -#[tauri::command] -pub async fn agent_runtime_reject_handoff( - app_state: tauri::State<'_, Arc>, - runtime: tauri::State<'_, AgentRuntimeServer>, - id: String, -) -> Result { - update_handoff_status(app_state.inner().as_ref(), runtime.state(), &id, HandoffStatus::Rejected).await -} - -pub fn start(app: AppHandle, app_state: Arc) -> AgentRuntimeServer { - let token = uuid::Uuid::new_v4().to_string(); - let state = AgentRuntimeState { - token: token.clone(), - snapshot: Arc::new(RwLock::new(AgentRuntimeSnapshot::default())), - handoffs: Arc::new(RwLock::new(Vec::new())), - app_state: Some(app_state), - }; - let discovery_path = - app.path().app_data_dir().map(|dir| dir.join(DISCOVERY_FILE)).unwrap_or_else(|_| PathBuf::from(DISCOVERY_FILE)); - let (shutdown_tx, shutdown_rx) = oneshot::channel(); - let server_state = state.clone(); - - tauri::async_runtime::spawn(async move { - run_server(app, server_state, shutdown_rx).await; - }); - - AgentRuntimeServer { state, discovery_path, shutdown: std::sync::Mutex::new(Some(shutdown_tx)) } -} - -async fn run_server(app: AppHandle, state: AgentRuntimeState, mut shutdown: oneshot::Receiver<()>) { - let listener = match TcpListener::bind(BIND_ADDR).await { - Ok(listener) => listener, - Err(err) => { - log::warn!("Agent runtime failed to bind {BIND_ADDR}: {err}"); - return; - } - }; - let port = listener.local_addr().map(|addr| addr.port()).unwrap_or(0); - let discovery_path = match app.path().app_data_dir() { - Ok(dir) => match write_discovery_file(&dir, port, &state.token) { - Ok(path) => Some(path), - Err(err) => { - log::warn!("Agent runtime discovery write failed: {err}"); - None - } - }, - Err(err) => { - log::warn!("Agent runtime app data dir unavailable: {err}"); - None - } - }; - log::info!("Agent runtime listening on 127.0.0.1:{port}"); - - loop { - tokio::select! { - _ = &mut shutdown => { - if let Some(path) = discovery_path.as_deref() { - cleanup_discovery_file(path); - } - break; - } - accepted = listener.accept() => { - let Ok((stream, _)) = accepted else { continue }; - let st = state.clone(); - tokio::spawn(async move { - handle_connection(stream, st).await; - }); - } - } - } -} - -async fn handle_connection(mut stream: TcpStream, state: AgentRuntimeState) { - let request = match read_request(&mut stream).await { - Ok(Some(request)) => request, - Ok(None) => return, - Err(response) => { - respond_json(&mut stream, response.status, response.body).await; - return; - } - }; - - if !is_authorized_headers(&request.headers, &state.token) { - respond_json(&mut stream, "401 Unauthorized", serde_json::json!({"error": "unauthorized"})).await; - return; - } - - let response = route_request(&request.first_line, &request.body, &state).await; - respond_json(&mut stream, response.status, response.body).await; -} - -async fn read_request(stream: &mut TcpStream) -> Result, RuntimeResponse> { - let mut buf = Vec::new(); - let header_end = loop { - if let Some(pos) = find_header_end(&buf) { - break pos; - } - if buf.len() >= MAX_HEADER_BYTES { - return Err(RuntimeResponse { - status: "431 Request Header Fields Too Large", - body: serde_json::json!({"error": "headers too large"}), - }); - } - - let mut chunk = [0u8; 8192]; - let n = stream.read(&mut chunk).await.map_err(|_| RuntimeResponse { - status: "400 Bad Request", - body: serde_json::json!({"error": "invalid request"}), - })?; - if n == 0 { - return if buf.is_empty() { - Ok(None) - } else { - Err(RuntimeResponse { - status: "400 Bad Request", - body: serde_json::json!({"error": "incomplete request"}), - }) - }; - } - buf.extend_from_slice(&chunk[..n]); - }; - - let header_text = String::from_utf8_lossy(&buf[..header_end]); - let mut lines = header_text.lines(); - let first_line = lines.next().unwrap_or("").to_string(); - let headers: Vec<(String, String)> = lines - .filter_map(|line| { - let (name, value) = line.split_once(':')?; - Some((name.trim().to_string(), value.trim().to_string())) - }) - .collect(); - let content_length = content_length(&headers)?; - if content_length > MAX_BODY_BYTES { - return Err(RuntimeResponse { - status: "413 Payload Too Large", - body: serde_json::json!({"error": "body too large"}), - }); - } - - let body_start = header_end + 4; - let body_end = body_start + content_length; - while buf.len() < body_end { - let mut chunk = [0u8; 8192]; - let n = stream.read(&mut chunk).await.map_err(|_| RuntimeResponse { - status: "400 Bad Request", - body: serde_json::json!({"error": "invalid request"}), - })?; - if n == 0 { - return Err(RuntimeResponse { - status: "400 Bad Request", - body: serde_json::json!({"error": "incomplete body"}), - }); - } - buf.extend_from_slice(&chunk[..n]); - } - - Ok(Some(RuntimeRequest { - first_line, - headers, - body: String::from_utf8_lossy(&buf[body_start..body_end]).to_string(), - })) -} - -fn find_header_end(buf: &[u8]) -> Option { - buf.windows(4).position(|window| window == b"\r\n\r\n") -} - -fn content_length(headers: &[(String, String)]) -> Result { - match headers.iter().find(|(name, _)| name.eq_ignore_ascii_case("content-length")) { - Some((_, value)) => value.parse::().map_err(|_| RuntimeResponse { - status: "400 Bad Request", - body: serde_json::json!({"error": "invalid content-length"}), - }), - None => Ok(0), - } -} - -#[cfg(test)] -fn is_authorized(request: &str, token: &str) -> bool { - request.lines().any(|line| { - let Some((name, value)) = line.split_once(':') else { - return false; - }; - name.trim().eq_ignore_ascii_case("authorization") && value.trim() == format!("Bearer {token}") - }) -} - -fn is_authorized_headers(headers: &[(String, String)], token: &str) -> bool { - headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") && value == &format!("Bearer {token}")) -} - -async fn route_request(first_line: &str, body: &str, state: &AgentRuntimeState) -> RuntimeResponse { - if first_line.starts_with("GET /context ") || first_line.starts_with("GET /context?") { - return RuntimeResponse { - status: "200 OK", - body: serde_json::to_value(&*state.snapshot.read().await).unwrap_or_else(|_| serde_json::json!({})), - }; - } - - if first_line.starts_with("GET /selection ") || first_line.starts_with("GET /selection?") { - let snapshot = state.snapshot.read().await; - return RuntimeResponse { - status: "200 OK", - body: snapshot.selection.clone().unwrap_or_else(|| serde_json::json!({"type": "none"})), - }; - } - - if first_line.starts_with("GET /result/current ") || first_line.starts_with("GET /result/current?") { - let snapshot = state.snapshot.read().await; - let mut body = snapshot.result.clone().unwrap_or_else(|| serde_json::json!({"columns": [], "rows": []})); - if let Some(limit) = query_limit(first_line) { - truncate_result_rows(&mut body, limit); - } - return RuntimeResponse { status: "200 OK", body }; - } - - if first_line.starts_with("POST /handoff ") { - let mut item = match serde_json::from_str::(body) { - Ok(item) => item, - Err(_) => { - return RuntimeResponse { - status: "400 Bad Request", - body: serde_json::json!({"error": "invalid handoff"}), - }; - } - }; - let snapshot = state.snapshot.read().await.clone(); - recompute_handoff_risk(&mut item, &snapshot, state.app_state.as_deref()).await; - item.status = dbx_core::handoff::HandoffStatus::Shown; - let id = item.id.clone(); - state.handoffs.write().await.push(item); - return RuntimeResponse { status: "200 OK", body: serde_json::json!({"id": id, "status": "shown"}) }; - } - - RuntimeResponse { status: "404 Not Found", body: serde_json::json!({"error": "not found"}) } -} - -async fn recompute_handoff_risk(item: &mut HandoffItem, snapshot: &AgentRuntimeSnapshot, app_state: Option<&AppState>) { - let connection = load_handoff_connection(app_state, &item.connection_id).await; - let risk = match connection { - ConnectionLookup::Found(config) => { - risk_for_connection(&item.sql, config.name.as_str(), config.color.as_deref()) - } - ConnectionLookup::Missing => match matching_snapshot_connection_name(item, snapshot) { - Some(connection_name) => risk_for_connection(&item.sql, connection_name, None), - None => conservative_production_risk(&item.sql), - }, - ConnectionLookup::ReadFailed => conservative_production_risk(&item.sql), - }; - item.operation_class = risk.operation_class; - item.risk_level = risk.risk_level; - item.is_production = risk.is_production; -} - -enum ConnectionLookup { - Found(ConnectionConfig), - Missing, - ReadFailed, -} - -async fn load_handoff_connection(app_state: Option<&AppState>, connection_id: &str) -> ConnectionLookup { - let connection_id = connection_id.trim(); - if connection_id.is_empty() { - return ConnectionLookup::Missing; - } - let Some(app_state) = app_state else { - return ConnectionLookup::Missing; - }; - - match app_state.storage.load_connections().await { - Ok(configs) => configs - .into_iter() - .find(|config| config.id == connection_id) - .map(ConnectionLookup::Found) - .unwrap_or(ConnectionLookup::Missing), - Err(err) => { - log::warn!("Agent runtime failed to load connection metadata for handoff risk: {err}"); - ConnectionLookup::ReadFailed - } - } -} - -fn matching_snapshot_connection_name<'a>(item: &HandoffItem, snapshot: &'a AgentRuntimeSnapshot) -> Option<&'a str> { - let handoff_connection_id = item.connection_id.trim(); - let active_connection_id = snapshot.active_connection_id.as_deref().map(str::trim)?; - if handoff_connection_id.is_empty() || handoff_connection_id != active_connection_id { - return None; - } - snapshot.active_connection_name.as_deref().map(str::trim).filter(|name| !name.is_empty()) -} - -fn conservative_production_risk(sql: &str) -> dbx_core::sql_safety::RiskMetadata { - let mut risk = - risk_for(sql, RiskContext { connection_name: "unknown", color: None, environment_label: Some("Production") }); - risk.is_production = true; - risk.risk_level = match classify_sql(sql) { - OperationClass::Ddl => RiskLevel::Critical, - OperationClass::Write if risk.risk_level == RiskLevel::Critical => RiskLevel::Critical, - _ => RiskLevel::High, - }; - risk -} - -async fn update_handoff_status( - app_state: &AppState, - runtime: &AgentRuntimeState, - id: &str, - status: HandoffStatus, -) -> Result { - let stored = app_state.storage.update_handoff_status(id, status.clone()).await?; - let runtime_updated = update_runtime_handoff_status(runtime, id, status).await; - Ok(stored || runtime_updated) -} - -async fn update_runtime_handoff_status(state: &AgentRuntimeState, id: &str, status: HandoffStatus) -> bool { - let mut handoffs = state.handoffs.write().await; - if let Some(item) = handoffs.iter_mut().find(|item| item.id == id) { - if !can_update_runtime_handoff_status(&item.status, &status) { - return false; - } - item.status = status; - return true; - } - false -} - -fn can_update_runtime_handoff_status(current: &HandoffStatus, next: &HandoffStatus) -> bool { - matches!( - (current, next), - (HandoffStatus::Queued | HandoffStatus::Shown, HandoffStatus::Shown) - | (HandoffStatus::Queued | HandoffStatus::Shown, HandoffStatus::Rejected) - | (HandoffStatus::Queued | HandoffStatus::Shown | HandoffStatus::Approved, HandoffStatus::Approved) - | (HandoffStatus::Approved | HandoffStatus::Executed, HandoffStatus::Executed) - | (HandoffStatus::Approved | HandoffStatus::Executed | HandoffStatus::Failed, HandoffStatus::Failed) - ) -} - -async fn pending_runtime_handoffs(state: &AgentRuntimeState) -> Vec { - state - .handoffs - .read() - .await - .iter() - .filter(|item| matches!(item.status, HandoffStatus::Queued | HandoffStatus::Shown)) - .cloned() - .collect() -} - -fn query_limit(first_line: &str) -> Option { - let target = first_line.split_whitespace().nth(1)?; - let query = target.split_once('?')?.1; - query.split('&').find_map(|pair| { - let (key, value) = pair.split_once('=')?; - (key == "limit").then(|| value.parse::().ok()).flatten() - }) -} - -fn truncate_result_rows(result: &mut serde_json::Value, limit: usize) { - if let Some(rows) = result.get_mut("rows").and_then(|rows| rows.as_array_mut()) { - rows.truncate(limit); - } -} - -fn write_discovery_file(dir: &Path, port: u16, token: &str) -> Result { - std::fs::create_dir_all(dir).map_err(|err| err.to_string())?; - let path = dir.join(DISCOVERY_FILE); - let temp_path = dir.join(format!("{DISCOVERY_FILE}.{}.tmp", uuid::Uuid::new_v4())); - let payload = serde_json::json!({ "port": port, "token": token }); - let body = serde_json::to_vec(&payload).map_err(|err| err.to_string())?; - - let mut options = OpenOptions::new(); - options.create_new(true).write(true); - #[cfg(unix)] - { - use std::os::unix::fs::OpenOptionsExt; - options.mode(0o600); - options.custom_flags(libc::O_NOFOLLOW); - } - let mut file = options.open(&temp_path).map_err(|err| err.to_string())?; - file.write_all(&body).map_err(|err| err.to_string())?; - file.sync_all().map_err(|err| err.to_string())?; - - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let mut permissions = file.metadata().map_err(|err| err.to_string())?.permissions(); - permissions.set_mode(0o600); - file.set_permissions(permissions).map_err(|err| err.to_string())?; - } - drop(file); - - if let Err(err) = std::fs::rename(&temp_path, &path) { - let _ = std::fs::remove_file(&temp_path); - return Err(err.to_string()); - } - - Ok(path) -} - -fn cleanup_discovery_file(path: &Path) { - if path.exists() { - let _ = std::fs::remove_file(path); - } -} - -async fn respond_json(stream: &mut TcpStream, status: &str, body: serde_json::Value) { - let body = serde_json::to_string(&body).unwrap_or_else(|_| "{}".to_string()); - let resp = format!( - "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", - body.len() - ); - let _ = stream.write_all(resp.as_bytes()).await; -} - -#[cfg(test)] -mod tests { - use super::*; - use tokio::net::TcpListener; - - fn runtime_state() -> AgentRuntimeState { - AgentRuntimeState { - token: "secret-token".to_string(), - snapshot: Arc::new(RwLock::new(AgentRuntimeSnapshot::default())), - handoffs: Arc::new(RwLock::new(Vec::new())), - app_state: None, - } - } - - async fn runtime_state_with_connections( - configs: Vec, - ) -> AgentRuntimeState { - let db_path = std::env::temp_dir().join(format!("dbx-agent-runtime-test-{}.db", uuid::Uuid::new_v4())); - let storage = dbx_core::storage::Storage::open(&db_path).await.unwrap(); - storage.save_connections(&configs).await.unwrap(); - AgentRuntimeState { app_state: Some(Arc::new(AppState::new(storage))), ..runtime_state() } - } - - fn connection_config(id: &str, name: &str, color: Option<&str>) -> dbx_core::models::connection::ConnectionConfig { - dbx_core::models::connection::ConnectionConfig { - id: id.to_string(), - name: name.to_string(), - db_type: dbx_core::models::connection::DatabaseType::Postgres, - driver_profile: None, - driver_label: None, - url_params: None, - host: "127.0.0.1".to_string(), - port: 5432, - username: "postgres".to_string(), - password: "secret".to_string(), - database: Some("postgres".to_string()), - color: color.map(str::to_string), - ssh_enabled: false, - ssh_host: String::new(), - ssh_port: 22, - ssh_user: String::new(), - ssh_password: String::new(), - ssh_key_path: String::new(), - ssh_key_passphrase: String::new(), - ssh_expose_lan: false, - ssh_connect_timeout_secs: dbx_core::models::connection::default_ssh_connect_timeout_secs(), - proxy_enabled: false, - proxy_type: dbx_core::models::connection::ProxyType::Socks5, - proxy_host: String::new(), - proxy_port: 1080, - proxy_username: String::new(), - proxy_password: String::new(), - ssl: false, - sysdba: false, - connection_string: None, - external_config: None, - jdbc_driver_class: None, - jdbc_driver_paths: Vec::new(), - } - } - - #[cfg(unix)] - fn mode(path: &Path) -> u32 { - use std::os::unix::fs::PermissionsExt; - - std::fs::metadata(path).unwrap().permissions().mode() & 0o777 - } - - async fn serve_once(state: AgentRuntimeState) -> std::net::SocketAddr { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - handle_connection(stream, state).await; - }); - addr - } - - #[test] - fn authorization_requires_exact_bearer_token() { - assert!(is_authorized("GET /context HTTP/1.1\r\nAuthorization: Bearer secret-token\r\n\r\n", "secret-token",)); - assert!(is_authorized("GET /context HTTP/1.1\r\nauthorization: Bearer secret-token\r\n\r\n", "secret-token",)); - assert!(!is_authorized("GET /context HTTP/1.1\r\nAuthorization: Bearer wrong\r\n\r\n", "secret-token",)); - assert!(!is_authorized("GET /context HTTP/1.1\r\n\r\n", "secret-token")); - } - - #[tokio::test] - async fn accepts_reqwest_lowercase_authorization_header() { - let state = runtime_state(); - *state.snapshot.write().await = AgentRuntimeSnapshot { - active_connection_id: Some("conn-1".to_string()), - ..AgentRuntimeSnapshot::default() - }; - let addr = serve_once(state).await; - - let response = reqwest::Client::new() - .get(format!("http://{addr}/context")) - .bearer_auth("secret-token") - .send() - .await - .unwrap(); - - assert_eq!(response.status(), reqwest::StatusCode::OK); - let body: serde_json::Value = response.json().await.unwrap(); - assert_eq!(body["activeConnectionId"], "conn-1"); - } - - #[tokio::test] - async fn routes_context_selection_result_and_handoff_from_shared_state() { - let state = runtime_state(); - *state.snapshot.write().await = AgentRuntimeSnapshot { - active_connection_id: Some("conn-1".to_string()), - active_connection_name: Some("Local".to_string()), - selection: Some(serde_json::json!({"type": "grid-cells", "cells": [[1]]})), - result: Some(serde_json::json!({"columns": ["id"], "rows": [[1]]})), - ..AgentRuntimeSnapshot::default() - }; - - let context = route_request("GET /context HTTP/1.1", "", &state).await; - assert_eq!(context.status, "200 OK"); - assert_eq!(context.body["activeConnectionId"], "conn-1"); - - let selection = route_request("GET /selection HTTP/1.1", "", &state).await; - assert_eq!(selection.status, "200 OK"); - assert_eq!(selection.body["type"], "grid-cells"); - - let result = route_request("GET /result/current?limit=50 HTTP/1.1", "", &state).await; - assert_eq!(result.status, "200 OK"); - assert_eq!(result.body["columns"][0], "id"); - - let item = dbx_core::handoff::HandoffItem::queued( - "conn-1".to_string(), - "Local".to_string(), - Some("main".to_string()), - "Review SQL".to_string(), - None, - "update users set name = 'a'".to_string(), - dbx_core::sql_safety::OperationClass::Write, - dbx_core::sql_safety::RiskLevel::Medium, - false, - ); - let handoff = route_request("POST /handoff HTTP/1.1", &serde_json::to_string(&item).unwrap(), &state).await; - assert_eq!(handoff.status, "200 OK"); - assert_eq!(handoff.body["id"], item.id); - assert_eq!(state.handoffs.read().await.len(), 1); - } - - #[tokio::test] - async fn handoff_recomputes_risk_from_sql_instead_of_trusting_client_fields() { - let state = runtime_state(); - *state.snapshot.write().await = AgentRuntimeSnapshot { - active_connection_id: Some("conn-1".to_string()), - active_connection_name: Some("prod-main".to_string()), - ..AgentRuntimeSnapshot::default() - }; - let item = dbx_core::handoff::HandoffItem::queued( - "conn-1".to_string(), - "client-supplied-dev".to_string(), - Some("main".to_string()), - "Review SQL".to_string(), - None, - "drop table users".to_string(), - dbx_core::sql_safety::OperationClass::Read, - dbx_core::sql_safety::RiskLevel::Low, - false, - ); - - let handoff = route_request("POST /handoff HTTP/1.1", &serde_json::to_string(&item).unwrap(), &state).await; - - assert_eq!(handoff.status, "200 OK"); - let stored = state.handoffs.read().await; - assert_eq!(stored[0].operation_class, dbx_core::sql_safety::OperationClass::Ddl); - assert_eq!(stored[0].risk_level, dbx_core::sql_safety::RiskLevel::Critical); - assert!(stored[0].is_production); - } - - #[tokio::test] - async fn handoff_uses_conservative_production_when_connection_metadata_is_unavailable() { - let state = runtime_state(); - let item = dbx_core::handoff::HandoffItem::queued( - "conn-1".to_string(), - "client-supplied-dev".to_string(), - None, - "Review SQL".to_string(), - None, - "update users set name = 'a' where id = 1".to_string(), - dbx_core::sql_safety::OperationClass::Read, - dbx_core::sql_safety::RiskLevel::Low, - false, - ); - - let handoff = route_request("POST /handoff HTTP/1.1", &serde_json::to_string(&item).unwrap(), &state).await; - - assert_eq!(handoff.status, "200 OK"); - let stored = state.handoffs.read().await; - assert_eq!(stored[0].operation_class, dbx_core::sql_safety::OperationClass::Write); - assert_eq!(stored[0].risk_level, dbx_core::sql_safety::RiskLevel::High); - assert!(stored[0].is_production); - } - - #[tokio::test] - async fn handoff_does_not_use_active_snapshot_when_target_connection_differs() { - let state = runtime_state_with_connections(vec![connection_config("target-dev", "dev-target", None)]).await; - *state.snapshot.write().await = AgentRuntimeSnapshot { - active_connection_id: Some("active-prod".to_string()), - active_connection_name: Some("prod-main".to_string()), - ..AgentRuntimeSnapshot::default() - }; - let item = dbx_core::handoff::HandoffItem::queued( - "target-dev".to_string(), - "dev-target".to_string(), - Some("main".to_string()), - "Review SQL".to_string(), - None, - "update users set name = 'a' where id = 1".to_string(), - dbx_core::sql_safety::OperationClass::Read, - dbx_core::sql_safety::RiskLevel::Low, - false, - ); - - let handoff = route_request("POST /handoff HTTP/1.1", &serde_json::to_string(&item).unwrap(), &state).await; - - assert_eq!(handoff.status, "200 OK"); - let stored = state.handoffs.read().await; - assert_eq!(stored[0].operation_class, dbx_core::sql_safety::OperationClass::Write); - assert_eq!(stored[0].risk_level, dbx_core::sql_safety::RiskLevel::Medium); - assert!(!stored[0].is_production); - } - - #[tokio::test] - async fn result_current_limit_truncates_rows() { - let state = runtime_state(); - *state.snapshot.write().await = AgentRuntimeSnapshot { - result: Some(serde_json::json!({"columns": ["id"], "rows": [[1], [2], [3]]})), - ..AgentRuntimeSnapshot::default() - }; - - let result = route_request("GET /result/current?limit=2 HTTP/1.1", "", &state).await; - - assert_eq!(result.status, "200 OK"); - assert_eq!(result.body["rows"], serde_json::json!([[1], [2]])); - } - - #[tokio::test] - async fn reads_fragmented_handoff_body_until_content_length() { - let state = runtime_state(); - let addr = serve_once(state.clone()).await; - let item = dbx_core::handoff::HandoffItem::queued( - "conn-1".to_string(), - "Local".to_string(), - Some("main".to_string()), - "Review SQL".to_string(), - None, - "select ".to_string() + &"1".repeat(70_000), - dbx_core::sql_safety::OperationClass::Read, - dbx_core::sql_safety::RiskLevel::Low, - false, - ); - let body = serde_json::to_string(&item).unwrap(); - let head = format!( - "POST /handoff HTTP/1.1\r\nHost: {addr}\r\nAuthorization: Bearer secret-token\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n", - body.len() - ); - let split_at = body.len() / 2; - let mut stream = TcpStream::connect(addr).await.unwrap(); - - stream.write_all(head.as_bytes()).await.unwrap(); - stream.write_all(body[..split_at].as_bytes()).await.unwrap(); - tokio::task::yield_now().await; - stream.write_all(body[split_at..].as_bytes()).await.unwrap(); - let mut response = Vec::new(); - stream.read_to_end(&mut response).await.unwrap(); - - let response = String::from_utf8(response).unwrap(); - assert!(response.starts_with("HTTP/1.1 200 OK"), "{response}"); - assert_eq!(state.handoffs.read().await.len(), 1); - } - - #[tokio::test] - async fn rejects_body_larger_than_limit() { - let state = runtime_state(); - let addr = serve_once(state).await; - let body_len = 1_048_577; - let request = format!( - "POST /handoff HTTP/1.1\r\nHost: {addr}\r\nAuthorization: Bearer secret-token\r\nContent-Length: {body_len}\r\n\r\n" - ); - let mut stream = TcpStream::connect(addr).await.unwrap(); - - stream.write_all(request.as_bytes()).await.unwrap(); - let mut response = Vec::new(); - stream.read_to_end(&mut response).await.unwrap(); - - let response = String::from_utf8(response).unwrap(); - assert!(response.starts_with("HTTP/1.1 413 Payload Too Large"), "{response}"); - } - - #[tokio::test] - async fn runtime_handoff_status_updates_filter_pending_items() { - let state = runtime_state(); - let item = dbx_core::handoff::HandoffItem::queued( - "conn-1".to_string(), - "Local".to_string(), - Some("main".to_string()), - "Review SQL".to_string(), - None, - "update users set name = 'a'".to_string(), - dbx_core::sql_safety::OperationClass::Write, - dbx_core::sql_safety::RiskLevel::Medium, - false, - ); - let id = item.id.clone(); - state.handoffs.write().await.push(item); - - assert!(update_runtime_handoff_status(&state, &id, dbx_core::handoff::HandoffStatus::Shown).await); - assert_eq!(pending_runtime_handoffs(&state).await[0].status, dbx_core::handoff::HandoffStatus::Shown); - - assert!(update_runtime_handoff_status(&state, &id, dbx_core::handoff::HandoffStatus::Rejected).await); - assert!(pending_runtime_handoffs(&state).await.is_empty()); - assert!(!update_runtime_handoff_status(&state, &id, dbx_core::handoff::HandoffStatus::Shown).await); - } - - #[test] - fn discovery_file_is_owner_only_and_removed_on_cleanup() { - let dir = std::env::temp_dir().join(format!("dbx-agent-runtime-test-{}", uuid::Uuid::new_v4())); - std::fs::create_dir_all(&dir).unwrap(); - - let path = write_discovery_file(&dir, 4321, "secret-token").unwrap(); - let payload: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); - assert_eq!(payload["port"], 4321); - assert_eq!(payload["token"], "secret-token"); - #[cfg(unix)] - assert_eq!(mode(&path), 0o600); - - cleanup_discovery_file(&path); - assert!(!path.exists()); - - let _ = std::fs::remove_dir_all(dir); - } - - #[cfg(unix)] - #[test] - fn discovery_file_replaces_existing_symlink() { - let dir = std::env::temp_dir().join(format!("dbx-agent-runtime-symlink-test-{}", uuid::Uuid::new_v4())); - std::fs::create_dir_all(&dir).unwrap(); - let target = dir.join("target.json"); - let link = dir.join(DISCOVERY_FILE); - std::fs::write(&target, "{}").unwrap(); - std::os::unix::fs::symlink(&target, &link).unwrap(); - - let path = write_discovery_file(&dir, 4321, "secret-token").unwrap(); - - assert!(!std::fs::symlink_metadata(&path).unwrap().file_type().is_symlink()); - assert_eq!(mode(&path), 0o600); - - let _ = std::fs::remove_dir_all(dir); - } -} diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index 69f61492d..57442c8b2 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -1,4 +1,3 @@ -pub mod agent_runtime; pub mod ai; pub mod connection; #[allow(dead_code, unused_imports)] diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 5d0ff8f50..cc5c97a11 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -38,9 +38,7 @@ pub fn run() { app.manage(state.clone()); let app_handle = app.handle().clone(); - commands::mcp_bridge::start(app_handle, state.clone()); - let runtime_state = commands::agent_runtime::start(app.handle().clone(), state.clone()); - app.manage(runtime_state); + commands::mcp_bridge::start(app_handle, state); #[cfg(not(target_os = "macos"))] { @@ -59,10 +57,6 @@ pub fn run() { } }) .invoke_handler(tauri::generate_handler![ - commands::agent_runtime::agent_runtime_update_snapshot, - commands::agent_runtime::agent_runtime_load_handoffs, - commands::agent_runtime::agent_runtime_mark_handoff_shown, - commands::agent_runtime::agent_runtime_reject_handoff, commands::ai::ai_complete, commands::ai::ai_stream, commands::ai::ai_cancel_stream, @@ -151,12 +145,6 @@ pub fn run() { .build(tauri::generate_context!()) .expect("error while building tauri application") .run(|app_handle, event| { - if let RunEvent::ExitRequested { .. } = &event { - if let Some(runtime) = app_handle.try_state::() { - runtime.cleanup(); - } - } - #[cfg(target_os = "macos")] if let RunEvent::Reopen { has_visible_windows, .. } = &event { if !has_visible_windows { diff --git a/src/App.vue b/src/App.vue index 4db821941..bf8558af8 100644 --- a/src/App.vue +++ b/src/App.vue @@ -15,7 +15,6 @@ import UpdateDialog from "@/components/layout/UpdateDialog.vue"; import LoginPage from "@/components/auth/LoginPage.vue"; import { useConnectionStore } from "@/stores/connectionStore"; import { useQueryStore } from "@/stores/queryStore"; -import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore"; import { useSettingsStore } from "@/stores/settingsStore"; import { useSavedSqlStore } from "@/stores/savedSqlStore"; import { useToast } from "@/composables/useToast"; @@ -38,7 +37,6 @@ import { isCloseTabShortcut, isExecuteSqlShortcut } from "@/lib/keyboardShortcut import { isPreviewTab } from "@/lib/tabPresentation"; import { SQL_FILE_UNSUPPORTED_TYPES } from "@/lib/databaseCapabilities"; import { classifyAiSqlExecution } from "@/lib/aiSqlExecutionPolicy"; -import { restoreStartupAgentRuntime } from "@/lib/appStartup"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; @@ -48,7 +46,6 @@ import type { HistoryEntry } from "@/lib/tauri"; const { t } = useI18n(); const connectionStore = useConnectionStore(); const queryStore = useQueryStore(); -const agentRuntimeStore = useAgentRuntimeStore(); const settingsStore = useSettingsStore(); const savedSqlStore = useSavedSqlStore(); const { message: toastMessage, visible: toastVisible, toast } = useToast(); @@ -161,7 +158,6 @@ watch( () => queryStore.activeTabId, () => { selectedSql.value = ""; - agentRuntimeStore.setSelectedSql(""); activeOutputView.value = "result"; }, ); @@ -449,14 +445,15 @@ function onLoginSuccess() { } function initApp() { - restoreStartupAgentRuntime({ - initSavedSql: () => savedSqlStore.initFromStorage(), - initConnections: () => connectionStore.initFromDisk(), - reconnectRestoredTabs, - scheduleSync: () => agentRuntimeStore.scheduleSync(), - }).catch((e: any) => { - toast(t("connection.loadFailed", { message: e?.message || String(e) }), 5000); - }); + savedSqlStore + .initFromStorage() + .then(() => connectionStore.initFromDisk()) + .then(() => { + reconnectRestoredTabs(); + }) + .catch((e: any) => { + toast(t("connection.loadFailed", { message: e?.message || String(e) }), 5000); + }); settingsStore.initAiConfig(); } @@ -617,12 +614,7 @@ onUnmounted(() => { if (queryStore.activeTabId) queryStore.updateSql(queryStore.activeTabId, v); } " - @editor-selection-change=" - (v: string) => { - selectedSql = v; - agentRuntimeStore.setSelectedSql(v); - } - " + @editor-selection-change="(v: string) => (selectedSql = v)" @editor-cursor-change="(p: number) => (cursorPos = p)" @format-error="toast(t('toolbar.formatSqlFailed'))" @reload=" diff --git a/src/components/agent/AgentHandoffDialog.vue b/src/components/agent/AgentHandoffDialog.vue deleted file mode 100644 index 610cd143c..000000000 --- a/src/components/agent/AgentHandoffDialog.vue +++ /dev/null @@ -1,78 +0,0 @@ - - - diff --git a/src/components/grid/DataGrid.vue b/src/components/grid/DataGrid.vue index a17e17a3e..12c5a7e1a 100644 --- a/src/components/grid/DataGrid.vue +++ b/src/components/grid/DataGrid.vue @@ -69,10 +69,9 @@ import { quoteTableIdentifier, } from "@/lib/tableSelectSql"; import { isHiddenGridColumn, usesSyntheticRowIdKey } from "@/lib/tableEditing"; -import { displayCellValue, type CellValue } from "@/lib/cellValue"; import { formatGridSqlLiteral } from "@/lib/dataGridSql"; import { matchesRowStatusFilter, type RowStatus, type RowStatusFilter } from "@/lib/gridRowStatus"; -import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore"; +import { displayCellValue, type CellValue } from "@/lib/cellValue"; import { useToast } from "@/composables/useToast"; import { useDataGridExport } from "@/composables/useDataGridExport"; @@ -82,7 +81,6 @@ import { useDataGridEditor } from "@/composables/useDataGridEditor"; const { t } = useI18n(); const { toast } = useToast(); -const agentRuntimeStore = useAgentRuntimeStore(); const props = defineProps<{ result: QueryResult; @@ -1136,22 +1134,6 @@ const activeCellDetail = computed(() => { const detailEditValue = ref(""); const isEditingDetail = ref(false); -watch( - selectedCells, - (data) => { - agentRuntimeStore.setSelection( - selectedCellCount.value > 0 - ? { - type: "grid-cells", - range: selectedRange.value, - data, - } - : { type: "none" }, - ); - }, - { deep: true }, -); - function startDetailEdit() { const detail = activeCellDetail.value; if (!detail || !detail.isEditable) return; @@ -1673,7 +1655,6 @@ watch( onUnmounted(() => { cleanupFrames(); - agentRuntimeStore.setSelection({ type: "none" }); onDdlResizeEnd(); finishCellSelection(); clearTimeout(_searchTimer); diff --git a/src/components/layout/AppDialogs.vue b/src/components/layout/AppDialogs.vue index 159fc25f0..0061e9d12 100644 --- a/src/components/layout/AppDialogs.vue +++ b/src/components/layout/AppDialogs.vue @@ -4,7 +4,6 @@ import { useI18n } from "vue-i18n"; import { Button } from "@/components/ui/button"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import ConnectionDialog from "@/components/connection/ConnectionDialog.vue"; -import AgentHandoffDialog from "@/components/agent/AgentHandoffDialog.vue"; import EditorSettingsDialog from "@/components/editor/EditorSettingsDialog.vue"; import DangerConfirmDialog from "@/components/editor/DangerConfirmDialog.vue"; const DataTransferDialog = defineAsyncComponent(() => import("@/components/transfer/DataTransferDialog.vue")); @@ -93,7 +92,6 @@ watch(