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;
#[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)
}
}
#[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()
);
}
}