xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Revert "feat(cli): add dbx-cli runtime and agent integration (#229)" 

This reverts commit 58c47adef5, reversing
changes made to 3bcec8ef7f.
This commit is contained in:
t8y2 2026-05-12 17:31:54 +08:00
parent 58c47adef5
commit bf58045a28
37 changed files with 56 additions and 6679 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -40,7 +40,6 @@ Thumbs.db
# Temp
tmp/
.worktrees/
# Logs
*.log

16
Cargo.lock generated
View File

@ -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"

2
Cargo.toml
View File

@ -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"

21
crates/dbx-cli/Cargo.toml
View File

@ -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"

1617
crates/dbx-cli/src/commands.rs
View File

File diff suppressed because it is too large Load Diff

10
crates/dbx-cli/src/main.rs
View File

@ -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);
}
}

190
crates/dbx-cli/src/runtime_client.rs
View File

@ -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<RuntimeDiscovery> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<reqwest::Url, String> {
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");
}
}

345
crates/dbx-core/src/cli.rs
View File

@ -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<T> {
Success { ok: bool, source: CliSource, data: T },
Failure { ok: bool, source: CliSource, error: CliError },
}
impl<'de, T> Deserialize<'de> for CliEnvelope<T>
where
T: Deserialize<'de>,
{
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
deserializer.deserialize_map(CliEnvelopeVisitor { marker: PhantomData })
}
}
struct CliEnvelopeVisitor<T> {
marker: PhantomData<T>,
}
impl<'de, T> Visitor<'de> for CliEnvelopeVisitor<T>
where
T: Deserialize<'de>,
{
type Value = CliEnvelope<T>;
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
formatter.write_str("a CLI envelope with consistent ok/data/error fields")
}
fn visit_map<M>(self, mut map: M) -> Result<Self::Value, M::Error>
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::<CliEnvelopeField>()? {
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::<IgnoredAny>()?;
}
}
}
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<D>(deserializer: D) -> Result<Self, D::Error>
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<E>(self, value: &str) -> Result<Self::Value, E>
where
E: de::Error,
{
Ok(match value {
"ok" => CliEnvelopeField::Ok,
"source" => CliEnvelopeField::Source,
"data" => CliEnvelopeField::Data,
"error" => CliEnvelopeField::Error,
_ => CliEnvelopeField::Ignore,
})
}
}
pub fn ok<T>(source: CliSource, data: T) -> CliEnvelope<T> {
CliEnvelope::Success { ok: true, source, data }
}
pub fn fail<T>(source: CliSource, code: CliErrorCode, message: impl Into<String>, recoverable: bool) -> CliEnvelope<T> {
CliEnvelope::Failure { ok: false, source, error: CliError { code, message: message.into(), recoverable } }
}
pub fn fail_safe<T>(
source: CliSource,
fallback_code: CliErrorCode,
message: impl AsRef<str>,
recoverable: bool,
) -> CliEnvelope<T> {
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::Value> = 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::Value> = 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::<CliEnvelope<serde_json::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::<CliEnvelope<serde_json::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::<CliEnvelope<serde_json::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.");
}
}

27
crates/dbx-core/src/db/mysql.rs
View File

@ -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<QueryResult, String> {
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<QueryResult, String> {
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));

42
crates/dbx-core/src/db/postgres.rs
View File

