pub mod agent_driver; pub mod clickhouse_driver; pub mod cloudberry; pub mod cloudflare_d1; pub use cloudflare_d1 as cloudflare_d1_driver; pub(crate) mod ddl_scan; pub mod document_result; pub mod dolt; pub mod doris; pub mod duckdb_sql; #[cfg(feature = "duckdb-sidecar")] pub mod duckdb_worker_process; #[cfg(feature = "duckdb-sidecar")] pub mod duckdb_worker_protocol; #[cfg(feature = "dynamodb")] #[path = "dynamodb_driver.rs"] pub mod dynamodb_driver; #[cfg(not(feature = "dynamodb"))] #[path = "dynamodb_driver_disabled.rs"] pub mod dynamodb_driver; pub mod easysearch_driver; pub mod elasticsearch_driver; pub mod elasticsearch_sql; pub mod file_validator; pub mod hbase_driver; pub mod http_tunnel; pub mod influxdb_driver; pub mod manticoresearch; pub mod meilisearch_driver; pub mod mongo_driver; pub mod mysql; pub mod mysql_compatible; pub mod ob_oracle; pub mod oceanbase_mysql; pub mod opentenbase; pub mod postgres; pub mod proxy_tunnel; pub mod questdb; pub mod redis_driver; pub mod rqlite_driver; pub mod sqlite; pub mod sqlserver; pub mod ssh_host_key; pub mod ssh_prompt; pub mod ssh_tunnel; pub mod starrocks; pub mod tdsql_mysql; pub mod tidb; pub mod transport_layer_tunnel; pub mod turso_driver; pub mod vector_driver; pub mod victoriametrics_driver; pub mod wkb; use reqwest::ClientBuilder; use std::fmt; use std::future::Future; use std::time::Duration; // Re-export types so that `db::QueryResult` etc. work within dbx-core pub use crate::mysql_event_sql::MysqlEventInfo; pub use crate::types::*; pub use file_validator::validate_file_path; pub const CONNECTION_TIMEOUT_SECS: u64 = 5; pub const TCP_PROBE_TIMEOUT_SECS: u64 = 3; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PoolCheckoutStage { Wait, Create, Recycle, Unknown, } impl PoolCheckoutStage { pub fn as_str(self) -> &'static str { match self { Self::Wait => "wait", Self::Create => "create", Self::Recycle => "recycle", Self::Unknown => "unknown", } } } #[derive(Debug, Clone, PartialEq, Eq)] pub enum PoolCheckoutError { Timeout { database: &'static str, stage: PoolCheckoutStage, timeout: Duration }, Failed { database: &'static str, stage: PoolCheckoutStage, detail: String }, Canceled, } impl PoolCheckoutError { pub fn stage(&self) -> Option { match self { Self::Timeout { stage, .. } | Self::Failed { stage, .. } => Some(*stage), Self::Canceled => None, } } pub fn is_pool_saturation(&self) -> bool { matches!(self, Self::Timeout { stage: PoolCheckoutStage::Wait, .. }) } } impl fmt::Display for PoolCheckoutError { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::Timeout { database, stage, timeout } => write!( formatter, "{database} connection pool checkout timed out [stage={}, timeout_ms={}]", stage.as_str(), timeout.as_millis() ), Self::Failed { database, stage, detail } => { write!(formatter, "{database} connection pool checkout failed [stage={}]: {detail}", stage.as_str()) } Self::Canceled => formatter.write_str(crate::query::QUERY_CANCELED), } } } pub fn connection_timeout() -> Duration { Duration::from_secs(CONNECTION_TIMEOUT_SECS) } pub fn http_client_builder(timeout: Duration) -> ClientBuilder { reqwest::Client::builder().connect_timeout(timeout).no_proxy() } const JS_MAX_SAFE_INTEGER: i64 = 9_007_199_254_740_991; pub fn safe_i64_to_json(v: i64) -> serde_json::Value { if !(-JS_MAX_SAFE_INTEGER..=JS_MAX_SAFE_INTEGER).contains(&v) { serde_json::Value::String(v.to_string()) } else { serde_json::Value::Number(v.into()) } } pub fn safe_u64_to_json(v: u64) -> serde_json::Value { if v > JS_MAX_SAFE_INTEGER as u64 { serde_json::Value::String(v.to_string()) } else { serde_json::Value::Number(v.into()) } } pub fn json_value_for_js(value: serde_json::Value) -> serde_json::Value { match value { serde_json::Value::Number(number) => { if let Some(value) = number.as_i64() { safe_i64_to_json(value) } else if let Some(value) = number.as_u64() { safe_u64_to_json(value) } else { let value = number.to_string(); if value.bytes().all(|byte| byte == b'-' || byte.is_ascii_digit()) { serde_json::Value::String(value) } else { serde_json::Value::Number(number) } } } serde_json::Value::Array(values) => { serde_json::Value::Array(values.into_iter().map(json_value_for_js).collect()) } serde_json::Value::Object(entries) => { serde_json::Value::Object(entries.into_iter().map(|(key, value)| (key, json_value_for_js(value))).collect()) } value => value, } } pub(crate) fn hex_encode(bytes: &[u8]) -> String { const HEX: &[u8; 16] = b"0123456789abcdef"; let mut out = String::with_capacity(bytes.len() * 2); for &byte in bytes { out.push(HEX[(byte >> 4) as usize] as char); out.push(HEX[(byte & 0x0f) as usize] as char); } out } pub(crate) fn binary_value_to_json(bytes: &[u8]) -> serde_json::Value { serde_json::Value::String(format!("0x{}", hex_encode(bytes))) } pub fn tcp_probe_timeout() -> Duration { Duration::from_secs(TCP_PROBE_TIMEOUT_SECS) } pub fn parse_connect_timeout(url: &str) -> Duration { parse_connect_timeout_with_fallback(url, connection_timeout()) } pub fn parse_connect_timeout_with_fallback(url: &str, fallback: Duration) -> Duration { let Some(query) = url.split('?').nth(1) else { return fallback; }; for param in query.split('&') { let trimmed = param.trim(); if trimmed.is_empty() { continue; } let (key, value) = match trimmed.split_once('=') { Some(pair) => pair, None => continue, }; if key.eq_ignore_ascii_case("connect_timeout") || key.eq_ignore_ascii_case("connectTimeout") || key.eq_ignore_ascii_case("connection_timeout") || key.eq_ignore_ascii_case("connectionTimeout") { if let Ok(v) = value.parse::() { if (1..=300).contains(&v) { return Duration::from_secs(v); } } } } fallback } pub async fn with_connection_timeout(label: &str, timeout: Duration, future: F) -> Result where F: Future>, { tokio::time::timeout(timeout, future) .await .map_err(|_| format!("{label} connection timed out ({}s)", timeout.as_secs()))? } pub async fn probe_tcp_endpoint(label: &str, host: &str, port: u16, timeout: Duration) -> Result<(), String> { tokio::time::timeout(timeout, tokio::net::TcpStream::connect((host, port))) .await .map_err(|_| format!("{label} TCP connection timed out ({}s)", timeout.as_secs()))? .map(|_| ()) .map_err(|e| format!("{label} TCP connection failed: {e}")) } #[cfg(test)] mod tests { use super::*; #[test] fn binary_values_are_displayed_as_prefixed_hex() { assert_eq!(binary_value_to_json(&[0x00, 0x01, 0xab, 0xff]), serde_json::json!("0x0001abff")); } #[test] fn arbitrary_precision_integers_are_strings_for_js() { let value = serde_json::from_str( r#"{ "positive": 20240302001417986771, "negative": -20240302001417986771 }"#, ) .unwrap(); assert_eq!( json_value_for_js(value), serde_json::json!({ "positive": "20240302001417986771", "negative": "-20240302001417986771" }) ); } #[test] fn json_value_for_js_preserves_existing_value_behavior_recursively() { let value = serde_json::from_str( r#"{ "safe": [9007199254740991, -9007199254740991], "unsafe": [9007199254740992, -9007199254740992], "non_integer": [1.25, 1e-2], "other": [null, true, "text"] }"#, ) .unwrap(); let expected: serde_json::Value = serde_json::from_str( r#"{ "safe": [9007199254740991, -9007199254740991], "unsafe": ["9007199254740992", "-9007199254740992"], "non_integer": [1.25, 1e-2], "other": [null, true, "text"] }"#, ) .unwrap(); assert_eq!(json_value_for_js(value), expected); } }