use async_trait::async_trait;
use cratefield_core::{Database, DbError, Row, Rows, Statement};
use sea_query::Value as SeaValue;
use serde_json::Value as Json;
use worker::D1Database as WorkerD1;
use worker::send::IntoSendFuture;
pub struct D1Database(pub WorkerD1);
#[derive(Debug)]
enum Bind {
Null,
Json(Json),
Bytes(Vec<u8>),
}
fn decide_bind(value: &SeaValue) -> Bind {
if let SeaValue::Bytes(Some(bytes)) = value {
return Bind::Bytes(bytes.to_vec());
}
let json = sea_to_json(value);
if json.is_null() {
Bind::Null
} else {
Bind::Json(json)
}
}
fn bind_statement(
db: &WorkerD1,
stmt: &Statement,
) -> worker::Result<worker::d1::D1PreparedStatement> {
let prepared = db.prepare(&stmt.sql);
if stmt.values.0.is_empty() {
return Ok(prepared);
}
let js_values: Vec<worker::wasm_bindgen::JsValue> = stmt
.values
.0
.iter()
.map(decide_bind)
.map(|bind| match bind {
Bind::Null => Ok(worker::wasm_bindgen::JsValue::NULL),
Bind::Json(json) => worker::d1::serde_wasm_bindgen::to_value(&json),
Bind::Bytes(bytes) => Ok(worker::js_sys::Uint8Array::from(bytes.as_slice()).into()),
})
.collect::<Result<Vec<_>, _>>()
.map_err(worker::Error::SerdeWasmBindgenError)?;
prepared.bind(&js_values)
}
fn sea_to_json(value: &SeaValue) -> serde_json::Value {
match value {
SeaValue::Bool(Some(v)) => serde_json::Value::Bool(*v),
SeaValue::TinyInt(Some(v)) => (*v).into(),
SeaValue::SmallInt(Some(v)) => (*v).into(),
SeaValue::Int(Some(v)) => (*v).into(),
SeaValue::BigInt(Some(v)) => (*v).into(),
SeaValue::TinyUnsigned(Some(v)) => (*v).into(),
SeaValue::SmallUnsigned(Some(v)) => (*v).into(),
SeaValue::Unsigned(Some(v)) => (*v).into(),
SeaValue::BigUnsigned(Some(v)) => (*v).into(),
SeaValue::Float(Some(v)) => f64::from(*v).into(),
SeaValue::Double(Some(v)) => (*v).into(),
SeaValue::String(Some(v)) => v.as_str().into(),
_ => serde_json::Value::Null,
}
}
fn json_to_sea(value: &Json) -> SeaValue {
match value {
Json::Null => SeaValue::String(None),
Json::Bool(v) => SeaValue::Bool(Some(*v)),
Json::Number(n) => {
if let Some(i) = n.as_i64() {
SeaValue::BigInt(Some(i))
} else {
SeaValue::Double(Some(n.as_f64().unwrap_or_default()))
}
}
Json::String(s) => SeaValue::String(Some(Box::new(s.clone()))),
Json::Array(items) => match array_to_bytes(items) {
Some(bytes) => SeaValue::Bytes(Some(Box::new(bytes))),
None => SeaValue::String(Some(Box::from(value.to_string()))),
},
Json::Object(_) => SeaValue::String(Some(Box::from(value.to_string()))),
}
}
fn array_to_bytes(items: &[Json]) -> Option<Vec<u8>> {
items
.iter()
.map(|item| item.as_u64().and_then(|n| u8::try_from(n).ok()))
.collect()
}
fn d1_rows_to_rows(result: &worker::d1::D1Result) -> Result<Rows, DbError> {
let values: Vec<Json> = result.results().map_err(|err| {
log_d1(&err);
DbError::Query(err.to_string())
})?;
let rows = values
.into_iter()
.map(|value| {
let Json::Object(map) = value else {
return Row::new(Vec::new());
};
Row::new(
map.into_iter()
.map(|(column, value)| (column, json_to_sea(&value)))
.collect(),
)
})
.collect();
Ok(Rows::new(rows))
}
fn log_d1(err: &worker::Error) {
worker::console_error!("[d1] {err}");
}
#[async_trait]
impl Database for D1Database {
async fn execute(&self, stmt: &Statement) -> Result<u64, DbError> {
let prepared = bind_statement(&self.0, stmt).map_err(|err| {
log_d1(&err);
DbError::Execute(err.to_string())
})?;
let results = self
.0
.batch(vec![prepared])
.into_send()
.await
.map_err(|err| {
log_d1(&err);
DbError::Execute(err.to_string())
})?;
let changed = results
.first()
.and_then(|result| result.meta().ok().flatten())
.and_then(|meta| meta.changes)
.unwrap_or_default();
Ok(u64::try_from(changed).unwrap_or(u64::MAX))
}
async fn query(&self, stmt: &Statement) -> Result<Rows, DbError> {
let prepared = bind_statement(&self.0, stmt).map_err(|err| {
log_d1(&err);
DbError::Query(err.to_string())
})?;
let result = prepared.all().into_send().await.map_err(|err| {
log_d1(&err);
DbError::Query(err.to_string())
})?;
d1_rows_to_rows(&result)
}
async fn batch_atomic(&self, stmts: &[Statement]) -> Result<(), DbError> {
let mut prepared = Vec::with_capacity(stmts.len());
for stmt in stmts {
prepared.push(bind_statement(&self.0, stmt).map_err(|err| {
log_d1(&err);
DbError::Batch(err.to_string())
})?);
}
self.0.batch(prepared).into_send().await.map_err(|err| {
log_d1(&err);
DbError::Batch(err.to_string())
})?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn hostile_payload() -> Vec<u8> {
vec![0x00, 0x41, 0xFF, 0x80, 0xFE]
}
fn blob_row(payload: &[u8]) -> Json {
Json::Array(payload.iter().map(|byte| Json::from(*byte)).collect())
}
fn decoded(value: &Json) -> Vec<u8> {
match json_to_sea(value) {
SeaValue::Bytes(Some(bytes)) => *bytes,
SeaValue::String(Some(text)) => panic!("stringified, not decoded: {text}"),
other => panic!("unexpected sea value: {other:?}"),
}
}
#[test]
fn bind_sends_every_byte_of_a_blob() {
let payload = hostile_payload();
match decide_bind(&SeaValue::Bytes(Some(Box::new(payload.clone())))) {
Bind::Bytes(bytes) => assert_eq!(bytes, payload),
other => panic!("expected Bind::Bytes, got {other:?}"),
}
}
#[test]
fn bind_maps_bytes_none_to_null() {
assert!(matches!(decide_bind(&SeaValue::Bytes(None)), Bind::Null));
}
#[test]
fn decode_survives_nul_and_invalid_utf8() {
let payload = hostile_payload();
assert_eq!(decoded(&blob_row(&payload)), payload);
}
#[test]
fn decode_empty_blob_to_empty_bytes() {
assert_eq!(decoded(&blob_row(&[])), Vec::<u8>::new());
}
#[test]
fn non_byte_arrays_are_not_blobs() {
for row in [
Json::Array(vec![Json::from(256)]),
Json::Array(vec![Json::from(-1)]),
Json::Array(vec![Json::from(2.5)]),
Json::Array(vec![Json::from(1), Json::Array(vec![Json::from(2)])]),
] {
match json_to_sea(&row) {
SeaValue::String(Some(_)) => {}
other => panic!("{row} decoded as {other:?}"),
}
}
}
#[test]
fn non_blob_binds_are_unchanged() {
for (sea, json) in [
(SeaValue::Bool(Some(true)), Json::Bool(true)),
(SeaValue::Int(Some(7)), Json::from(7)),
(
SeaValue::BigInt(Some(i64::from(i32::MAX))),
Json::from(i64::from(i32::MAX)),
),
(SeaValue::Double(Some(2.5)), Json::from(2.5)),
(
SeaValue::String(Some(Box::new("x".to_owned()))),
Json::String("x".to_owned()),
),
] {
match decide_bind(&sea) {
Bind::Json(value) => assert_eq!(value, json),
other => panic!("{sea:?} bound as {other:?}"),
}
}
}
#[test]
fn non_blob_decodes_are_unchanged() {
assert!(matches!(json_to_sea(&Json::Null), SeaValue::String(None)));
assert!(matches!(
json_to_sea(&Json::Bool(false)),
SeaValue::Bool(Some(false))
));
assert!(matches!(
json_to_sea(&Json::from(7)),
SeaValue::BigInt(Some(7))
));
assert!(matches!(
json_to_sea(&Json::from(2.5)),
SeaValue::Double(Some(f)) if (f - 2.5).abs() < f64::EPSILON
));
assert!(matches!(
json_to_sea(&Json::String("x".to_owned())),
SeaValue::String(Some(_))
));
assert!(matches!(
json_to_sea(&Json::String("[1,2]".to_owned())),
SeaValue::String(Some(_))
));
}
}