dbx/crates/dbx-core/src/db/mod.rs

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<PoolCheckoutStage> {
        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::<u64>() {
                if (1..=300).contains(&v) {
                    return Duration::from_secs(v);
                }
            }
        }
    }
    fallback
}

pub async fn with_connection_timeout<T, F>(label: &str, timeout: Duration, future: F) -> Result<T, String>
where
    F: Future<Output = Result<T, String>>,
{
    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);
    }
}