saddle-db 0.1.1

Saddle managed asynchronous database access and transactions
Documentation
use std::fmt;

use sqlx::{MySql, query::Query};

use crate::{Result, error::invalid_statement};

pub const MAX_OPERATION_BYTES: usize = 128;
pub const MAX_SQL_BYTES: usize = 65_536;
pub const MAX_PARAMETERS: usize = 256;
pub const MAX_PARAMETER_BYTES: usize = 1_048_576;
pub const MAX_PARAMETERS_TOTAL_BYTES: usize = 2_097_152;

/// A scalar parameter accepted by the controlled V1 database API.
///
/// SQL parameters are never included in framework logs or errors.
#[derive(Clone, PartialEq)]
#[non_exhaustive]
pub enum DbValue {
    Null,
    Bool(bool),
    I64(i64),
    U64(u64),
    F64(f64),
    String(String),
    Bytes(Vec<u8>),
}

impl fmt::Debug for DbValue {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::Null => "Null",
            Self::Bool(_) => "Bool(<redacted>)",
            Self::I64(_) => "I64(<redacted>)",
            Self::U64(_) => "U64(<redacted>)",
            Self::F64(_) => "F64(<redacted>)",
            Self::String(_) => "String(<redacted>)",
            Self::Bytes(_) => "Bytes(<redacted>)",
        })
    }
}

impl From<bool> for DbValue {
    fn from(value: bool) -> Self {
        Self::Bool(value)
    }
}
impl From<i32> for DbValue {
    fn from(value: i32) -> Self {
        Self::I64(i64::from(value))
    }
}
impl From<i64> for DbValue {
    fn from(value: i64) -> Self {
        Self::I64(value)
    }
}
impl From<u32> for DbValue {
    fn from(value: u32) -> Self {
        Self::U64(u64::from(value))
    }
}
impl From<u64> for DbValue {
    fn from(value: u64) -> Self {
        Self::U64(value)
    }
}
impl From<f64> for DbValue {
    fn from(value: f64) -> Self {
        Self::F64(value)
    }
}
impl From<String> for DbValue {
    fn from(value: String) -> Self {
        Self::String(value)
    }
}
impl From<&str> for DbValue {
    fn from(value: &str) -> Self {
        Self::String(value.to_owned())
    }
}
impl From<Vec<u8>> for DbValue {
    fn from(value: Vec<u8>) -> Self {
        Self::Bytes(value)
    }
}

/// A parameterized SQL statement and its operation name.
///
/// `operation` is a stable, low-cardinality name used for Trace records. It
/// must describe the operation (for example `orders.insert`), not contain SQL.
#[derive(Clone, PartialEq)]
pub struct Statement {
    operation: String,
    sql: String,
    parameters: Vec<DbValue>,
}

impl Statement {
    pub fn new(operation: impl AsRef<str>, sql: impl AsRef<str>) -> Result<Self> {
        let operation = operation.as_ref();
        let sql = sql.as_ref();
        validate_operation(operation)?;
        if sql.trim().is_empty() || sql.len() > MAX_SQL_BYTES {
            return Err(invalid_statement(
                "SQL must be non-empty and within the V1 size limit",
            ));
        }
        Ok(Self {
            operation: operation.to_owned(),
            sql: sql.to_owned(),
            parameters: Vec::new(),
        })
    }

    pub fn bind(mut self, value: impl Into<DbValue>) -> Result<Self> {
        let value = value.into();
        self.ensure_parameter_fits(&value)?;
        self.parameters.push(value);
        Ok(self)
    }

    pub fn bind_null(self) -> Result<Self> {
        self.bind(DbValue::Null)
    }

    pub fn operation(&self) -> &str {
        &self.operation
    }
    pub fn parameter_count(&self) -> usize {
        self.parameters.len()
    }

