quack_protocol 0.1.0

Rust client SDK for DuckDB's experimental Quack remote protocol
Documentation
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use serde_json::{Map, Number, Value as JsonValue, json};

use crate::errors::{QuackError, Result};
use crate::vector::{Row, TimeUnit, TimestampUnit, Value, decimal_to_string};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BigIntJsonMode {
    String,
    Number,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BytesJsonMode {
    Base64,
    Hex,
    Array,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TaggedJsonMode {
    Default,
    Tagged,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct JsonOptions {
    pub bigint: BigIntJsonMode,
    pub bytes: BytesJsonMode,
    pub decimal: TaggedJsonMode,
    pub date: TaggedJsonMode,
    pub time: TaggedJsonMode,
    pub timestamp: TaggedJsonMode,
}

impl Default for JsonOptions {
    fn default() -> Self {
        Self {
            bigint: BigIntJsonMode::String,
            bytes: BytesJsonMode::Base64,
            decimal: TaggedJsonMode::Default,
            date: TaggedJsonMode::Default,
            time: TaggedJsonMode::Default,
            timestamp: TaggedJsonMode::Default,
        }
    }
}

pub fn to_json_value(value: &Value, options: JsonOptions) -> Result<JsonValue> {
    Ok(match value {
        Value::Null => JsonValue::Null,
        Value::Bool(value) => JsonValue::Bool(*value),
        Value::Int(value) => bigint_to_json(*value as i128, options.bigint)?,
        Value::UInt(value) => unsigned_bigint_to_json(*value as u128, options.bigint)?,
        Value::HugeInt(value) => bigint_to_json(*value, options.bigint)?,
        Value::UHugeInt(value) => unsigned_bigint_to_json(*value, options.bigint)?,
        Value::Float(value) => Number::from_f64(*value as f64)
            .map(JsonValue::Number)
            .unwrap_or(JsonValue::Null),
        Value::Double(value) => Number::from_f64(*value)
            .map(JsonValue::Number)
            .unwrap_or(JsonValue::Null),
        Value::String(value) => JsonValue::String(value.clone()),
        Value::Bytes(value) => bytes_to_json(value, options.bytes),
        Value::Decimal(value) => match options.decimal {
            TaggedJsonMode::Tagged => json!({
                "kind": "decimal",
                "value": value.value.to_string(),
                "width": value.width,
                "scale": value.scale
            }),
            TaggedJsonMode::Default => JsonValue::String(decimal_to_string(value)),
        },
        Value::Date(value) => match options.date {
            TaggedJsonMode::Tagged => json!({ "kind": "date", "days": value.days }),
            TaggedJsonMode::Default => JsonValue::String(date_to_iso(value.days)),
        },
        Value::Time(value) => match options.time {
            TaggedJsonMode::Tagged => json!({
                "kind": "time",
                "unit": match value.unit { TimeUnit::Micros => "micros", TimeUnit::Nanos => "nanos" },
                "value": value.value.to_string()
            }),
            TaggedJsonMode::Default => JsonValue::String(time_to_string(value.value, value.unit)),
        },
        Value::TimeTz(value) => json!({ "kind": "time_tz", "bits": value.bits.to_string() }),
        Value::Timestamp(value) => match options.timestamp {
            TaggedJsonMode::Tagged => {
                let unit = match value.unit {
                    TimestampUnit::Seconds => "seconds",
                    TimestampUnit::Millis => "millis",
                    TimestampUnit::Micros => "micros",
                    TimestampUnit::Nanos => "nanos",
                };
                if value.timezone_utc {
                    json!({ "kind": "timestamp", "unit": unit, "value": value.value.to_string(), "timezone": "utc" })
                } else {
                    json!({ "kind": "timestamp", "unit": unit, "value": value.value.to_string() })
                }
            }
            TaggedJsonMode::Default => {
                JsonValue::String(timestamp_to_iso(value.value, value.unit)?)
            }
        },
        Value::Interval(value) => json!({
            "kind": "interval",
            "months": value.months,
            "days": value.days,
            "micros": value.micros.to_string()
        }),
        Value::List(values) => JsonValue::Array(
            values
                .iter()
                .map(|value| to_json_value(value, options))
                .collect::<Result<Vec<_>>>()?,
        ),
        Value::Struct(row) => {
            let mut object = Map::new();
            for (key, value) in row {
                object.insert(key.clone(), to_json_value(value, options)?);
            }
            JsonValue::Object(object)
        }
    })
}

pub fn to_json_row(row: &Row, options: JsonOptions) -> Result<JsonValue> {
    let mut object = Map::new();
    for (key, value) in row {
        object.insert(key.clone(), to_json_value(value, options)?);
    }
    Ok(JsonValue::Object(object))
}

pub fn to_json_rows(rows: &[Row], options: JsonOptions) -> Result<Vec<JsonValue>> {
    rows.iter().map(|row| to_json_row(row, options)).collect()
}

fn bigint_to_json(value: i128, mode: BigIntJsonMode) -> Result<JsonValue> {
    match mode {
        BigIntJsonMode::String => Ok(JsonValue::String(value.to_string())),
        BigIntJsonMode::Number => i64::try_from(value)
            .ok()
            .map(Number::from)
            .map(JsonValue::Number)
            .ok_or_else(|| {
                QuackError::protocol(format!("bigint value {value} exceeds JSON number range"))
            }),
    }
}

fn unsigned_bigint_to_json(value: u128, mode: BigIntJsonMode) -> Result<JsonValue> {
    match mode {
        BigIntJsonMode::String => Ok(JsonValue::String(value.to_string())),
        BigIntJsonMode::Number => u64::try_from(value)
            .ok()
            .map(|value| JsonValue::Number(Number::from(value)))
            .ok_or_else(|| {
                QuackError::protocol(format!("bigint value {value} exceeds JSON number range"))
            }),
    }
}

fn bytes_to_json(value: &[u8], mode: BytesJsonMode) -> JsonValue {
    match mode {
        BytesJsonMode::Base64 => JsonValue::String(STANDARD.encode(value)),
        BytesJsonMode::Hex => {
            JsonValue::String(value.iter().map(|byte| format!("{byte:02x}")).collect())
        }
        BytesJsonMode::Array => JsonValue::Array(value.iter().map(|byte| json!(byte)).collect()),
    }
}

fn date_to_iso(days: i32) -> String {
    let epoch = chrono::NaiveDate::from_ymd_opt(1970, 1, 1).expect("valid epoch");
    (epoch + chrono::Duration::days(days as i64))
        .format("%Y-%m-%d")
        .to_string()
}

fn time_to_string(value: i64, unit: TimeUnit) -> String {
    let nanos = match unit {
        TimeUnit::Micros => value as i128 * 1_000,
        TimeUnit::Nanos => value as i128,
    };
    let total_seconds = nanos.div_euclid(1_000_000_000);
    let fraction = nanos.rem_euclid(1_000_000_000);
    let hours = total_seconds / 3_600;
    let minutes = (total_seconds % 3_600) / 60;
    let seconds = total_seconds % 60;
    let fraction_text = if fraction == 0 {
        String::new()
    } else {
        format!(".{:09}", fraction)
            .trim_end_matches('0')
            .to_string()
    };
    format!("{hours:02}:{minutes:02}:{seconds:02}{fraction_text}")
}

fn timestamp_to_iso(value: i64, unit: TimestampUnit) -> Result<String> {
    let nanos = match unit {
        TimestampUnit::Seconds => value as i128 * 1_000_000_000,
        TimestampUnit::Millis => value as i128 * 1_000_000,
        TimestampUnit::Micros => value as i128 * 1_000,
        TimestampUnit::Nanos => value as i128,
    };
    let seconds = nanos.div_euclid(1_000_000_000);
    let fraction = nanos.rem_euclid(1_000_000_000);
    let datetime = chrono::DateTime::from_timestamp(seconds as i64, fraction as u32)
        .ok_or_else(|| QuackError::protocol(format!("timestamp value {value} is outside range")))?;
    let fraction_text = if fraction == 0 {
        String::new()
    } else {
        format!(".{:09}", fraction)
            .trim_end_matches('0')
            .to_string()
    };
    Ok(format!(
        "{}{}Z",
        datetime.format("%Y-%m-%dT%H:%M:%S"),
        fraction_text
    ))
}