github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 14:52:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(duckdb): add connection bootstrap and remote catalogs

Browse Source
This commit is contained in:
Komh authored and GitHub committed 2026-07-13 17:03:33 +08:00
1 parent fc5e991a75
commit 35ca4e3a37
41 files changed
+893 -194
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/agent_connection.rs
+1
View File
@@ -590,6 +590,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: default_connect_timeout_secs(),
crates/dbx-core/src/cloud_sync.rs
+10
View File
@@ -30,6 +30,7 @@ const SECRET_KEYS: &[&str] = &[
"proxy_password",
"redis_sentinel_password",
"connection_string",
"init_script",
MQ_AUTH_TOKEN_KEY,
MQ_AUTH_PASSWORD_KEY,
MQ_AUTH_API_KEY_VALUE_KEY,
@@ -597,6 +598,7 @@ fn scrub_connection_secrets(config: &mut ConnectionConfig) {
}
config.redis_sentinel_password.clear();
config.connection_string = None;
config.init_script = None;
scrub_mq_external_config_secrets(config);
scrub_nacos_auth_secrets(config);
}
@@ -617,6 +619,7 @@ async fn build_sensitive_payload(
let mut connection_secrets = Vec::new();
for config in connections {
push_secret(&mut connection_secrets, &config.id, "password", &config.password);
push_secret(&mut connection_secrets, &config.id, "init_script", config.init_script.as_deref().unwrap_or(""));
for (index, layer) in config.transport_layers.iter().enumerate() {
match layer {
TransportLayerConfig::Ssh(ssh) => {
@@ -1063,6 +1066,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
@@ -1115,6 +1119,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
@@ -1204,6 +1209,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: Some("CREATE SECRET (TYPE quack, TOKEN 'token-value');".to_string()),
color: None,
transport_layers: vec![
TransportLayerConfig::Ssh(crate::models::connection::SshTunnelConfig {
@@ -1279,6 +1285,10 @@ mod tests {
}
assert!(config.redis_sentinel_password.is_empty());
assert!(config.connection_string.is_none());
assert!(config.init_script.is_none());
let public_json = serde_json::to_string(&config).unwrap();
assert!(!public_json.contains("token-value"));
assert!(super::SECRET_KEYS.contains(&"init_script"));
}
#[test]
crates/dbx-core/src/connection.rs
+90 -69
View File
@@ -610,6 +610,49 @@ impl AppState {
}
}
#[cfg(feature = "duckdb-bundled")]
async fn create_duckdb_pool(&self, config: &ConnectionConfig) -> Result<PoolKind, String> {
if self.duckdb_worker_process_isolation.load(Ordering::Relaxed) {
let attached_databases = config
.attached_databases
.iter()
.map(|attached| crate::models::connection::AttachedDatabaseConfig {
name: attached.name.clone(),
path: expand_tilde(&attached.path),
})
.collect();
let client = db::duckdb_worker_process::DuckDbWorkerClient::open_with_process_limit(
expand_tilde(&config.host),
attached_databases,
config.init_script.clone(),
self.duckdb_worker_max_processes.load(Ordering::Relaxed),
)
.await?;
Ok(PoolKind::DuckDbWorker(Arc::new(client)))
} else {
let con = db::duckdb_driver::connect_path(&expand_tilde(&config.host))?;
{
let locked = con.lock().map_err(|e| e.to_string())?;
for attached in &config.attached_databases {
crate::schema::duckdb_attach_database(&locked, &attached.name, &expand_tilde(&attached.path))?;
}
if let Some(script) = config.init_script.as_deref() {
db::duckdb_driver::run_init_script(&locked, script)?;
}
}
Ok(PoolKind::DuckDb(con))
}
}
#[cfg(feature = "duckdb-bundled")]
pub async fn test_duckdb_connection_config(&self, config: &ConnectionConfig) -> Result<(), String> {
// Test the submitted form as a fresh session so unsaved ATTACH/init
// changes cannot be masked by a pool created from older settings.
let pool = self.create_duckdb_pool(config).await?;
close_pool_kind(pool).await;
Ok(())
}
pub async fn test_external_driver(&self, driver_id: &str, config: &ConnectionConfig) -> Result<String, String> {
let params = serde_json::json!({ "connection": config });
let env = self.external_driver_runtime_env(driver_id)?;
@@ -1108,38 +1151,7 @@ impl AppState {
PoolKind::Redis(con)
}
#[cfg(feature = "duckdb-bundled")]
DatabaseType::DuckDb => {
if self.duckdb_worker_process_isolation.load(Ordering::Relaxed) {
let attached_databases = db_config
.attached_databases
.iter()
.map(|attached| crate::models::connection::AttachedDatabaseConfig {
name: attached.name.clone(),
path: expand_tilde(&attached.path),
})
.collect();
let client = db::duckdb_worker_process::DuckDbWorkerClient::open_with_process_limit(
expand_tilde(&db_config.host),
attached_databases,
self.duckdb_worker_max_processes.load(Ordering::Relaxed),
)
.await?;
PoolKind::DuckDbWorker(Arc::new(client))
} else {
let con = db::duckdb_driver::connect_path(&expand_tilde(&db_config.host))?;
{
let locked = con.lock().map_err(|e| e.to_string())?;
for attached in &db_config.attached_databases {
crate::schema::duckdb_attach_database(
&locked,
&attached.name,
&expand_tilde(&attached.path),
)?;
}
}
PoolKind::DuckDb(con)
}
}
DatabaseType::DuckDb => self.create_duckdb_pool(&db_config).await?,
#[cfg(not(feature = "duckdb-bundled"))]
DatabaseType::DuckDb => {
return Err("DuckDB support is not compiled in this build. Rebuild with default features.".to_string());
@@ -2259,39 +2271,6 @@ impl AppState {
keys
}
#[cfg(feature = "duckdb-bundled")]
pub async fn duckdb_existing_pool_is_usable_for_config(&self, config: &ConnectionConfig) -> Result<bool, String> {
if config.db_type != DatabaseType::DuckDb {
return Ok(false);
}
let matches_existing_config = {
let configs = self.configs.read().await;
configs.get(&config.id).is_some_and(|existing| {
existing.db_type == DatabaseType::DuckDb && duckdb_paths_match(&existing.host, &config.host)
})
};
if !matches_existing_config {
return Ok(false);
}
let duckdb_pool = {
let conns = self.connections.read().await;
match conns.get(&config.id) {
Some(PoolKind::DuckDb(con)) => Some(con.clone()),
_ => None,
}
};
let Some(con) = duckdb_pool else {
return Ok(false);
};
let locked = con.lock().map_err(|e| e.to_string())?;
locked.execute_batch("SELECT 1;").map_err(|e| format!("DuckDb connection failed: {e}"))?;
Ok(true)
}
pub async fn reset_connection_transport(&self, connection_id: &str) {
let layer_count = {
let configs = self.configs.read().await;
@@ -3270,8 +3249,8 @@ mod tests {
use crate::database_capabilities;
use crate::db;
use crate::models::connection::{
default_connect_timeout_secs, default_redis_key_separator, ConnectionConfig, DatabaseType, HttpTunnelConfig,
ProxyTunnelConfig, ProxyType, SshTunnelConfig, TransportLayerConfig,
default_connect_timeout_secs, default_redis_key_separator, AttachedDatabaseConfig, ConnectionConfig,
DatabaseType, HttpTunnelConfig, ProxyTunnelConfig, ProxyType, SshTunnelConfig, TransportLayerConfig,
};
use crate::query;
use crate::schema;
@@ -3295,6 +3274,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: default_connect_timeout_secs(),
@@ -4270,7 +4250,7 @@ mod tests {
#[cfg(feature = "duckdb-bundled")]
#[tokio::test]
async fn duckdb_existing_pool_can_be_used_for_connection_test() {
async fn duckdb_connection_test_does_not_reuse_stale_pool_config() {
let (state, dir) = test_app_state().await;
let db_path = dir.join("app.duckdb");
duckdb::Connection::open(&db_path).unwrap();
@@ -4284,7 +4264,48 @@ mod tests {
state.configs.write().await.insert(config.id.clone(), config.clone());
state.get_or_create_pool("duckdb-conn", None).await.unwrap();
assert!(state.duckdb_existing_pool_is_usable_for_config(&config).await.unwrap());
config.init_script = Some("SELECT definitely_invalid_syntax(".to_string());
let error = state.test_duckdb_connection_config(&config).await.expect_err("invalid submitted script must fail");
assert!(error.contains("Connection init script statement 1 failed"));
let _ = std::fs::remove_dir_all(dir);
}
#[cfg(feature = "duckdb-bundled")]
#[tokio::test]
async fn duckdb_connection_test_validates_attached_databases() {
let (state, dir) = test_app_state().await;
let mut config = mysql_config(None);
config.id = "duckdb-conn".to_string();
config.name = "DuckDB".to_string();
config.db_type = DatabaseType::DuckDb;
config.host = ":memory:".to_string();
config.port = 0;
config.attached_databases.push(AttachedDatabaseConfig {
name: "missing".to_string(),
path: dir.join("missing-parent").join("missing.duckdb").to_string_lossy().to_string(),
});
let error = state.test_duckdb_connection_config(&config).await.expect_err("invalid attach must fail");
assert!(error.to_ascii_lowercase().contains("attach"));
let _ = std::fs::remove_dir_all(dir);
}
#[cfg(feature = "duckdb-bundled")]
#[tokio::test]
async fn duckdb_connection_test_supports_memory_bootstrap() {
let (state, dir) = test_app_state().await;
let mut config = mysql_config(None);
config.id = "duckdb-memory".to_string();
config.name = "DuckDB memory".to_string();
config.db_type = DatabaseType::DuckDb;
config.host = ":memory:".to_string();
config.port = 0;
config.init_script = Some(r#"CREATE TABLE probe AS SELECT E'it\'s;ok' AS value;"#.to_string());
state.test_duckdb_connection_config(&config).await.expect("memory bootstrap succeeds");
let _ = std::fs::remove_dir_all(dir);
}
crates/dbx-core/src/connection_secrets.rs
+41 -2
View File
@@ -11,6 +11,7 @@ pub const TRANSPORT_LAYER_SECRET_PREFIX: &str = "transport_layers.";
pub const PROXY_PASSWORD_KEY: &str = "proxy_password";
pub const REDIS_SENTINEL_PASSWORD_KEY: &str = "redis_sentinel_password";
pub const CONNECTION_STRING_KEY: &str = "connection_string";
pub const INIT_SCRIPT_KEY: &str = "init_script";
pub const MQ_AUTH_SECRET_PREFIX: &str = "mq.auth.";
pub const MQ_AUTH_TOKEN_KEY: &str = "mq.auth.token";
pub const MQ_AUTH_PASSWORD_KEY: &str = "mq.auth.password";
@@ -104,6 +105,7 @@ pub fn save_connections_to_file(
}
persist_secret(store, &config.id, REDIS_SENTINEL_PASSWORD_KEY, &config.redis_sentinel_password)?;
persist_optional_secret(store, &config.id, CONNECTION_STRING_KEY, config.connection_string.as_deref())?;
persist_optional_secret(store, &config.id, INIT_SCRIPT_KEY, config.init_script.as_deref())?;
persist_mq_auth_secrets(store, config)?;
persist_mq_token_signing_secret(store, config)?;
@@ -160,6 +162,18 @@ pub fn load_connections_from_file(
}
}
}
match config.init_script.as_deref().filter(|secret| !secret.is_empty()) {
Some(secret) => {
store.set_secret(&config.id, INIT_SCRIPT_KEY, secret)?;
needs_rewrite = true;
}
None => {
if let Some(secret) = store.get_secret(&config.id, INIT_SCRIPT_KEY)? {
config.init_script = Some(secret);
}
}
}
hydrate_mq_auth_secrets(store, config, &mut needs_rewrite)?;
hydrate_mq_token_signing_secret(store, config, &mut needs_rewrite)?;
}
@@ -337,6 +351,7 @@ fn delete_removed_connection_secrets(
delete_secret_prefix(store, &config.id, SSH_TUNNEL_SECRET_PREFIX)?;
delete_secret_prefix(store, &config.id, TRANSPORT_LAYER_SECRET_PREFIX)?;
store.delete_secret(&config.id, CONNECTION_STRING_KEY)?;
store.delete_secret(&config.id, INIT_SCRIPT_KEY)?;
delete_secret_prefix(store, &config.id, MQ_AUTH_SECRET_PREFIX)?;
delete_secret_prefix(store, &config.id, MQ_TOKEN_SIGNING_SECRET_PREFIX)?;
}
@@ -637,6 +652,7 @@ fn sanitize_connections(configs: &[ConnectionConfig]) -> Vec<ConnectionConfig> {
}
config.redis_sentinel_password.clear();
config.connection_string = None;
config.init_script = None;
scrub_mq_auth_secrets(&mut config);
scrub_mq_token_signing_secret(&mut config);
config
@@ -652,8 +668,8 @@ pub fn secret_account(connection_id: &str, key: &str) -> String {
mod tests {
use super::{
load_connections_from_file, save_connections_to_file, ConnectionSecretStore, CONNECTION_STRING_KEY,
MAIN_PASSWORD_KEY, MQ_AUTH_PASSWORD_KEY, MQ_AUTH_TOKEN_KEY, MQ_TOKEN_SIGNING_KEY, REDIS_SENTINEL_PASSWORD_KEY,
SSH_PASSWORD_KEY,
INIT_SCRIPT_KEY, MAIN_PASSWORD_KEY, MQ_AUTH_PASSWORD_KEY, MQ_AUTH_TOKEN_KEY, MQ_TOKEN_SIGNING_KEY,
REDIS_SENTINEL_PASSWORD_KEY, SSH_PASSWORD_KEY,
};
use crate::models::connection::{
ConnectionConfig, DatabaseType, HttpTunnelConfig, SshTunnelConfig, TransportLayerConfig,
@@ -733,6 +749,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: crate::models::connection::default_connect_timeout_secs(),
@@ -961,6 +978,28 @@ mod tests {
assert_eq!(loaded[0].connection_string.as_deref(), Some("mongodb://user:secret@localhost/app"));
}
#[test]
fn save_connections_moves_init_script_to_secret_store_and_restores_it() {
let path = temp_connections_file("init-script");
let store = MemorySecretStore::default();
let mut config = connection("duck", "", "");
config.db_type = DatabaseType::DuckDb;
config.init_script = Some("CREATE SECRET (TYPE quack, TOKEN 'token-value');".to_string());
save_connections_to_file(&path, &[config], &store).unwrap();
assert_eq!(
store.get_existing("duck", INIT_SCRIPT_KEY).as_deref(),
Some("CREATE SECRET (TYPE quack, TOKEN 'token-value');")
);
let persisted = read_configs(&path);
assert_eq!(persisted[0].init_script, None);
assert!(!std::fs::read_to_string(&path).unwrap().contains("token-value"));
let loaded = load_connections_from_file(&path, &store).unwrap();
assert_eq!(loaded[0].init_script.as_deref(), Some("CREATE SECRET (TYPE quack, TOKEN 'token-value');"));
}
#[test]
fn save_connections_moves_mq_auth_secrets_to_secret_store_and_restores_them() {
let path = temp_connections_file("mq-auth");
crates/dbx-core/src/db/duckdb_driver.rs
+25
View File
@@ -64,6 +64,19 @@ pub fn connect_path(path: &str) -> Result<Arc<DuckDbConnection>, String> {
Ok(Arc::new(DuckDbConnection::new(connection)))
}
/// Runs a connection init script statement-by-statement so a failure can
/// point at the offending statement instead of the whole batch.
pub fn run_init_script(con: &duckdb::Connection, script: &str) -> Result<(), String> {
for (index, statement) in crate::db::duckdb_sql::split_sql_statements(script).iter().enumerate() {
if crate::db::duckdb_sql::strip_leading_comments(statement).is_empty() {
continue;
}
con.execute_batch(statement)
.map_err(|e| format!("Connection init script statement {} failed: {e}", index + 1))?;
}
Ok(())
}
fn validate_duckdb_path(path: &str) -> Result<(), String> {
if path.is_empty() {
return Err("Database file path cannot be empty".to_string());
@@ -168,4 +181,16 @@ mod tests {
assert!(std::fs::metadata(&path).expect("database metadata").len() > 0);
let _ = std::fs::remove_file(path);
}
#[test]
fn run_init_script_executes_duckdb_escape_strings() {
let con = connect_path(":memory:").expect("connect in-memory DuckDB");
let locked = con.lock().expect("lock connection");
run_init_script(&locked, r#"CREATE TABLE probe AS SELECT E'it\'s;ok' AS value; SELECT 2;"#)
.expect("run init script");
let value: String = locked.query_row("SELECT value FROM probe", [], |row| row.get(0)).expect("select value");
assert_eq!(value, "it's;ok");
}
}
crates/dbx-core/src/db/duckdb_sql.rs
+291
View File
@@ -0,0 +1,291 @@
//! Quote-aware helpers for inspecting DuckDB SQL snippets without a full
//! parser. Shared by the DuckDB worker (execute-time ATTACH tracking) and the
//! connection layer (init-script handling), and kept free of `duckdb-bundled`
//! so config-derived metadata works in every build.
use crate::models::connection::ConnectionConfig;
/// Splits a SQL script into statements on `;`, ignoring separators inside
/// single/double quotes, `--` line comments, and `/* */` block comments.
/// Statements are trimmed; empty fragments are dropped.
pub fn split_sql_statements(script: &str) -> Vec<String> {
let bytes = script.as_bytes();
let mut statements = Vec::new();
let mut start = 0;
let mut i = 0;
let mut in_single = false;
let mut in_double = false;
let mut in_escape_single = false;
while i < bytes.len() {
match bytes[i] {
b'\\' if in_single && in_escape_single => {
// DuckDB E/e strings use backslash escapes, so an escaped
// quote or semicolon must not change statement boundaries.
i = (i + 2).min(bytes.len());
continue;
}
b'\'' if !in_double => {
if in_single && i + 1 < bytes.len() && bytes[i + 1] == b'\'' {
i += 2;
continue;
}
if !in_single {
in_escape_single = is_escape_string_quote(bytes, i);
}
in_single = !in_single;
if !in_single {
in_escape_single = false;
}
}
b'"' if !in_single => {
if in_double && i + 1 < bytes.len() && bytes[i + 1] == b'"' {
i += 2;
continue;
}
in_double = !in_double;
}
b'-' if !in_single && !in_double && bytes.get(i + 1) == Some(&b'-') => {
while i < bytes.len() && bytes[i] != b'\n' {
i += 1;
}
continue;
}
b'/' if !in_single && !in_double && bytes.get(i + 1) == Some(&b'*') => {
i += 2;
while i < bytes.len() && !(bytes[i] == b'*' && bytes.get(i + 1) == Some(&b'/')) {
i += 1;
}
i = (i + 2).min(bytes.len());
continue;
}
b'$' if !in_single && !in_double => {
if let Some(end) = dollar_quote_end(bytes, i) {
i = end;
continue;
}
}
b';' if !in_single && !in_double => {
push_statement(&mut statements, &script[start..i]);
start = i + 1;
}
_ => {}
}
i += 1;
}
push_statement(&mut statements, &script[start..]);
statements
}
fn is_escape_string_quote(bytes: &[u8], quote_index: usize) -> bool {
if quote_index == 0 || !matches!(bytes[quote_index - 1], b'E' | b'e') {
return false;
}
quote_index == 1 || !(bytes[quote_index - 2].is_ascii_alphanumeric() || bytes[quote_index - 2] == b'_')
}
fn push_statement(statements: &mut Vec<String>, fragment: &str) {
let trimmed = fragment.trim();
if !trimmed.is_empty() {
statements.push(trimmed.to_string());
}
}
/// Returns the index right after the closing `$tag$` when `start` opens a
/// dollar-quoted string (Postgres rules: empty tag or an identifier that does
/// not start with a digit, so `$1` placeholders never match). An unterminated
/// dollar quote swallows the rest of the script, like an unterminated quote.
fn dollar_quote_end(bytes: &[u8], start: usize) -> Option<usize> {
let mut j = start + 1;
while j < bytes.len() && (bytes[j] == b'_' || bytes[j].is_ascii_alphanumeric()) {
j += 1;
}
if j >= bytes.len() || bytes[j] != b'$' {
return None;
}
if j > start + 1 && bytes[start + 1].is_ascii_digit() {
return None;
}
let tag = &bytes[start..=j];
let mut k = j + 1;
while k + tag.len() <= bytes.len() {
if &bytes[k..k + tag.len()] == tag {
return Some(k + tag.len());
}
k += 1;
}
Some(bytes.len())
}
/// Returns the alias of an `ATTACH ... AS <name>` statement, or `None` when
/// the statement is not an ATTACH or carries no alias. Leading comments are
/// skipped so `-- note\nATTACH ...` is still recognized.
pub fn attached_name_from_attach_sql(sql: &str) -> Option<String> {
let trimmed = strip_leading_comments(sql);
let first_word = trimmed.split(|ch: char| ch.is_whitespace() || ch == ';').next().unwrap_or_default();
if !first_word.eq_ignore_ascii_case("ATTACH") {
return None;
}
let as_index = find_as_keyword_outside_quotes(trimmed)?;
parse_identifier_after_as(&trimmed[as_index + 2..])
}
pub(crate) fn strip_leading_comments(sql: &str) -> &str {
let mut rest = sql.trim_start();
loop {
if let Some(after) = rest.strip_prefix("--") {
rest = after.split_once('\n').map(|(_, tail)| tail).unwrap_or("").trim_start();
} else if let Some(after) = rest.strip_prefix("/*") {
rest = after.split_once("*/").map(|(_, tail)| tail).unwrap_or("").trim_start();
} else {
return rest;
}
}
}
/// Collects the aliases of every `ATTACH ... AS <name>` statement in a script.
pub fn init_script_attached_names(script: &str) -> Vec<String> {
split_sql_statements(script).iter().filter_map(|statement| attached_name_from_attach_sql(statement)).collect()
}
/// Attached-database names for a connection: the configured attach list plus
/// aliases attached by the connection init script (case-insensitive dedup).
pub fn config_attached_names(config: &ConnectionConfig) -> Vec<String> {
let mut names: Vec<String> = config.attached_databases.iter().map(|database| database.name.clone()).collect();
if let Some(script) = config.init_script.as_deref() {
for name in init_script_attached_names(script) {
if !names.iter().any(|existing| existing.eq_ignore_ascii_case(&name)) {
names.push(name);
}
}
}
names
}
fn find_as_keyword_outside_quotes(sql: &str) -> Option<usize> {
let bytes = sql.as_bytes();
let mut i = 0;
let mut in_single = false;
let mut in_double = false;
while i < bytes.len() {
match bytes[i] {
b'\'' if !in_double => {
if in_single && i + 1 < bytes.len() && bytes[i + 1] == b'\'' {
i += 2;
continue;
}
in_single = !in_single;
}
b'"' if !in_single => {
if in_double && i + 1 < bytes.len() && bytes[i + 1] == b'"' {
i += 2;
continue;
}
in_double = !in_double;
}
b'a' | b'A' if !in_single && !in_double && i + 1 < bytes.len() => {
if (bytes[i + 1] == b's' || bytes[i + 1] == b'S')
&& is_sql_word_boundary(bytes.get(i.wrapping_sub(1)).copied())
&& is_sql_word_boundary(bytes.get(i + 2).copied())
{
return Some(i);
}
}
_ => {}
}
i += 1;
}
None
}
fn is_sql_word_boundary(byte: Option<u8>) -> bool {
!matches!(byte, Some(b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'_'))
}
fn parse_identifier_after_as(input: &str) -> Option<String> {
let input = input.trim_start();
if input.is_empty() {
return None;
}
if let Some(rest) = input.strip_prefix('"') {
let mut name = String::new();
let mut chars = rest.chars().peekable();
while let Some(ch) = chars.next() {
if ch == '"' {
if chars.peek() == Some(&'"') {
name.push('"');
chars.next();
continue;
}
return (!name.trim().is_empty()).then_some(name);
}
name.push(ch);
}
return None;
}
let name = input.split(|ch: char| ch.is_whitespace() || ch == ';').next().unwrap_or_default().trim();
(!name.is_empty()).then(|| name.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn attach_sql_alias_parser_handles_generated_sql() {
assert_eq!(
attached_name_from_attach_sql("ATTACH 'D:\\tmp\\sales.duckdb' AS \"sales db\";"),
Some("sales db".to_string())
);
assert_eq!(attached_name_from_attach_sql("select 'not attach' as value"), None);
assert_eq!(
attached_name_from_attach_sql("-- note; with semicolon\n/* block */ ATTACH 'x.duckdb' AS warehouse"),
Some("warehouse".to_string())
);
}
#[test]
fn split_ignores_separators_in_quotes_and_comments() {
// Comment text stays inside the statement (DuckDB parses it); the
// splitter only refuses to break on `;` inside quotes or comments.
let script = "SET s3_endpoint='host;port'; -- comment; with semicolon\nATTACH 'db;x.duckdb' AS a; /* multi;\nline */ SELECT 1";
let statements = split_sql_statements(script);
assert_eq!(
statements,
vec![
"SET s3_endpoint='host;port'".to_string(),
"-- comment; with semicolon\nATTACH 'db;x.duckdb' AS a".to_string(),
"/* multi;\nline */ SELECT 1".to_string()
]
);
}
#[test]
fn split_ignores_semicolons_in_dollar_quoted_strings() {
let script = "SELECT $$a;b$$; SELECT $body$c;d$body$; SELECT 1";
assert_eq!(
split_sql_statements(script),
vec!["SELECT $$a;b$$".to_string(), "SELECT $body$c;d$body$".to_string(), "SELECT 1".to_string()]
);
// $1 placeholders are not dollar-quote openers
assert_eq!(split_sql_statements("SELECT $1; SELECT $2").len(), 2);
// unterminated dollar quote swallows the rest instead of splitting it
assert_eq!(split_sql_statements("SELECT $$a;b").len(), 1);
}
#[test]
fn split_ignores_escaped_quotes_and_semicolons_in_escape_strings() {
let script = r#"SELECT E'it\'s;ok'; SELECT e'path\\;name'; SELECT 2"#;
assert_eq!(
split_sql_statements(script),
vec![r#"SELECT E'it\'s;ok'"#.to_string(), r#"SELECT e'path\\;name'"#.to_string(), "SELECT 2".to_string()]
);
}
#[test]
fn init_script_names_collects_attach_aliases() {
let script = "INSTALL httpfs; LOAD httpfs;\nCREATE SECRET (TYPE s3, KEY_ID 'k', SECRET 's');\nATTACH 'ducklake:postgres:dbname=lake' AS lake;\nATTACH 'sales.duckdb' AS \"sales db\";";
assert_eq!(init_script_attached_names(script), vec!["lake".to_string(), "sales db".to_string()]);
}
}
crates/dbx-core/src/db/duckdb_worker_process.rs
+15 -4
View File
@@ -63,29 +63,37 @@ struct WorkerProcessState {
}
impl DuckDbWorkerClient {
pub async fn open(path: String, attached_databases: Vec<AttachedDatabaseConfig>) -> Result<Self, String> {
pub async fn open(
path: String,
attached_databases: Vec<AttachedDatabaseConfig>,
init_script: Option<String>,
) -> Result<Self, String> {
let executable = std::env::current_exe().map_err(|e| e.to_string())?;
Self::open_with_executable(executable, path, attached_databases).await
Self::open_with_executable(executable, path, attached_databases, init_script).await
}
pub async fn open_with_process_limit(
path: String,
attached_databases: Vec<AttachedDatabaseConfig>,
init_script: Option<String>,
process_limit: usize,
) -> Result<Self, String> {
let executable = std::env::current_exe().map_err(|e| e.to_string())?;
Self::open_with_executable_and_process_limit(executable, path, attached_databases, process_limit).await
Self::open_with_executable_and_process_limit(executable, path, attached_databases, init_script, process_limit)
.await
}
pub async fn open_with_executable(
executable: PathBuf,
path: String,
attached_databases: Vec<AttachedDatabaseConfig>,
init_script: Option<String>,
) -> Result<Self, String> {
Self::open_with_executable_and_process_limit(
executable,
path,
attached_databases,
init_script,
DUCKDB_WORKER_MAX_PROCESSES_DEFAULT,
)
.await
@@ -95,12 +103,14 @@ impl DuckDbWorkerClient {
executable: PathBuf,
path: String,
attached_databases: Vec<AttachedDatabaseConfig>,
init_script: Option<String>,
process_limit: usize,
) -> Result<Self, String> {
let client = Self::new_unconnected_with_timeouts(
executable,
path,
attached_databases,
init_script,
process_limit,
DEFAULT_WORKER_REQUEST_TIMEOUT,
DEFAULT_WORKER_START_WAIT,
@@ -114,6 +124,7 @@ impl DuckDbWorkerClient {
executable: PathBuf,
path: String,
attached_databases: Vec<AttachedDatabaseConfig>,
init_script: Option<String>,
process_limit: usize,
request_timeout: Duration,
worker_start_timeout: Duration,
@@ -128,7 +139,7 @@ impl DuckDbWorkerClient {
process_limiter,
process_limit,
executable,
connect_params: DuckDbWorkerConnectParams { path, attached_databases },
connect_params: DuckDbWorkerConnectParams { path, attached_databases, init_script },
request_timeout,
worker_start_timeout,
next_id: AtomicU64::new(1),
crates/dbx-core/src/db/duckdb_worker_protocol.rs
+3
View File
@@ -100,6 +100,8 @@ pub struct DuckDbWorkerConnectParams {
pub path: String,
#[serde(default)]
pub attached_databases: Vec<AttachedDatabaseConfig>,
#[serde(default)]
pub init_script: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
@@ -148,6 +150,7 @@ mod tests {
name: "analytics".to_string(),
path: "D:/tmp/analytics.duckdb".to_string(),
}],
init_script: None,
},
)
.expect("serialize request");
crates/dbx-core/src/db/duckdb_worker_runtime.rs
+119 -92
View File
@@ -30,6 +30,20 @@ impl DuckDbWorkerSession {
crate::schema::duckdb_attach_database(&locked, &attached.name, &path)?;
attached_names.push(attached.name.clone());
}
if let Some(script) = params.init_script.as_deref() {
// Track script attachments by diffing the actual catalog list
// instead of parsing the SQL, so ATTACH without an alias,
// dynamic statements, and extension-specific syntax all count.
let before = duckdb_catalog_list(&locked)?;
db::duckdb_driver::run_init_script(&locked, script)?;
for name in duckdb_catalog_list(&locked)? {
if !before.iter().any(|existing| existing.eq_ignore_ascii_case(&name))
&& !attached_names.iter().any(|attached| attached.eq_ignore_ascii_case(&name))
{
attached_names.push(name);
}
}
}
}
self.connection = Some(connection);
self.attached_names = attached_names;
@@ -46,7 +60,7 @@ impl DuckDbWorkerSession {
&params.sql,
params.max_rows,
)?;
if let Some(name) = duckdb_attached_name_from_attach_sql(&params.sql) {
if let Some(name) = crate::db::duckdb_sql::attached_name_from_attach_sql(&params.sql) {
if !self.attached_names.iter().any(|attached| attached.eq_ignore_ascii_case(&name)) {
self.attached_names.push(name);
}
@@ -125,6 +139,12 @@ impl DuckDbWorkerSession {
}
}
fn duckdb_catalog_list(con: &duckdb::Connection) -> Result<Vec<String>, String> {
let mut stmt = con.prepare("SHOW DATABASES").map_err(|e| e.to_string())?;
let rows = stmt.query_map([], |row| row.get::<_, String>(0)).map_err(|e| e.to_string())?;
Ok(rows.filter_map(|row| row.ok()).collect())
}
#[derive(Clone, Default)]
struct DuckDbWorkerRuntime {
session: Arc<Mutex<DuckDbWorkerSession>>,
@@ -334,83 +354,6 @@ async fn write_response(stdout: Arc<tokio::sync::Mutex<Stdout>>, response: &Duck
let _ = stdout.flush().await;
}
fn duckdb_attached_name_from_attach_sql(sql: &str) -> Option<String> {
let trimmed = sql.trim_start();
let first_word = trimmed.split(|ch: char| ch.is_whitespace() || ch == ';').next().unwrap_or_default();
if !first_word.eq_ignore_ascii_case("ATTACH") {
return None;
}
let as_index = find_as_keyword_outside_quotes(trimmed)?;
parse_identifier_after_as(&trimmed[as_index + 2..])
}
fn find_as_keyword_outside_quotes(sql: &str) -> Option<usize> {
let bytes = sql.as_bytes();
let mut i = 0;
let mut in_single = false;
let mut in_double = false;
while i < bytes.len() {
match bytes[i] {
b'\'' if !in_double => {
if in_single && i + 1 < bytes.len() && bytes[i + 1] == b'\'' {
i += 2;
continue;
}
in_single = !in_single;
}
b'"' if !in_single => {
if in_double && i + 1 < bytes.len() && bytes[i + 1] == b'"' {
i += 2;
continue;
}
in_double = !in_double;
}
b'a' | b'A' if !in_single && !in_double && i + 1 < bytes.len() => {
if (bytes[i + 1] == b's' || bytes[i + 1] == b'S')
&& is_sql_word_boundary(bytes.get(i.wrapping_sub(1)).copied())
&& is_sql_word_boundary(bytes.get(i + 2).copied())
{
return Some(i);
}
}
_ => {}
}
i += 1;
}
None
}
fn is_sql_word_boundary(byte: Option<u8>) -> bool {
!matches!(byte, Some(b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'_'))
}
fn parse_identifier_after_as(input: &str) -> Option<String> {
let input = input.trim_start();
if input.is_empty() {
return None;
}
if let Some(rest) = input.strip_prefix('"') {
let mut name = String::new();
let mut chars = rest.chars().peekable();
while let Some(ch) = chars.next() {
if ch == '"' {
if chars.peek() == Some(&'"') {
name.push('"');
chars.next();
continue;
}
return (!name.trim().is_empty()).then_some(name);
}
name.push(ch);
}
return None;
}
let name = input.split(|ch: char| ch.is_whitespace() || ch == ';').next().unwrap_or_default().trim();
(!name.is_empty()).then(|| name.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
@@ -420,7 +363,11 @@ mod tests {
let mut session = DuckDbWorkerSession::default();
session
.connect(DuckDbWorkerConnectParams { path: ":memory:".to_string(), attached_databases: Vec::new() })
.connect(DuckDbWorkerConnectParams {
path: ":memory:".to_string(),
attached_databases: Vec::new(),
init_script: None,
})
.expect("connect");
assert_eq!(session.list_databases().expect("list databases")[0].name, "main");
@@ -430,7 +377,11 @@ mod tests {
fn worker_session_executes_select_query() {
let mut session = DuckDbWorkerSession::default();
session
.connect(DuckDbWorkerConnectParams { path: ":memory:".to_string(), attached_databases: Vec::new() })
.connect(DuckDbWorkerConnectParams {
path: ":memory:".to_string(),
attached_databases: Vec::new(),
init_script: None,
})
.expect("connect");
let result = session
@@ -445,7 +396,11 @@ mod tests {
fn worker_session_lists_tables_and_columns() {
let mut session = DuckDbWorkerSession::default();
session
.connect(DuckDbWorkerConnectParams { path: ":memory:".to_string(), attached_databases: Vec::new() })
.connect(DuckDbWorkerConnectParams {
path: ":memory:".to_string(),
attached_databases: Vec::new(),
init_script: None,
})
.expect("connect");
session
.execute(DuckDbWorkerExecuteParams {
@@ -484,7 +439,11 @@ mod tests {
}
let mut session = DuckDbWorkerSession::default();
session
.connect(DuckDbWorkerConnectParams { path: ":memory:".to_string(), attached_databases: Vec::new() })
.connect(DuckDbWorkerConnectParams {
path: ":memory:".to_string(),
attached_databases: Vec::new(),
init_script: None,
})
.expect("connect");
session
@@ -504,6 +463,79 @@ mod tests {
let _ = std::fs::remove_dir_all(dir);
}
#[test]
fn worker_session_runs_init_script_and_tracks_attached_alias() {
let dir = std::env::temp_dir().join(format!("dbx-duckdb-worker-{}", uuid::Uuid::new_v4()));
std::fs::create_dir_all(&dir).expect("create temp dir");
let attached_path = dir.join("warehouse.duckdb");
{
let con = duckdb::Connection::open(&attached_path).expect("create attached db");
con.execute_batch("CREATE TABLE facts (id INTEGER);").expect("create attached table");
}
let mut session = DuckDbWorkerSession::default();
let script = format!(
"SET memory_limit='512MB'; -- session tuning\nATTACH '{}' AS warehouse;\nCREATE TABLE staging (id INTEGER);",
attached_path.to_string_lossy().replace('\'', "''")
);
session
.connect(DuckDbWorkerConnectParams {
path: ":memory:".to_string(),
attached_databases: Vec::new(),
init_script: Some(script),
})
.expect("connect");
assert!(session.attached_names.iter().any(|name| name == "warehouse"));
let databases = session.list_databases().expect("list databases");
assert!(databases.iter().any(|database| database.name == "warehouse"));
let tables = session
.list_tables(DuckDbWorkerTableParams { database: "warehouse".to_string(), schema: "main".to_string() })
.expect("list tables");
assert!(tables.iter().any(|table| table.name == "facts"));
let _ = std::fs::remove_dir_all(dir);
}
#[test]
fn worker_session_tracks_init_script_attach_without_alias() {
let dir = std::env::temp_dir().join(format!("dbx-duckdb-worker-{}", uuid::Uuid::new_v4()));
std::fs::create_dir_all(&dir).expect("create temp dir");
let attached_path = dir.join("noalias.duckdb");
{
let con = duckdb::Connection::open(&attached_path).expect("create attached db");
con.execute_batch("CREATE TABLE items (id INTEGER);").expect("create attached table");
}
let mut session = DuckDbWorkerSession::default();
session
.connect(DuckDbWorkerConnectParams {
path: ":memory:".to_string(),
attached_databases: Vec::new(),
init_script: Some(format!("ATTACH '{}';", attached_path.to_string_lossy().replace('\'', "''"))),
})
.expect("connect");
assert!(session.attached_names.iter().any(|name| name == "noalias"), "names: {:?}", session.attached_names);
let tables = session
.list_tables(DuckDbWorkerTableParams { database: "noalias".to_string(), schema: "main".to_string() })
.expect("list tables");
assert!(tables.iter().any(|table| table.name == "items"));
let _ = std::fs::remove_dir_all(dir);
}
#[test]
fn worker_session_reports_failing_init_script_statement() {
let mut session = DuckDbWorkerSession::default();
let err = session
.connect(DuckDbWorkerConnectParams {
path: ":memory:".to_string(),
attached_databases: Vec::new(),
init_script: Some("SELECT 1; SELECT * FROM missing_table;".to_string()),
})
.expect_err("init script should fail");
assert!(err.contains("statement 2"), "unexpected error: {err}");
}
#[test]
fn worker_session_tracks_attach_sql_alias() {
let dir = std::env::temp_dir().join(format!("dbx-duckdb-worker-{}", uuid::Uuid::new_v4()));
@@ -515,7 +547,11 @@ mod tests {
}
let mut session = DuckDbWorkerSession::default();
session
.connect(DuckDbWorkerConnectParams { path: ":memory:".to_string(), attached_databases: Vec::new() })
.connect(DuckDbWorkerConnectParams {
path: ":memory:".to_string(),
attached_databases: Vec::new(),
init_script: None,
})
.expect("connect");
session
@@ -554,13 +590,4 @@ mod tests {
assert!(!response.ok);
assert_eq!(response.error.expect("error").code, "duckdb_worker_busy");
}
#[test]
fn attach_sql_alias_parser_handles_generated_sql() {
assert_eq!(
duckdb_attached_name_from_attach_sql("ATTACH 'D:\\tmp\\sales.duckdb' AS \"sales db\";"),
Some("sales db".to_string())
);
assert_eq!(duckdb_attached_name_from_attach_sql("select 'not attach' as value"), None);
}
}
crates/dbx-core/src/db/mod.rs
+1
View File
@@ -3,6 +3,7 @@ pub mod clickhouse_driver;
pub mod cloudflare_d1;
pub use cloudflare_d1 as cloudflare_d1_driver;
pub mod duckdb_driver;
pub mod duckdb_sql;
#[cfg(feature = "duckdb-bundled")]
pub mod duckdb_worker_process;
#[cfg(feature = "duckdb-bundled")]
crates/dbx-core/src/db/redis_driver.rs
+1
View File
@@ -3078,6 +3078,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: crate::models::connection::default_connect_timeout_secs(),
crates/dbx-core/src/models/connection.rs
+8
View File
@@ -27,6 +27,10 @@ pub struct ConnectionConfig {
pub visible_schemas: Option<HashMap<String, Vec<String>>>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub attached_databases: Vec<AttachedDatabaseConfig>,
/// SQL statements executed right after the connection is established
/// (DuckDB only for now): INSTALL/LOAD, SET, CREATE SECRET, ATTACH, ...
#[serde(default, skip_serializing_if = "Option::is_none")]
pub init_script: Option<String>,
#[serde(default)]
pub color: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
@@ -477,6 +481,8 @@ struct ConnectionConfigData {
#[serde(default)]
pub attached_databases: Vec<AttachedDatabaseConfig>,
#[serde(default)]
pub init_script: Option<String>,
#[serde(default)]
pub color: Option<String>,
#[serde(default)]
pub transport_layers: Vec<TransportLayerConfig>,
@@ -560,6 +566,7 @@ impl From<ConnectionConfigData> for ConnectionConfig {
visible_databases: data.visible_databases,
visible_schemas: data.visible_schemas,
attached_databases: data.attached_databases,
init_script: data.init_script,
color: data.color,
transport_layers: data.transport_layers,
connect_timeout_secs: data.connect_timeout_secs,
@@ -1902,6 +1909,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: super::default_connect_timeout_secs(),
crates/dbx-core/src/mq/config.rs
+1
View File
@@ -114,6 +114,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
crates/dbx-core/src/mq/service.rs
+1
View File
@@ -562,6 +562,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
crates/dbx-core/src/nacos/config.rs
+1
View File
@@ -139,6 +139,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
crates/dbx-core/src/nacos/service.rs
+2
View File
@@ -183,6 +183,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
@@ -250,6 +251,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
crates/dbx-core/src/production_safety.rs
+1
View File
@@ -579,6 +579,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: vec![],
init_script: None,
color: None,
transport_layers: vec![],
connect_timeout_secs: 10,
crates/dbx-core/src/query.rs
+3 -3
View File
@@ -1107,9 +1107,7 @@ pub async fn do_execute(
let (_duckdb_attached_names, read_only_connection) = {
let configs = state.configs.read().await;
let config = crate::connection::config_for_pool_key(pool_key, &configs);
let attached = config
.map(|c| c.attached_databases.iter().map(|db| db.name.clone()).collect::<Vec<_>>())
.unwrap_or_default();
let attached = config.map(crate::db::duckdb_sql::config_attached_names).unwrap_or_default();
let connection = config.filter(|config| config.read_only).map(|config| (config.name.clone(), config.db_type));
(attached, connection)
};
@@ -3159,6 +3157,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 10,
@@ -4029,6 +4028,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
crates/dbx-core/src/schema.rs
+154 -11
View File
@@ -74,7 +74,13 @@ pub fn duckdb_query_tables_in_database_with_attached(
})
})
.map_err(|e| e.to_string())?;
Ok(rows.filter_map(|r| r.ok()).collect())
let tables: Vec<db::TableInfo> = rows.filter_map(|r| r.ok()).collect();
if tables.is_empty() && duckdb_is_quack_catalog(con, &database) {
if let Some(remote) = duckdb_quack_remote_tables(con, &database, schema) {
return Ok(remote);
}
}
Ok(tables)
}
#[cfg(feature = "duckdb-bundled")]
@@ -85,7 +91,7 @@ pub fn duckdb_attach_database(con: &duckdb::Connection, name: &str, path: &str)
return Err("DuckDB attached database name and path are required".to_string());
}
let sql = format!("ATTACH {} AS {}", duckdb_quote_string(path), duckdb_quote_ident(name));
con.execute_batch(&sql).map_err(|e| e.to_string())
con.execute_batch(&sql).map_err(|e| format!("Failed to attach database \"{name}\": {e}"))
}
#[cfg(feature = "duckdb-bundled")]
@@ -127,7 +133,143 @@ pub fn duckdb_list_schemas_with_attached(
)
.map_err(|e| e.to_string())?;
let rows = stmt.query_map([database.as_str()], |row| row.get::<_, String>(0)).map_err(|e| e.to_string())?;
Ok(rows.filter_map(|r| r.ok()).collect())
let schemas: Vec<String> = rows.filter_map(|r| r.ok()).collect();
if schemas.is_empty() && duckdb_is_quack_catalog(con, &database) {
if let Some(remote) = duckdb_quack_remote_schemas(con, &database) {
return Ok(remote);
}
}
Ok(schemas)
}
/// 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
/// derived catalog name (e.g. `localhost:9494`) that no parser sees.
#[cfg(feature = "duckdb-bundled")]
fn duckdb_is_quack_catalog(con: &duckdb::Connection, database: &str) -> bool {
con.query_row("SELECT type FROM duckdb_databases() WHERE lower(database_name) = lower(?)", [database], |row| {
row.get::<_, String>(0)
})
.map(|catalog_type| catalog_type.eq_ignore_ascii_case("quack"))
.unwrap_or(false)
}
/// Quack-attached catalogs expose no metadata on the client side (the beta
/// extension implements name resolution and streaming scans, but not catalog
/// enumeration: `information_schema`, `duckdb_tables()` and `SHOW ALL TABLES`
/// all skip the remote catalog). The extension's `quack_query_by_name` table
/// function runs SQL on the remote server, so when a metadata query for an
/// attached catalog comes back empty we ask the remote for its own
/// information_schema. For non-quack catalogs (or when the extension is not
/// loaded) the function call errors and the fallback quietly yields `None`.
#[cfg(feature = "duckdb-bundled")]
fn duckdb_quack_remote_query<T>(
con: &duckdb::Connection,
alias: &str,
remote_sql: &str,
map_row: impl Fn(&duckdb::Row<'_>) -> Result<T, duckdb::Error>,
) -> Option<Vec<T>> {
let sql = format!(
"SELECT * FROM quack_query_by_name({}, {})",
duckdb_quote_string(alias),
duckdb_quote_string(remote_sql)
);
let mut stmt = match con.prepare(&sql) {
Ok(stmt) => stmt,
Err(e) => {
log::debug!("quack metadata fallback unavailable for catalog {alias}: {e}");
return None;
}
};
let rows = match stmt.query_map([], |row| map_row(row)) {
Ok(rows) => rows,
Err(e) => {
log::debug!("quack metadata fallback query failed for catalog {alias}: {e}");
return None;
}
};
Some(rows.filter_map(|r| r.ok()).collect())
}
#[cfg(feature = "duckdb-bundled")]
fn duckdb_quack_remote_schemas(con: &duckdb::Connection, alias: &str) -> Option<Vec<String>> {
duckdb_quack_remote_query(
con,
alias,
"SELECT DISTINCT schema_name FROM information_schema.schemata WHERE schema_name NOT IN ('information_schema', 'pg_catalog') ORDER BY schema_name",
|row| row.get::<_, String>(0),
)
}
#[cfg(feature = "duckdb-bundled")]
fn duckdb_quack_remote_tables(con: &duckdb::Connection, alias: &str, schema: &str) -> Option<Vec<db::TableInfo>> {
let remote_sql = format!(
"SELECT table_name, table_type FROM information_schema.tables WHERE table_schema = {} ORDER BY table_name",
duckdb_quote_string(schema)
);
duckdb_quack_remote_query(con, alias, &remote_sql, |row| {
Ok(db::TableInfo {
name: row.get::<_, String>(0)?,
table_type: row.get::<_, String>(1)?,
comment: None,
parent_schema: None,
parent_name: None,
})
})
}
#[cfg(feature = "duckdb-bundled")]
fn duckdb_quack_remote_columns(
con: &duckdb::Connection,
alias: &str,
schema: &str,
table: &str,
) -> Option<Vec<db::ColumnInfo>> {
let pk_sql = format!(
"SELECT kcu.column_name
FROM information_schema.table_constraints tc
JOIN information_schema.key_column_usage kcu
ON tc.constraint_name = kcu.constraint_name
AND tc.table_schema = kcu.table_schema
AND tc.table_name = kcu.table_name
WHERE tc.constraint_type = 'PRIMARY KEY'
AND tc.table_schema = {schema}
AND tc.table_name = {table}
ORDER BY kcu.ordinal_position",
schema = duckdb_quote_string(schema),
table = duckdb_quote_string(table),
);
let primary_keys: std::collections::HashSet<String> =
duckdb_quack_remote_query(con, alias, &pk_sql, |row| row.get::<_, String>(0))
.map(|names| names.into_iter().collect())
.unwrap_or_default();
let columns_sql = format!(
"SELECT column_name, data_type, is_nullable, column_default
FROM information_schema.columns
WHERE table_schema = {schema} AND table_name = {table}
ORDER BY ordinal_position",
schema = duckdb_quote_string(schema),
table = duckdb_quote_string(table),
);
duckdb_quack_remote_query(con, alias, &columns_sql, |row| {
let name = row.get::<_, String>(0)?;
Ok(db::ColumnInfo {
is_primary_key: primary_keys.contains(&name),
name,
data_type: row.get::<_, String>(1)?,
is_nullable: row.get::<_, String>(2).unwrap_or_default() == "YES",
column_default: row.get::<_, Option<String>>(3)?,
extra: None,
comment: None,
numeric_precision: None,
numeric_scale: None,
character_maximum_length: None,
enum_values: None,
..Default::default()
})
})
}
#[cfg(feature = "duckdb-bundled")]
@@ -243,7 +385,13 @@ pub fn duckdb_query_columns_in_database_with_attached(
})
})
.map_err(|e| e.to_string())?;
Ok(rows.filter_map(|r| r.ok()).collect())
let columns: Vec<db::ColumnInfo> = rows.filter_map(|r| r.ok()).collect();
if columns.is_empty() && duckdb_is_quack_catalog(con, &database) {
if let Some(remote) = duckdb_quack_remote_columns(con, &database, schema, table) {
return Ok(remote);
}
}
Ok(columns)
}
#[cfg(feature = "duckdb-bundled")]
@@ -454,13 +602,7 @@ fn duckdb_completion_like_pattern(request: &db::CompletionAssistantRequest) -> S
#[cfg(feature = "duckdb-bundled")]
async fn duckdb_attached_database_names(state: &AppState, connection_id: &str) -> Vec<String> {
state
.configs
.read()
.await
.get(connection_id)
.map(|config| config.attached_databases.iter().map(|database| database.name.clone()).collect())
.unwrap_or_default()
state.configs.read().await.get(connection_id).map(crate::db::duckdb_sql::config_attached_names).unwrap_or_default()
}
fn clickhouse_metadata_database<'a>(database: &'a str, schema: &'a str) -> &'a str {
@@ -2326,6 +2468,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
crates/dbx-core/src/schema/duckdb_metadata.rs
+1 -1
View File
@@ -198,6 +198,6 @@ pub(super) async fn duckdb_attached_database_names(state: &AppState, connection_
.read()
.await
.get(connection_id)
.map(|config| config.attached_databases.iter().map(|database| database.name.clone()).collect())
.map(crate::db::duckdb_sql::config_attached_names)
.unwrap_or_default()
}
crates/dbx-core/src/schema/normalization.rs
+1
View File
@@ -160,6 +160,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
crates/dbx-core/src/storage.rs
+14
View File
@@ -1265,6 +1265,7 @@ impl Storage {
scrub_transport_layer_secrets(&mut sanitized);
sanitized.redis_sentinel_password = String::new();
sanitized.connection_string = None;
sanitized.init_script = None;
scrub_mq_auth_secrets(&mut sanitized);
scrub_mq_token_signing_secret(&mut sanitized);
scrub_nacos_auth_secrets(&mut sanitized);
@@ -1302,6 +1303,7 @@ impl Storage {
scrub_transport_layer_secrets(&mut sanitized);
sanitized.redis_sentinel_password = String::new();
sanitized.connection_string = None;
sanitized.init_script = None;
scrub_mq_auth_secrets(&mut sanitized);
scrub_mq_token_signing_secret(&mut sanitized);
scrub_nacos_auth_secrets(&mut sanitized);
@@ -1360,6 +1362,15 @@ impl Storage {
)
.map_err(|e| e.to_string())?;
}
if let Some(script) = &config.init_script {
persist_secret_in_tx(&tx, &config.id, "init_script", script)?;
} else {
tx.execute(
"DELETE FROM connection_secrets WHERE connection_id = ?1 AND key = ?2",
params![config.id, "init_script"],
)
.map_err(|e| e.to_string())?;
}
persist_mq_auth_secrets_in_tx(&tx, &config)?;
persist_mq_token_signing_secret_in_tx(&tx, &config)?;
persist_nacos_auth_secrets_in_tx(&tx, &config)?;
@@ -1447,6 +1458,7 @@ impl Storage {
}
config.redis_sentinel_password = self.get_secret(&id, "redis_sentinel_password").await?.unwrap_or_default();
config.connection_string = self.get_secret(&id, "connection_string").await?;
config.init_script = self.get_secret(&id, "init_script").await?;
let needs_mq_auth_rewrite = self.hydrate_mq_auth_secrets(&id, &mut config).await?;
let needs_mq_token_signing_rewrite = self.hydrate_mq_token_signing_secret(&id, &mut config).await?;
let needs_nacos_auth_rewrite = self.hydrate_nacos_auth_secret(&id, &mut config).await?;
@@ -2585,6 +2597,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 30,
@@ -2644,6 +2657,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 30,
crates/dbx-core/src/transfer.rs
+1
View File
@@ -4756,6 +4756,7 @@ mod tests {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
crates/dbx-core/tests/duckdb_worker_process.rs
+12 -5
View File
@@ -26,7 +26,7 @@ async fn worker_process_recovers_immediately_after_cancelled_long_query() {
let _ = std::fs::remove_file(&db_path);
let client =
DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new())
DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new(), None)
.await
.expect("worker process connects");
@@ -77,7 +77,7 @@ async fn worker_process_recovers_after_registered_cancel_interrupt() {
let _ = std::fs::remove_file(&db_path);
let client =
DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new())
DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new(), None)
.await
.expect("worker process connects");
@@ -140,7 +140,7 @@ async fn worker_process_recovers_after_parser_error() {
let _ = std::fs::remove_file(&db_path);
let client =
DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new())
DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new(), None)
.await
.expect("worker process connects");
@@ -190,7 +190,7 @@ async fn worker_process_keeps_session_state_after_benign_error() {
let _ = std::fs::remove_file(&db_path);
let client =
DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new())
DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new(), None)
.await
.expect("worker process connects");
@@ -251,7 +251,11 @@ async fn worker_process_exits_when_stdin_closes_during_active_query() {
DuckDbWorkerRequest::new(
"connect",
DuckDbWorkerMethod::Connect,
DuckDbWorkerConnectParams { path: db_path.to_string_lossy().to_string(), attached_databases: Vec::new() },
DuckDbWorkerConnectParams {
path: db_path.to_string_lossy().to_string(),
attached_databases: Vec::new(),
init_script: None,
},
)
.expect("connect request"),
)
@@ -298,6 +302,7 @@ async fn worker_process_is_killed_after_connect_timeout() {
executable,
db_path.to_string_lossy().to_string(),
Vec::new(),
None,
4,
Duration::from_secs(1),
Duration::from_secs(5),
@@ -332,6 +337,7 @@ async fn worker_process_is_killed_after_connect_error() {
executable,
db_path.to_string_lossy().to_string(),
Vec::new(),
None,
4,
Duration::from_secs(5),
Duration::from_secs(5),
@@ -386,6 +392,7 @@ async fn worker_process_retries_connect_after_transient_file_lock() {
executable,
db_path.to_string_lossy().to_string(),
Vec::new(),
None,
4,
Duration::from_secs(5),
Duration::from_secs(5),
crates/dbx-core/tests/live_postgres_query_result_export.rs
+1
View File
@@ -31,6 +31,7 @@ fn live_postgres_config(
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 10,
crates/dbx-core/tests/live_postgres_transfer.rs
+1
View File
@@ -25,6 +25,7 @@ fn postgres_test_config(id: &str, database: &str) -> ConnectionConfig {
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
crates/dbx-core/tests/live_sqlserver_completion.rs
+1
View File
@@ -28,6 +28,7 @@ fn live_sqlserver_config(id: &str, database: &str) -> dbx_core::models::connecti
visible_databases: None,
visible_schemas: None,
attached_databases: Vec::new(),
init_script: None,
color: None,
transport_layers: Vec::new(),
connect_timeout_secs: 10,