    pub(crate) fn validate(&self) -> Result<()> {
        validate_operation(&self.operation)?;
        if self.sql.trim().is_empty() || self.sql.len() > MAX_SQL_BYTES {
            return Err(invalid_statement(
                "SQL must be non-empty and within the V1 size limit",
            ));
        }
        if self.parameters.len() > MAX_PARAMETERS {
            return Err(invalid_statement("statement has too many parameters"));
        }
        let mut total_bytes = 0_usize;
        for parameter in &self.parameters {
            let bytes = parameter.encoded_size();
            if bytes > MAX_PARAMETER_BYTES {
                return Err(invalid_statement(
                    "statement parameter exceeds the V1 byte limit",
                ));
            }
            total_bytes = total_bytes.saturating_add(bytes);
            if total_bytes > MAX_PARAMETERS_TOTAL_BYTES {
                return Err(invalid_statement(
                    "statement parameters exceed the V1 total byte limit",
                ));
            }
        }
        Ok(())
    }

    fn ensure_parameter_fits(&self, value: &DbValue) -> Result<()> {
        if self.parameters.len() == MAX_PARAMETERS {
            return Err(invalid_statement("statement has too many parameters"));
        }
        let bytes = value.encoded_size();
        if bytes > MAX_PARAMETER_BYTES {
            return Err(invalid_statement(
                "statement parameter exceeds the V1 byte limit",
            ));
        }
        let current_bytes = self
            .parameters
            .iter()
            .map(DbValue::encoded_size)
            .fold(0_usize, usize::saturating_add);
        if current_bytes.saturating_add(bytes) > MAX_PARAMETERS_TOTAL_BYTES {
            return Err(invalid_statement(
                "statement parameters exceed the V1 total byte limit",
            ));
        }
        Ok(())
    }

    pub(crate) fn query(&self) -> Query<'_, MySql, sqlx::mysql::MySqlArguments> {
        let mut query = sqlx::query(&self.sql);
        for value in &self.parameters {
            query = match value {
                DbValue::Null => query.bind(Option::<String>::None),
                DbValue::Bool(value) => query.bind(*value),
                DbValue::I64(value) => query.bind(*value),
                DbValue::U64(value) => query.bind(*value),
                DbValue::F64(value) => query.bind(*value),
                DbValue::String(value) => query.bind(value.as_str()),
                DbValue::Bytes(value) => query.bind(value.as_slice()),
            };
        }
        query
    }
}

impl DbValue {
    fn encoded_size(&self) -> usize {
        match self {
            Self::Null => 0,
            Self::Bool(_) => 1,
            Self::I64(_) | Self::U64(_) | Self::F64(_) => 8,
            Self::String(value) => value.len(),
            Self::Bytes(value) => value.len(),
        }
    }
}

pub(crate) fn validate_operation(operation: &str) -> Result<()> {
    if operation.is_empty()
        || operation.len() > MAX_OPERATION_BYTES
        || !operation
            .bytes()
            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-'))
    {
        return Err(invalid_statement(
            "operation must be a bounded stable identifier",
        ));
    }
    Ok(())
}

impl fmt::Debug for Statement {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("Statement")
            .field("operation", &self.operation)
            .field("sql_bytes", &self.sql.len())
            .field("parameter_count", &self.parameters.len())
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn debug_and_validation_do_not_expose_sql_or_parameters() {
        let statement = Statement::new("orders.insert", "INSERT secret")
            .unwrap()
            .bind("password")
            .unwrap();
        statement.validate().unwrap();
        let output = format!("{statement:?}");
        assert!(!output.contains("secret"));
        assert!(!output.contains("password"));
        assert!(!format!("{:?}", DbValue::from("password")).contains("password"));
        assert!(Statement::new("raw SQL", "SELECT 1").is_err());
        assert!(Statement::new("orders.query", "x".repeat(MAX_SQL_BYTES + 1)).is_err());
        assert!(
            Statement::new("orders.insert", "INSERT")
                .unwrap()
                .bind(vec![0_u8; MAX_PARAMETER_BYTES + 1])
                .is_err()
        );
        let mut statement = Statement::new("orders.insert", "INSERT").unwrap();
        for _ in 0..MAX_PARAMETERS {
            statement = statement.bind_null().unwrap();
        }
        assert!(statement.bind_null().is_err());
        assert!(
            Statement::new("orders.insert", "INSERT")
                .unwrap()
                .bind(vec![0_u8; MAX_PARAMETER_BYTES])
                .unwrap()
                .bind(vec![0_u8; MAX_PARAMETER_BYTES])
                .unwrap()
                .bind(1_u64)
                .is_err()
        );
    }
}