From be1a2c623693978c2bc639954c0ccc69a9e6b663 Mon Sep 17 00:00:00 2001 From: Guoyu Su Date: Fri, 17 Jul 2026 21:38:14 +0800 Subject: [PATCH] feat(duckdb): support editing view definitions --- .../assets/database-drivers.manifest.json | 2 +- .../dbx-core/src/db/duckdb_worker_process.rs | 18 ++- .../dbx-core/src/db/duckdb_worker_protocol.rs | 34 +++++ .../dbx-core/src/db/duckdb_worker_runtime.rs | 51 ++++++- crates/dbx-core/src/object_source_sql.rs | 44 ++++++ crates/dbx-core/src/schema.rs | 129 +++++++++++++++++- .../dbx-core/tests/duckdb_worker_process.rs | 39 ++++++ 7 files changed, 312 insertions(+), 5 deletions(-) diff --git a/crates/dbx-core/assets/database-drivers.manifest.json b/crates/dbx-core/assets/database-drivers.manifest.json index 50560770ac..d6308638a1 100644 --- a/crates/dbx-core/assets/database-drivers.manifest.json +++ b/crates/dbx-core/assets/database-drivers.manifest.json @@ -187,7 +187,7 @@ "queryExecution": true, "metadataBrowse": true, "objectBrowser": true, - "objectSource": false, + "objectSource": true, "schemaSearch": true, "diagram": false, "tableDataEdit": true, diff --git a/crates/dbx-core/src/db/duckdb_worker_process.rs b/crates/dbx-core/src/db/duckdb_worker_process.rs index 1ebfdb39b5..da01c7c7a1 100644 --- a/crates/dbx-core/src/db/duckdb_worker_process.rs +++ b/crates/dbx-core/src/db/duckdb_worker_process.rs @@ -14,8 +14,8 @@ use tokio_util::sync::CancellationToken; use crate::db; use crate::db::duckdb_worker_protocol::{ - DuckDbWorkerConnectParams, DuckDbWorkerError, DuckDbWorkerExecuteParams, DuckDbWorkerMethod, DuckDbWorkerRequest, - DuckDbWorkerResponse, + DuckDbWorkerConnectParams, DuckDbWorkerError, DuckDbWorkerExecuteParams, DuckDbWorkerMethod, + DuckDbWorkerObjectSourceParams, DuckDbWorkerRequest, DuckDbWorkerResponse, }; use crate::models::connection::AttachedDatabaseConfig; use crate::storage::{normalize_duckdb_worker_max_processes, DUCKDB_WORKER_MAX_PROCESSES_DEFAULT}; @@ -246,6 +246,20 @@ impl DuckDbWorkerClient { .await } + pub async fn get_object_source( + &self, + database: String, + schema: String, + name: String, + object_type: db::ObjectSourceKind, + ) -> Result { + self.metadata_request( + DuckDbWorkerMethod::GetObjectSource, + DuckDbWorkerObjectSourceParams { database, schema, name, object_type }, + ) + .await + } + pub async fn attach_database(&self, attached: AttachedDatabaseConfig) -> Result<(), String> { self.metadata_request::(DuckDbWorkerMethod::AttachDatabase, attached).await?; Ok(()) diff --git a/crates/dbx-core/src/db/duckdb_worker_protocol.rs b/crates/dbx-core/src/db/duckdb_worker_protocol.rs index c10ea0203a..0610f2c01e 100644 --- a/crates/dbx-core/src/db/duckdb_worker_protocol.rs +++ b/crates/dbx-core/src/db/duckdb_worker_protocol.rs @@ -3,6 +3,7 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize}; use crate::models::connection::AttachedDatabaseConfig; +use crate::types::ObjectSourceKind; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct DuckDbWorkerRequest { @@ -34,6 +35,7 @@ pub enum DuckDbWorkerMethod { ListSchemas, ListTables, ListColumns, + GetObjectSource, AttachDatabase, Cancel, Shutdown, @@ -135,6 +137,15 @@ pub struct DuckDbWorkerColumnParams { pub table: String, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct DuckDbWorkerObjectSourceParams { + pub database: String, + pub schema: String, + pub name: String, + pub object_type: ObjectSourceKind, +} + #[cfg(test)] mod tests { use super::*; @@ -178,4 +189,27 @@ mod tests { assert!(!parsed.ok); assert_eq!(parsed.error.expect("error").code, "duckdb_open_failed"); } + + #[test] + fn protocol_round_trips_object_source_request() { + let request = DuckDbWorkerRequest::new( + "req-3", + DuckDbWorkerMethod::GetObjectSource, + DuckDbWorkerObjectSourceParams { + database: "main".to_string(), + schema: "main".to_string(), + name: "active_orders".to_string(), + object_type: ObjectSourceKind::View, + }, + ) + .expect("serialize request"); + + let json = serde_json::to_string(&request).expect("request json"); + let parsed: DuckDbWorkerRequest = serde_json::from_str(&json).expect("parse request"); + let params: DuckDbWorkerObjectSourceParams = parsed.parse_params().expect("parse params"); + + assert_eq!(parsed.method, DuckDbWorkerMethod::GetObjectSource); + assert_eq!(params.name, "active_orders"); + assert_eq!(params.object_type, ObjectSourceKind::View); + } } diff --git a/crates/dbx-core/src/db/duckdb_worker_runtime.rs b/crates/dbx-core/src/db/duckdb_worker_runtime.rs index 595a3471ac..84bdbc3510 100644 --- a/crates/dbx-core/src/db/duckdb_worker_runtime.rs +++ b/crates/dbx-core/src/db/duckdb_worker_runtime.rs @@ -7,7 +7,8 @@ use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Stdout}; use crate::db; use crate::db::duckdb_worker_protocol::{ DuckDbWorkerColumnParams, DuckDbWorkerConnectParams, DuckDbWorkerDatabaseParams, DuckDbWorkerError, - DuckDbWorkerExecuteParams, DuckDbWorkerMethod, DuckDbWorkerRequest, DuckDbWorkerResponse, DuckDbWorkerTableParams, + DuckDbWorkerExecuteParams, DuckDbWorkerMethod, DuckDbWorkerObjectSourceParams, DuckDbWorkerRequest, + DuckDbWorkerResponse, DuckDbWorkerTableParams, }; use crate::models::connection::AttachedDatabaseConfig; use crate::path_utils::expand_tilde; @@ -103,6 +104,19 @@ impl DuckDbWorkerSession { ) } + pub fn get_object_source(&self, params: DuckDbWorkerObjectSourceParams) -> Result { + let connection = self.connection.as_ref().ok_or("DuckDB worker is not connected")?.clone(); + let locked = connection.lock().map_err(|e| e.to_string())?; + crate::schema::duckdb_object_source_with_attached( + &locked, + ¶ms.database, + ¶ms.schema, + ¶ms.name, + ¶ms.object_type, + &self.attached_names, + ) + } + pub fn attach_database(&mut self, attached: AttachedDatabaseConfig) -> Result<(), String> { let connection = self.connection.as_ref().ok_or("DuckDB worker is not connected")?.clone(); let locked = connection.lock().map_err(|e| e.to_string())?; @@ -182,6 +196,10 @@ impl DuckDbWorkerRuntime { let params = request.parse_params()?; Ok(DuckDbWorkerResponse::ok(request.id, session.list_columns(params)?)) }), + DuckDbWorkerMethod::GetObjectSource => self.handle_session_request(request, |session, request| { + let params = request.parse_params()?; + Ok(DuckDbWorkerResponse::ok(request.id, session.get_object_source(params)?)) + }), DuckDbWorkerMethod::AttachDatabase => self.handle_session_request(request, |session, request| { session.attach_database(request.parse_params()?)?; Ok(DuckDbWorkerResponse::ok_empty(request.id)) @@ -428,6 +446,37 @@ mod tests { assert!(columns.iter().any(|column| column.name == "id" && column.is_primary_key)); } + #[test] + fn worker_session_reads_view_source() { + let mut session = DuckDbWorkerSession::default(); + session + .connect(DuckDbWorkerConnectParams { + path: ":memory:".to_string(), + attached_databases: Vec::new(), + init_script: None, + }) + .expect("connect"); + session + .execute(DuckDbWorkerExecuteParams { + sql: "CREATE VIEW active_orders AS SELECT 1 AS id".to_string(), + database: None, + max_rows: None, + }) + .expect("create view"); + + let source = session + .get_object_source(DuckDbWorkerObjectSourceParams { + database: "main".to_string(), + schema: "main".to_string(), + name: "active_orders".to_string(), + object_type: db::ObjectSourceKind::View, + }) + .expect("get view source"); + + assert!(source.starts_with("CREATE VIEW")); + assert!(source.contains("active_orders")); + } + #[test] fn worker_session_attaches_database() { let dir = std::env::temp_dir().join(format!("dbx-duckdb-worker-{}", uuid::Uuid::new_v4())); diff --git a/crates/dbx-core/src/object_source_sql.rs b/crates/dbx-core/src/object_source_sql.rs index 21f7fdd522..b4afe6dac5 100644 --- a/crates/dbx-core/src/object_source_sql.rs +++ b/crates/dbx-core/src/object_source_sql.rs @@ -129,6 +129,10 @@ pub fn build_executable_object_source_statements(input: EditableObjectSourceSqlI return Ok(vec![replace_sqlserver_create_with_alter(source)]); } + if input.database_type == DatabaseType::DuckDb && input.object_type == ObjectSourceKind::View { + return Ok(vec![executable_duckdb_view_ddl(input.schema.as_deref(), &input.name, source)]); + } + if matches!( input.database_type, DatabaseType::Postgres @@ -813,6 +817,16 @@ fn build_sqlserver_alter_view_sql(schema: Option<&str>, name: &str, source: &str format!("ALTER VIEW {} AS\n{}", sqlserver_qualified_name(schema, name), ensure_semicolon(source)) } +fn executable_duckdb_view_ddl(schema: Option<&str>, name: &str, source: &str) -> String { + let existing_view_statement = + Regex::new(r"(?i)^CREATE\s+(?:OR\s+REPLACE\s+)?VIEW(?:\s+IF\s+NOT\s+EXISTS)?\s+").unwrap(); + if existing_view_statement.is_match(source) { + return ensure_semicolon(&existing_view_statement.replace(source, "CREATE OR REPLACE VIEW ")); + } + + format!("CREATE OR REPLACE VIEW {} AS\n{}", postgres_qualified_name(schema, name), ensure_semicolon(source)) +} + fn parse_object_source_kind(value: &str) -> Option { if value.eq_ignore_ascii_case("VIEW") { Some(ObjectSourceKind::View) @@ -977,6 +991,36 @@ mod tests { assert_eq!(sql, "ALTER VIEW [dbo].[new_view] AS\nSELECT\n *\nFROM AppInfo;"); } + #[test] + fn duckdb_view_create_source_saves_as_create_or_replace_view() { + let sql = build_executable_object_source_sql(EditableObjectSourceSqlInput { + database_type: DatabaseType::DuckDb, + object_type: ObjectSourceKind::View, + schema: Some("main".to_string()), + name: "active_orders".to_string(), + source: "CREATE VIEW active_orders AS SELECT id FROM orders WHERE active".to_string(), + }) + .unwrap(); + + assert_eq!(sql, "CREATE OR REPLACE VIEW active_orders AS SELECT id FROM orders WHERE active;"); + } + + #[test] + fn duckdb_view_body_saves_as_qualified_create_or_replace_view() { + let sql = build_editable_object_source(EditableObjectSourceSqlInput { + database_type: DatabaseType::DuckDb, + object_type: ObjectSourceKind::View, + schema: Some("reporting".to_string()), + name: "active orders".to_string(), + source: "SELECT id FROM orders WHERE active".to_string(), + }); + + assert_eq!( + sql, + "CREATE OR REPLACE VIEW \"reporting\".\"active orders\" AS\nSELECT id FROM orders WHERE active;" + ); + } + #[test] fn postgres_view_body_opens_as_create_or_replace_view() { let sql = build_executable_object_source_sql(EditableObjectSourceSqlInput { diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 2f2804f296..7a79e3bf8a 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -142,6 +142,32 @@ pub fn duckdb_list_schemas_with_attached( Ok(schemas) } +#[cfg(feature = "duckdb-bundled")] +pub fn duckdb_object_source_with_attached( + con: &duckdb::Connection, + database: &str, + schema: &str, + name: &str, + object_type: &db::ObjectSourceKind, + attached_names: &[String], +) -> Result { + if !matches!(object_type, db::ObjectSourceKind::View) { + return Err("DuckDB object source only supports views".to_string()); + } + + let database = duckdb_catalog_name(con, database, attached_names)?; + let mut stmt = con + .prepare("SELECT sql FROM duckdb_views() WHERE database_name = ? AND schema_name = ? AND view_name = ?") + .map_err(|e| e.to_string())?; + let mut rows = + stmt.query_map((database.as_str(), schema, name), |row| row.get::<_, String>(0)).map_err(|e| e.to_string())?; + + match rows.next() { + Some(row) => row.map_err(|e| e.to_string()), + None => Err(format!("DuckDB view source not found: {database}.{schema}.{name}")), + } +} + /// Identifies quack catalogs by the storage-extension `type` that /// `duckdb_databases()` reports, rather than by the attach aliases parsed from /// the init script: `ATTACH 'quack:host:port'` without an `AS` alias gets a @@ -2430,7 +2456,7 @@ mod tests { #[cfg(feature = "duckdb-bundled")] use super::{ duckdb_attach_database, duckdb_completion_assistant_search, duckdb_list_databases, - duckdb_query_tables_in_database, + duckdb_object_source_with_attached, duckdb_query_tables_in_database, }; use crate::models::connection::{ConnectionConfig, DatabaseType}; use std::collections::HashMap; @@ -2955,6 +2981,83 @@ mod tests { assert_eq!(name.comment.as_deref(), Some("Display name")); } + #[cfg(feature = "duckdb-bundled")] + #[test] + fn duckdb_object_source_filters_by_schema_and_view_name() { + let con = duckdb::Connection::open_in_memory().unwrap(); + con.execute_batch( + "CREATE SCHEMA reporting; \ + CREATE VIEW main.active_orders AS SELECT 'main' AS origin; \ + CREATE VIEW reporting.active_orders AS SELECT 'reporting' AS origin;", + ) + .unwrap(); + + let main_source = + duckdb_object_source_with_attached(&con, "main", "main", "active_orders", &db::ObjectSourceKind::View, &[]) + .unwrap(); + let reporting_source = duckdb_object_source_with_attached( + &con, + "main", + "reporting", + "active_orders", + &db::ObjectSourceKind::View, + &[], + ) + .unwrap(); + + assert!(main_source.contains("'main'")); + assert!(reporting_source.contains("'reporting'")); + } + + #[cfg(feature = "duckdb-bundled")] + #[test] + fn duckdb_object_source_rejects_unsupported_or_missing_objects() { + let con = duckdb::Connection::open_in_memory().unwrap(); + + let unsupported = duckdb_object_source_with_attached( + &con, + "main", + "main", + "calculate_total", + &db::ObjectSourceKind::Function, + &[], + ) + .unwrap_err(); + let missing = + duckdb_object_source_with_attached(&con, "main", "main", "missing_view", &db::ObjectSourceKind::View, &[]) + .unwrap_err(); + + assert_eq!(unsupported, "DuckDB object source only supports views"); + assert!(missing.contains("DuckDB view source not found")); + } + + #[cfg(feature = "duckdb-bundled")] + #[test] + fn duckdb_catalog_view_source_can_replace_existing_view() { + let con = duckdb::Connection::open_in_memory().unwrap(); + con.execute_batch("CREATE VIEW active_orders AS SELECT 1 AS value;").unwrap(); + let source = + duckdb_object_source_with_attached(&con, "main", "main", "active_orders", &db::ObjectSourceKind::View, &[]) + .unwrap(); + let edited_source = source.replace("SELECT 1", "SELECT 2"); + assert_ne!(edited_source, source); + + let statements = crate::object_source_sql::build_executable_object_source_statements( + crate::object_source_sql::EditableObjectSourceSqlInput { + database_type: DatabaseType::DuckDb, + object_type: db::ObjectSourceKind::View, + schema: Some("main".to_string()), + name: "active_orders".to_string(), + source: edited_source, + }, + ) + .unwrap(); + con.execute_batch(&statements.join("\n")).unwrap(); + + let value = con.query_row("SELECT value FROM active_orders", [], |row| row.get::<_, i32>(0)).unwrap(); + assert_eq!(value, 2); + } + #[cfg(feature = "duckdb-bundled")] #[test] fn duckdb_completion_assistant_searches_catalog_metadata_with_limit() { @@ -5194,6 +5297,8 @@ async fn get_object_source_once( ) -> Result { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; + #[cfg(feature = "duckdb-bundled")] + let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await; let source = { let connections = state.connections.read().await; if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { @@ -5262,6 +5367,28 @@ async fn get_object_source_once( PoolKind::Sqlite(pool) => first_string_cell( db::sqlite::execute_query(pool, &sqlite_object_source_sql(name, &object_type)).await?, )?, + #[cfg(feature = "duckdb-bundled")] + PoolKind::DuckDb(con) => { + let con = con.lock().map_err(|e| e.to_string())?; + duckdb_object_source_with_attached( + &con, + database, + schema, + name, + &object_type, + &duckdb_attached_names, + )? + } + #[cfg(feature = "duckdb-bundled")] + PoolKind::DuckDbWorker(client) => { + let client = client.clone(); + let database = database.to_string(); + let schema = schema.to_string(); + let name = name.to_string(); + let object_type = object_type.clone(); + drop(connections); + client.get_object_source(database, schema, name, object_type).await? + } PoolKind::Rqlite(client) => { return db::rqlite_driver::object_source(client, name, &object_type).await; } diff --git a/crates/dbx-core/tests/duckdb_worker_process.rs b/crates/dbx-core/tests/duckdb_worker_process.rs index 6dc017fb9c..415418a3e7 100644 --- a/crates/dbx-core/tests/duckdb_worker_process.rs +++ b/crates/dbx-core/tests/duckdb_worker_process.rs @@ -18,6 +18,45 @@ use tokio_util::sync::CancellationToken; static TEMP_DB_COUNTER: AtomicU64 = AtomicU64::new(0); +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_process_reads_view_source() { + let _guard = duckdb_worker_process_test_guard().await; + let executable = PathBuf::from(env!("CARGO_BIN_EXE_duckdb-worker-test-host")); + let db_path = temp_duckdb_path(); + let _ = std::fs::remove_file(&db_path); + + let client = + DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new(), None) + .await + .expect("worker process connects"); + client + .execute( + None, + "CREATE VIEW active_orders AS SELECT 1 AS id".to_string(), + None, + None, + Some(Duration::from_secs(5)), + ) + .await + .expect("create view"); + + let source = client + .get_object_source( + "main".to_string(), + "main".to_string(), + "active_orders".to_string(), + dbx_core::db::ObjectSourceKind::View, + ) + .await + .expect("get view source"); + + assert!(source.starts_with("CREATE VIEW")); + assert!(source.contains("active_orders")); + + client.shutdown().await; + let _ = std::fs::remove_file(&db_path); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn worker_process_recovers_immediately_after_cancelled_long_query() { let _guard = duckdb_worker_process_test_guard().await;