use rusqlite::functions::FunctionFlags;
use rusqlite::types::{Value, ValueRef};
use rusqlite::{Connection, Error, Result};
use flate2::read::ZlibDecoder;
use flate2::write::ZlibEncoder;
use flate2::Compression;
use std::io::prelude::*;
pub fn init(db: &Connection) -> Result<()> {
sqlar_compress(db)?;
sqlar_uncompress(db)?;
Ok(())
}
fn sqlar_compress(db: &Connection) -> Result<()> {
db.create_scalar_function(
"rusty_sqlar_compress",
1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|ctx| {
assert_eq!(ctx.len(), 1, "called with unexpected number of arguments");
let value = ctx.get_raw(0);
let value = match value {
ValueRef::Blob(blob) => {
let mut enc = ZlibEncoder::new(Vec::new(), Compression::default());
enc.write_all(blob)
.map_err(|_| Error::UserFunctionError("error in compress()".into()))?;
match enc.finish() {
Ok(compressed) if compressed.len() < blob.len() => Value::Blob(compressed),
Ok(_) => Value::from(ValueRef::Blob(blob)),
Err(_) => {
return Err(Error::UserFunctionError("error in compress()".into()))
}
}
}
value => Value::from(value),
};
Ok(value)
},
)
}
fn sqlar_uncompress(db: &Connection) -> Result<()> {
db.create_scalar_function(
"rusty_sqlar_uncompress",
2,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|ctx| {
assert_eq!(ctx.len(), 2, "called with unexpected number of arguments");
let value = ctx.get_raw(0);
let size = ctx.get::<i32>(1)?;
if size <= 0 {
return Ok(Value::from(value));
}
let value = match value {
ValueRef::Blob(blob) if blob.len() == size as usize => Value::from(value),
ValueRef::Blob(blob) => {
let mut dec = ZlibDecoder::new(blob);
let mut out = Vec::new();
dec.read_to_end(&mut out)
.map_err(|_| Error::UserFunctionError("error in compress()".into()))?;
Value::Blob(out)
}
value => Value::from(value),
};
Ok(value)
},
)
}