@ -281,12 +281,7 @@ pub async fn get_columns(pool: &PgPool, schema: &str, table: &str) -> Result<Vec
}
pub async fn execute_query(pool: &PgPool, sql: &str) -> Result<QueryResult, String> {
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<QueryResult, 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(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<QueryResult, String> {
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<QueryResult, String> {
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<Ve
.collect())
}
#[cfg(test)]
mod tests {
use super::*;
#[allow(dead_code)]
fn postgres_limit_aware_execute_query_api_compiles(pool: &PgPool, sql: &str) {
let _ = execute_query_with_row_limit(pool, sql, 7);
}
#[allow(dead_code)]
fn postgres_schema_limit_aware_execute_query_api_compiles(pool: &PgPool, schema: &str, sql: &str) {
let _ = execute_query_with_schema_and_row_limit(pool, schema, sql, 7);
}
}
pub async fn list_foreign_keys(pool: &PgPool, schema: &str, table: &str) -> Result<Vec<ForeignKeyInfo>, String> {
let rows: Vec<PgRow> = sqlx::query(
"SELECT kcu.constraint_name, kcu.column_name, \

15
crates/dbx-core/src/db/sqlite.rs
View File

@ -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<QueryResult, String> {
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<QueryResult, String> {
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 {

69
crates/dbx-core/src/handoff.rs
View File

@ -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<Utc>,
pub created_by: String,
#[serde(default)]
pub connection_id: String,
pub connection_name: String,
pub database: Option<String>,
pub title: String,
pub description: Option<String>,
pub sql: String,
pub operation_class: OperationClass,
pub risk_level: RiskLevel,
pub is_production: bool,
pub status: HandoffStatus,
pub result_summary: Option<String>,
pub error: Option<String>,
}
impl HandoffItem {
pub fn queued(
connection_id: String,
connection_name: String,
database: Option<String>,
title: String,
description: Option<String>,
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,
}
}
}

4
crates/dbx-core/src/lib.rs
View File

@ -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;

233
crates/dbx-core/src/query.rs
View File

@ -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<db::QueryResult, String> {
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<db::QueryResult, String> {
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<serde_json::Value> = (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<CancellationToken>,
) -> Result<db::QueryResult, String> {
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<CancellationToken>,
row_limit: usize,
) -> Result<db::QueryResult, String> {
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::<db::QueryResult>("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<CancellationToken>,
) -> Result<db::QueryResult, String> {
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<CancellationToken>,
row_limit: usize,
) -> Result<db::QueryResult, String> {
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);
}
}

164
crates/dbx-core/src/schema_snapshot.rs
View File

@ -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<String>,
pub columns: Vec<types::ColumnInfo>,
pub indexes: Vec<types::IndexInfo>,
pub foreign_keys: Vec<types::ForeignKeyInfo>,
pub triggers: Vec<types::TriggerInfo>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SchemaSnapshot {
pub connection_id: String,
pub connection_name: String,
pub database: Option<String>,
pub database_type: DatabaseType,
pub driver_profile: Option<String>,
pub captured_at: DateTime<Utc>,
pub databases: Vec<types::DatabaseInfo>,
pub schemas: Vec<String>,
pub tables: Vec<TableSnapshot>,
}
pub async fn snapshot(
state: &AppState,
connection_id: &str,
database: Option<&str>,
schema_name: Option<&str>,
) -> Result<SchemaSnapshot, String> {
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<types::TableInfo>,
) -> Result<Vec<TableSnapshot>, 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<TableSnapshot, String> {
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<String, String> {
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
)
}

490
crates/dbx-core/src/sql_safety.rs
View File

@ -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<String>,
pub first_token: Option<String>,
}
#[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<String>) {
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<String> {
executable_statements(sql).into_iter().flatten().collect()
}
fn executable_statements(sql: &str) -> Vec<Vec<String>> {
scanned_executable_statements(sql)
.into_iter()
.map(|statement| statement.into_iter().map(|token| token.text).collect())
.collect()
}
fn scanned_executable_statements(sql: &str) -> Vec<Vec<SqlToken>> {
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<SqlToken>,
statements: &mut Vec<Vec<SqlToken>>,
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<SqlToken>, statements: &mut Vec<Vec<SqlToken>>) {
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<usize> {
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<String> {
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"));
}
}

295
crates/dbx-core/src/storage.rs
View File

@ -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<bool, String> {
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<Vec<HandoffItem>, 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<String, String> {
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<_>>(), 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<_>>(), 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<_>>(), vec![queued.id, shown.id]);
}
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

140
crates/dbx-core/tests/schema_snapshot.rs
View File

@ -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<_>>(), vec!["main"]);
let table_names = snapshot.tables.iter().map(|table| table.name.as_str()).collect::<Vec<_>>();
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}");
}

1537
docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md
View File

File diff suppressed because it is too large Load Diff

1
src-tauri/Cargo.toml
View File

@ -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" }

889
src-tauri/src/commands/agent_runtime.rs
View File

@ -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<String>,
pub active_connection_name: Option<String>,
pub database: Option<String>,
pub schema: Option<String>,
pub active_tab_id: Option<String>,
pub active_tab_title: Option<String>,
pub sql: Option<String>,
pub selected_sql: Option<String>,
pub selection: Option<serde_json::Value>,
pub result: Option<serde_json::Value>,
}
#[derive(Clone)]
pub struct AgentRuntimeState {
pub token: String,
pub snapshot: Arc<RwLock<AgentRuntimeSnapshot>>,
pub handoffs: Arc<RwLock<Vec<dbx_core::handoff::HandoffItem>>>,
pub app_state: Option<Arc<AppState>>,
}
pub struct AgentRuntimeServer {
state: AgentRuntimeState,
discovery_path: PathBuf,
shutdown: std::sync::Mutex<Option<oneshot::Sender<()>>>,
}
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<AppState>>,
runtime: tauri::State<'_, AgentRuntimeServer>,
) -> Result<Vec<dbx_core::handoff::HandoffItem>, 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<AppState>>,
runtime: tauri::State<'_, AgentRuntimeServer>,
id: String,
) -> Result<bool, String> {
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<AppState>>,
runtime: tauri::State<'_, AgentRuntimeServer>,
id: String,
) -> Result<bool, String> {
update_handoff_status(app_state.inner().as_ref(), runtime.state(), &id, HandoffStatus::Rejected).await
}
pub fn start(app: AppHandle, app_state: Arc<AppState>) -> 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<Option<RuntimeRequest>, 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<usize> {
buf.windows(4).position(|window| window == b"\r\n\r\n")
}
fn content_length(headers: &[(String, String)]) -> Result<usize, RuntimeResponse> {
match headers.iter().find(|(name, _)| name.eq_ignore_ascii_case("content-length")) {
Some((_, value)) => value.parse::<usize>().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::<dbx_core::handoff::HandoffItem>(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<bool, String> {
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<HandoffItem> {
state
.handoffs
.read()
.await
.iter()
.filter(|item| matches!(item.status, HandoffStatus::Queued | HandoffStatus::Shown))
.cloned()
.collect()
}
fn query_limit(first_line: &str) -> Option<usize> {
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::<usize>().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<PathBuf, String> {
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<dbx_core::models::connection::ConnectionConfig>,
) -> 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);
}
}

1
src-tauri/src/commands/mod.rs
View File

@ -1,4 +1,3 @@
pub mod agent_runtime;
pub mod ai;
pub mod connection;
#[allow(dead_code, unused_imports)]

14
src-tauri/src/lib.rs
View File

@ -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::<commands::agent_runtime::AgentRuntimeServer>() {
runtime.cleanup();
}
}
#[cfg(target_os = "macos")]
if let RunEvent::Reopen { has_visible_windows, .. } = &event {
if !has_visible_windows {

28
src/App.vue
View File

@ -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="

78
src/components/agent/AgentHandoffDialog.vue
View File

@ -1,78 +0,0 @@
<script setup lang="ts">
import { computed, onMounted, onUnmounted } from "vue";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore";
const agentRuntimeStore = useAgentRuntimeStore();
let refreshTimer: number | null = null;
const active = computed(() => agentRuntimeStore.activeHandoff);
const riskTone = computed(() => {
if (!active.value) return "secondary";
return active.value.riskLevel === "critical" || active.value.isProduction ? "destructive" : "secondary";
});
async function refresh() {
try {
await agentRuntimeStore.loadHandoffs();
} catch (err) {
console.debug("[DBX] Agent handoff refresh skipped:", err);
}
}
async function rejectActive() {
if (!active.value) return;
await agentRuntimeStore.rejectHandoff(active.value.id);
}
onMounted(() => {
void refresh();
refreshTimer = window.setInterval(() => void refresh(), 5000);
});
onUnmounted(() => {
if (refreshTimer) window.clearInterval(refreshTimer);
});
</script>
<template>
<Dialog v-model:open="agentRuntimeStore.handoffDialogOpen">
<DialogContent class="max-w-3xl">
<DialogHeader>
<DialogTitle>DBX Agent Handoff</DialogTitle>
</DialogHeader>
<div v-if="active" class="space-y-4">
<div class="flex flex-wrap items-center gap-2 text-xs text-muted-foreground">
<Badge :variant="riskTone">{{ active.riskLevel }}</Badge>
<span>{{ active.connectionName }}</span>
<span v-if="active.database">/ {{ active.database }}</span>
<span>/ {{ active.operationClass }}</span>
<span v-if="active.status === 'shown'">/ shown</span>
</div>
<div class="space-y-1">
<h3 class="text-sm font-semibold">{{ active.title }}</h3>
<p v-if="active.description" class="text-sm text-muted-foreground">{{ active.description }}</p>
</div>
<pre class="max-h-96 overflow-auto rounded-md border bg-muted/60 p-3 text-xs leading-relaxed">{{
active.sql
}}</pre>
<p class="text-xs text-muted-foreground">
Review only: DBX does not execute agent handoff SQL from this dialog.
</p>
</div>
<div v-else class="text-sm text-muted-foreground">No pending agent handoffs.</div>
<DialogFooter>
<Button variant="outline" @click="agentRuntimeStore.handoffDialogOpen = false">Close</Button>
<Button v-if="active" variant="destructive" @click="rejectActive">Reject</Button>
</DialogFooter>
</DialogContent>
</Dialog>
</template>

21
src/components/grid/DataGrid.vue
View File

@ -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);

2
src/components/layout/AppDialogs.vue
View File

@ -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(
</script>
<template>
<AgentHandoffDialog />
<ConnectionDialog
:open="showConnectionDialog"
:edit-config="editConfig"

48
src/lib/agentHandoff.ts
View File

@ -1,48 +0,0 @@
export type AgentHandoffStatus = "queued" | "shown" | "approved" | "rejected" | "executed" | "failed";
export interface AgentHandoffItem {
id: string;
createdAt: string;
createdBy: string;
connectionId: string;
connectionName: string;
database?: string;
title: string;
description?: string;
sql: string;
operationClass: string;
riskLevel: string;
isProduction: boolean;
status: AgentHandoffStatus;
resultSummary?: string;
error?: string;
}
export function isPendingHandoff(item: AgentHandoffItem): boolean {
return item.status === "queued" || item.status === "shown";
}
export function mergeLoadedHandoffs(
items: AgentHandoffItem[],
ignoredIds: ReadonlySet<string> = new Set(),
): AgentHandoffItem[] {
return items
.filter((item) => isPendingHandoff(item) && !ignoredIds.has(item.id))
.sort((a, b) => Date.parse(a.createdAt) - Date.parse(b.createdAt));
}
export function updateHandoffStatus(
items: AgentHandoffItem[],
id: string,
status: AgentHandoffStatus,
): AgentHandoffItem[] {
return mergeLoadedHandoffs(items.map((item) => (item.id === id ? { ...item, status } : item)));
}
export function deriveHandoffDialogState(items: AgentHandoffItem[], activeId: string | null) {
const active = (activeId ? items.find((item) => item.id === activeId) : undefined) ?? items[0] ?? null;
return {
open: items.length > 0,
active,
};
}

62
src/lib/agentRuntimeSnapshot.ts
View File

@ -1,62 +0,0 @@
import type { ConnectionConfig, QueryResult, QueryTab } from "@/types/database";
export const DEFAULT_RESULT_SAMPLE_LIMIT = 50;
export interface RuntimeResultSample {
columns: string[];
rows: QueryResult["rows"];
truncated: boolean;
executionTimeMs: number;
sampleLimit: number;
}
export interface AgentRuntimeSnapshot {
activeConnectionId?: string;
activeConnectionName?: string;
database?: string;
schema?: string;
activeTabId?: string;
activeTabTitle?: string;
sql?: string;
selectedSql?: string;
selection: unknown;
result?: RuntimeResultSample;
}
export interface BuildAgentRuntimeSnapshotOptions {
tabs: QueryTab[];
activeTabId: string | null;
getConnection: (connectionId: string) => Pick<ConnectionConfig, "name"> | undefined;
selectedSql?: string;
selection?: unknown;
resultSampleLimit?: number;
}
export function buildAgentRuntimeSnapshot(options: BuildAgentRuntimeSnapshotOptions): AgentRuntimeSnapshot {
const tab = options.tabs.find((item) => item.id === options.activeTabId);
const conn = tab ? options.getConnection(tab.connectionId) : undefined;
const selectedSql = options.selectedSql?.trim();
const limit = options.resultSampleLimit ?? DEFAULT_RESULT_SAMPLE_LIMIT;
const result = tab?.result
? {
columns: tab.result.columns,
rows: tab.result.rows.slice(0, limit),
truncated: tab.result.rows.length > limit || !!tab.result.truncated,
executionTimeMs: tab.result.execution_time_ms,
sampleLimit: limit,
}
: undefined;
return {
activeConnectionId: tab?.connectionId,
activeConnectionName: conn?.name,
database: tab?.database,
schema: tab?.schema,
activeTabId: tab?.id,
activeTabTitle: tab?.title,
sql: tab?.sql,
selectedSql: selectedSql || undefined,
selection: options.selection ?? { type: "none" },
result,
};
}

6
src/lib/api.ts
View File

@ -72,12 +72,6 @@ export const executeScript = forward("executeScript");
export const executeInTransaction = forward("executeInTransaction");
export const cancelQuery = forward("cancelQuery");
// Agent Runtime
export const agentRuntimeUpdateSnapshot = forward("agentRuntimeUpdateSnapshot");
export const agentRuntimeLoadHandoffs = forward("agentRuntimeLoadHandoffs");
export const agentRuntimeMarkHandoffShown = forward("agentRuntimeMarkHandoffShown");
export const agentRuntimeRejectHandoff = forward("agentRuntimeRejectHandoff");
// AI
export const aiComplete = forward("aiComplete");
export const aiStream = forward("aiStream");

13
src/lib/appStartup.ts
View File

@ -1,13 +0,0 @@
export interface RestoreStartupAgentRuntimeOptions {
initSavedSql: () => Promise<unknown>;
initConnections: () => Promise<unknown>;
reconnectRestoredTabs: () => Promise<unknown> | unknown;
scheduleSync: () => void;
}
export async function restoreStartupAgentRuntime(options: RestoreStartupAgentRuntimeOptions) {
await options.initSavedSql();
await options.initConnections();
await options.reconnectRestoredTabs();
options.scheduleSync();
}

16
src/lib/http.ts
View File

@ -39,8 +39,6 @@ import type {
TableImportSummary,
TableImportProgress,
} from "./tauri";
import type { AgentRuntimeSnapshot } from "@/lib/agentRuntimeSnapshot";
import type { AgentHandoffItem } from "@/lib/agentHandoff";
// ---------------------------------------------------------------------------
// Helpers
@ -292,20 +290,6 @@ export async function cancelQuery(executionId: string): Promise<boolean> {
return post("/api/query/cancel", { executionId });
}
export async function agentRuntimeUpdateSnapshot(_snapshot: AgentRuntimeSnapshot): Promise<void> {}
export async function agentRuntimeLoadHandoffs(): Promise<AgentHandoffItem[]> {
return [];
}
export async function agentRuntimeMarkHandoffShown(_id: string): Promise<boolean> {
return false;
}
export async function agentRuntimeRejectHandoff(_id: string): Promise<boolean> {
return false;
}
// ---------------------------------------------------------------------------
// AI
// ---------------------------------------------------------------------------

18
src/lib/tauri.ts
View File

@ -19,8 +19,6 @@ import type {
SavedSqlFolder,
SavedSqlLibrary,
} from "@/types/database";
import type { AgentRuntimeSnapshot } from "@/lib/agentRuntimeSnapshot";
import type { AgentHandoffItem } from "@/lib/agentHandoff";
import type { AiConfig } from "@/stores/settingsStore";
export interface AiMessage {
@ -82,22 +80,6 @@ export async function loadAiConfig(): Promise<AiConfig | null> {
return invoke("load_ai_config");
}
export async function agentRuntimeUpdateSnapshot(snapshot: AgentRuntimeSnapshot): Promise<void> {
return invoke("agent_runtime_update_snapshot", { snapshot });
}
export async function agentRuntimeLoadHandoffs(): Promise<AgentHandoffItem[]> {
return invoke("agent_runtime_load_handoffs");
}
export async function agentRuntimeMarkHandoffShown(id: string): Promise<boolean> {
return invoke("agent_runtime_mark_handoff_shown", { id });
}
export async function agentRuntimeRejectHandoff(id: string): Promise<boolean> {
return invoke("agent_runtime_reject_handoff", { id });
}
// --- AI Conversations ---
export interface AiChatMessage {

109
src/stores/agentRuntimeStore.ts
View File

@ -1,109 +0,0 @@
import { defineStore } from "pinia";
import { computed, ref } from "vue";
import * as api from "@/lib/api";
import {
deriveHandoffDialogState,
mergeLoadedHandoffs,
updateHandoffStatus,
type AgentHandoffItem,
} from "@/lib/agentHandoff";
import { buildAgentRuntimeSnapshot, DEFAULT_RESULT_SAMPLE_LIMIT } from "@/lib/agentRuntimeSnapshot";
import { useConnectionStore } from "@/stores/connectionStore";
import { useQueryStore } from "@/stores/queryStore";
export const useAgentRuntimeStore = defineStore("agentRuntime", () => {
const selection = ref<unknown>({ type: "none" });
const selectedSql = ref("");
const handoffs = ref<AgentHandoffItem[]>([]);
const activeHandoffId = ref<string | null>(null);
const handoffDialogOpen = ref(false);
const locallyClosedHandoffIds = new Set<string>();
let timer: ReturnType<typeof setTimeout> | null = null;
const activeHandoff = computed(() => handoffs.value.find((item) => item.id === activeHandoffId.value) ?? null);
function setSelection(value: unknown) {
selection.value = value;
scheduleSync();
}
function setSelectedSql(value: string) {
selectedSql.value = value;
scheduleSync();
}
function scheduleSync() {
if (timer) clearTimeout(timer);
timer = setTimeout(() => {
timer = null;
void syncNow();
}, 100);
}
async function syncNow() {
if (!globalThis.localStorage) return;
try {
const connectionStore = useConnectionStore();
const queryStore = useQueryStore();
const snapshot = buildAgentRuntimeSnapshot({
tabs: queryStore.tabs,
activeTabId: queryStore.activeTabId,
getConnection: (connectionId) => connectionStore.getConfig(connectionId),
selectedSql: selectedSql.value,
selection: selection.value,
resultSampleLimit: DEFAULT_RESULT_SAMPLE_LIMIT,
});
await api.agentRuntimeUpdateSnapshot(snapshot);
} catch (err) {
console.debug("[DBX] Agent runtime snapshot sync skipped:", err);
}
}
async function loadHandoffs() {
const loaded = mergeLoadedHandoffs(await api.agentRuntimeLoadHandoffs(), locallyClosedHandoffIds);
handoffs.value = loaded;
const state = deriveHandoffDialogState(loaded, activeHandoffId.value);
activeHandoffId.value = state.active?.id ?? null;
handoffDialogOpen.value = state.open;
if (state.active?.status === "queued") {
await markHandoffShown(state.active.id);
}
}
async function markHandoffShown(id: string) {
handoffs.value = updateHandoffStatus(handoffs.value, id, "shown");
await api.agentRuntimeMarkHandoffShown(id);
}
async function rejectHandoff(id: string) {
await api.agentRuntimeRejectHandoff(id);
locallyClosedHandoffIds.add(id);
handoffs.value = updateHandoffStatus(handoffs.value, id, "rejected");
const state = deriveHandoffDialogState(handoffs.value, activeHandoffId.value === id ? null : activeHandoffId.value);
activeHandoffId.value = state.active?.id ?? null;
handoffDialogOpen.value = state.open;
}
function setActiveHandoff(id: string) {
activeHandoffId.value = id;
const item = handoffs.value.find((handoff) => handoff.id === id);
if (item?.status === "queued") void markHandoffShown(id);
}
return {
selection,
selectedSql,
handoffs,
activeHandoff,
handoffDialogOpen,
setSelection,
setSelectedSql,
scheduleSync,
syncNow,
loadHandoffs,
markHandoffShown,
rejectHandoff,
setActiveHandoff,
};
});

35
src/stores/queryStore.ts
View File

@ -10,31 +10,22 @@ import { analyzeEditableQuery, allPrimaryKeysPresent } from "@/lib/sqlAnalysis";
import { restoreOpenTabsState, serializeOpenTabs } from "@/lib/openTabsPersistence";
import * as api from "@/lib/api";
import { useConnectionStore } from "@/stores/connectionStore";
import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore";
import { isTauriRuntime } from "@/lib/tauriRuntime";
import type { SavedSqlFile } from "@/types/database";
const STORAGE_KEY = "dbx-open-tabs";
const ACTIVE_TAB_KEY = "dbx-active-tab";
function browserStorage(): Storage | undefined {
return globalThis.localStorage;
}
function saveTabs(tabs: QueryTab[], activeTabId: string | null) {
const storage = browserStorage();
if (!storage) return;
try {
storage.setItem(STORAGE_KEY, JSON.stringify(serializeOpenTabs(tabs)));
storage.setItem(ACTIVE_TAB_KEY, activeTabId || "");
localStorage.setItem(STORAGE_KEY, JSON.stringify(serializeOpenTabs(tabs)));
localStorage.setItem(ACTIVE_TAB_KEY, activeTabId || "");
} catch {}
}
function loadSavedTabs(): { tabs: QueryTab[]; activeTabId: string | null } {
const storage = browserStorage();
if (!storage) return { tabs: [], activeTabId: null };
try {
return restoreOpenTabsState(storage.getItem(STORAGE_KEY), storage.getItem(ACTIVE_TAB_KEY), {
return restoreOpenTabsState(localStorage.getItem(STORAGE_KEY), localStorage.getItem(ACTIVE_TAB_KEY), {
queryOnly: isTauriRuntime(),
});
} catch {
@ -81,18 +72,11 @@ export const useQueryStore = defineStore("query", () => {
return tabs.value.find((t) => t.connectionId === connectionId && t.database === database && t.title === title);
}
function scheduleAgentRuntimeSync() {
useAgentRuntimeStore().scheduleSync();
}
watch(activeTabId, () => scheduleAgentRuntimeSync());
function createTab(connectionId: string, database: string, title?: string, mode: QueryTab["mode"] = "query") {
if (title) {
const existing = findTabByTitle(connectionId, database, title);
if (existing) {
activeTabId.value = existing.id;
scheduleAgentRuntimeSync();
return existing.id;
}
}
@ -111,7 +95,6 @@ export const useQueryStore = defineStore("query", () => {
};
tabs.value.push(tab);
activeTabId.value = id;
scheduleAgentRuntimeSync();
return id;
}
@ -162,7 +145,6 @@ export const useQueryStore = defineStore("query", () => {
if (activeTabId.value === id) {
activeTabId.value = tabs.value[Math.min(idx, tabs.value.length - 1)]?.id ?? null;
}
scheduleAgentRuntimeSync();
}
function closeOtherTabs(id: string) {
@ -171,7 +153,6 @@ export const useQueryStore = defineStore("query", () => {
const next = closeOtherTabsState(tabs.value, activeTabId.value, id);
tabs.value = next.tabs;
activeTabId.value = next.activeTabId;
scheduleAgentRuntimeSync();
}
function closeAllTabs() {
@ -180,7 +161,6 @@ export const useQueryStore = defineStore("query", () => {
const next = closeAllTabsState(tabs.value, activeTabId.value);
tabs.value = next.tabs;
activeTabId.value = next.activeTabId;
scheduleAgentRuntimeSync();
}
function updateSql(id: string, sql: string) {
@ -189,7 +169,6 @@ export const useQueryStore = defineStore("query", () => {
tab.sql = sql;
tab.resultSortedSql = undefined;
tab.resultBaseSql = undefined;
scheduleAgentRuntimeSync();
}
}
@ -204,7 +183,6 @@ export const useQueryStore = defineStore("query", () => {
const existing = tabs.value.find((tab) => tab.savedSqlId === file.id);
if (existing) {
activeTabId.value = existing.id;
scheduleAgentRuntimeSync();
return existing.id;
}
@ -224,7 +202,6 @@ export const useQueryStore = defineStore("query", () => {
};
tabs.value.push(tab);
activeTabId.value = id;
scheduleAgentRuntimeSync();
return id;
}
@ -247,7 +224,6 @@ export const useQueryStore = defineStore("query", () => {
tab.resultSortedSql = undefined;
clearExplain(tab);
tab.tableMeta = undefined;
scheduleAgentRuntimeSync();
}
function updateSchema(id: string, schema: string | undefined) {
@ -255,7 +231,6 @@ export const useQueryStore = defineStore("query", () => {
if (!tab || tab.schema === schema) return;
tab.schema = schema;
if (tab.mode === "objects") tab.objectBrowser = { ...tab.objectBrowser, schema };
scheduleAgentRuntimeSync();
}
function updateConnection(id: string, connectionId: string, database = "") {
@ -270,7 +245,6 @@ export const useQueryStore = defineStore("query", () => {
tab.resultSortedSql = undefined;
clearExplain(tab);
tab.tableMeta = undefined;
scheduleAgentRuntimeSync();
}
function setTableMeta(id: string, meta: NonNullable<QueryTab["tableMeta"]>) {
@ -436,7 +410,6 @@ export const useQueryStore = defineStore("query", () => {
console.info("[DBX][executeTabSql:metadata:start]", { traceId, elapsed: elapsed() });
await analyzeQueryMetadata(current, current.resultBaseSql);
console.info("[DBX][executeTabSql:metadata:done]", { traceId, elapsed: elapsed() });
scheduleAgentRuntimeSync();
} else {
console.warn("[DBX][executeTabSql:stale-result]", {
traceId,
@ -455,7 +428,6 @@ export const useQueryStore = defineStore("query", () => {
if (current.mode !== "data") current.tableMeta = undefined;
current.resultBaseSql = options?.resultBaseSql ?? sql;
current.resultSortedSql = options?.resultSortedSql;
scheduleAgentRuntimeSync();
}
} finally {
const current = tabs.value.find((t) => t.id === id);
@ -566,7 +538,6 @@ export const useQueryStore = defineStore("query", () => {
if (!tab?.results || index < 0 || index >= tab.results.length) return;
tab.activeResultIndex = index;
tab.result = tab.results[index];
scheduleAgentRuntimeSync();
}
function trimResultCache() {

74
tests/agentHandoff.test.ts
View File

@ -1,74 +0,0 @@
import { strict as assert } from "node:assert";
import test from "node:test";
import {
deriveHandoffDialogState,
mergeLoadedHandoffs,
updateHandoffStatus,
type AgentHandoffItem,
} from "../src/lib/agentHandoff.ts";
function handoff(id: string, status: AgentHandoffItem["status"]): AgentHandoffItem {
return {
id,
createdAt: `2026-05-10T00:00:0${id}.000Z`,
createdBy: "dbx-cli",
connectionId: `conn-${id}`,
connectionName: `Connection ${id}`,
database: "main",
title: `Review ${id}`,
sql: "UPDATE users SET active = 0",
operationClass: "write",
riskLevel: "high",
isProduction: true,
status,
};
}
test("mergeLoadedHandoffs keeps only queued and shown records in FIFO order", () => {
const merged = mergeLoadedHandoffs([
handoff("3", "rejected"),
handoff("2", "shown"),
handoff("1", "queued"),
handoff("4", "executed"),
]);
assert.deepEqual(
merged.map((item) => [item.id, item.status]),
[
["1", "queued"],
["2", "shown"],
],
);
});
test("deriveHandoffDialogState opens the first pending handoff", () => {
const state = deriveHandoffDialogState([handoff("1", "queued"), handoff("2", "shown")], null);
assert.equal(state.open, true);
assert.equal(state.active?.id, "1");
});
test("updateHandoffStatus marks shown and removes rejected handoffs from pending view", () => {
const shown = updateHandoffStatus([handoff("1", "queued")], "1", "shown");
assert.equal(shown[0].status, "shown");
const rejected = updateHandoffStatus(shown, "1", "rejected");
assert.deepEqual(rejected, []);
});
test("updateHandoffStatus does not let shown overwrite rejected handoffs", () => {
const rejected = updateHandoffStatus([handoff("1", "shown")], "1", "rejected");
const staleShown = updateHandoffStatus(rejected, "1", "shown");
assert.deepEqual(staleShown, []);
});
test("mergeLoadedHandoffs can ignore locally closed handoffs from stale loads", () => {
const merged = mergeLoadedHandoffs([handoff("1", "shown"), handoff("2", "queued")], new Set(["1"]));
assert.deepEqual(
merged.map((item) => item.id),
["2"],
);
});

102
tests/agentRuntimeSnapshot.test.ts
View File

@ -1,102 +0,0 @@
import { strict as assert } from "node:assert";
import test from "node:test";
import { buildAgentRuntimeSnapshot } from "../src/lib/agentRuntimeSnapshot.ts";
import { restoreStartupAgentRuntime } from "../src/lib/appStartup.ts";
test("builds runtime snapshot from active tab, selection, and limited result rows", () => {
const snapshot = buildAgentRuntimeSnapshot({
tabs: [
{
id: "tab-1",
title: "Orders",
connectionId: "conn-1",
database: "sales",
schema: "public",
sql: "select * from orders",
result: {
columns: ["id"],
rows: [[1], [2], [3]],
affected_rows: 0,
execution_time_ms: 12,
},
isExecuting: false,
mode: "query",
},
],
activeTabId: "tab-1",
getConnection: (id) => (id === "conn-1" ? { name: "Prod Sales" } : undefined),
selectedSql: "select *",
selection: {
type: "grid-cells",
data: { columns: ["id"], rows: [[1]] },
},
resultSampleLimit: 2,
});
assert.deepEqual(snapshot, {
activeConnectionId: "conn-1",
activeConnectionName: "Prod Sales",
database: "sales",
schema: "public",
activeTabId: "tab-1",
activeTabTitle: "Orders",
sql: "select * from orders",
selectedSql: "select *",
selection: {
type: "grid-cells",
data: { columns: ["id"], rows: [[1]] },
},
result: {
columns: ["id"],
rows: [[1], [2]],
truncated: true,
executionTimeMs: 12,
sampleLimit: 2,
},
});
});
test("omits empty selected SQL and marks empty UI selection as none", () => {
const snapshot = buildAgentRuntimeSnapshot({
tabs: [
{
id: "tab-2",
title: "Scratch",
connectionId: "conn-2",
database: "",
sql: "",
isExecuting: false,
mode: "query",
},
],
activeTabId: "tab-2",
getConnection: () => ({ name: "Local" }),
selectedSql: " ",
resultSampleLimit: 50,
});
assert.equal(snapshot.selectedSql, undefined);
assert.deepEqual(snapshot.selection, { type: "none" });
assert.equal(snapshot.result, undefined);
});
test("schedules startup runtime sync only after connections and restored tabs are ready", async () => {
const calls: string[] = [];
await restoreStartupAgentRuntime({
initSavedSql: async () => {
calls.push("saved-sql");
},
initConnections: async () => {
calls.push("connections");
},
reconnectRestoredTabs: async () => {
calls.push("reconnect");
},
scheduleSync: () => {
calls.push("schedule");
},
});
assert.deepEqual(calls, ["saved-sql", "connections", "reconnect", "schedule"]);
});