sqlite_compressions/bzip2.rs
1use std::io::{Read, Write};
2
3use bzip2::read::BzDecoder;
4use bzip2::write::BzEncoder;
5use bzip2::Compression;
6use rusqlite::Error::UserFunctionError;
7
8use crate::common::{register_compression, Encoder};
9use crate::rusqlite::{Connection, Result};
10
11/// Register the `bzip2` SQL functions with the given `SQLite` connection.
12/// The function takes a single argument and returns the [bzip2 compression](https://en.wikipedia.org/wiki/Bzip2) (blob) of that argument.
13/// The argument can be either a string or a blob.
14/// If the argument is `NULL`, the result is `NULL`.
15///
16/// # Example
17///
18/// ```
19/// # use sqlite_compressions::rusqlite::{Connection, Result};
20/// # use sqlite_compressions::register_bzip2_functions;
21/// # fn main() -> Result<()> {
22/// let db = Connection::open_in_memory()?;
23/// register_bzip2_functions(&db)?;
24/// let result: Vec<u8> = db.query_row("SELECT bzip2('hello')", [], |r| r.get(0))?;
25/// let expected = b"\x42\x5a\x68\x36\x31\x41\x59\x26\x53\x59\x19\x31\x65\x3d\x00\x00\x00\x81\x00\x02\x44\xa0\x00\x21\x9a\x68\x33\x4d\x07\x33\x8b\xb9\x22\x9c\x28\x48\x0c\x98\xb2\x9e\x80";
26/// assert_eq!(result, expected);
27/// let result: String = db.query_row("SELECT CAST(bzip2_decode(bzip2('world')) AS TEXT)", [], |r| r.get(0))?;
28/// let expected = "world";
29/// assert_eq!(result, expected);
30/// let result: bool = db.query_row("SELECT bzip2_test(bzip2('world'))", [], |r| r.get(0))?;
31/// let expected = true;
32/// assert_eq!(result, expected);
33/// # Ok(())
34/// # }
35/// ```
36pub fn register_bzip2_functions(conn: &Connection) -> Result<()> {
37 register_compression::<Bzip2Encoder>(conn)
38}
39
40pub struct Bzip2Encoder;
41
42impl Encoder for Bzip2Encoder {
43 fn enc_name() -> &'static str {
44 "bzip2"
45 }
46 fn dec_name() -> &'static str {
47 "bzip2_decode"
48 }
49 fn test_name() -> &'static str {
50 "bzip2_test"
51 }
52
53 fn encode(data: &[u8], quality: Option<u32>) -> Result<Vec<u8>> {
54 let quality = if let Some(param) = quality {
55 if param < Compression::fast().level() || param > Compression::best().level() {
56 return Err(UserFunctionError(
57 format!(
58 "The optional second argument to bzip2() must be between {} and {}",
59 Compression::fast().level(),
60 Compression::best().level()
61 )
62 .into(),
63 ));
64 }
65 Compression::new(param)
66 } else {
67 Compression::default()
68 };
69
70 let mut encoder = BzEncoder::new(Vec::new(), quality);
71 encoder
72 .write_all(data)
73 .map_err(|e| UserFunctionError(e.into()))?;
74 encoder.finish().map_err(|e| UserFunctionError(e.into()))
75 }
76
77 fn decode(data: &[u8]) -> Result<Vec<u8>> {
78 let mut decompressed = Vec::new();
79 BzDecoder::new(data)
80 .read_to_end(&mut decompressed)
81 .map_err(|e| UserFunctionError(e.into()))?;
82 Ok(decompressed)
83 }
84
85 fn test(data: &[u8]) -> bool {
86 // reuse the same buffer when decompressing
87 // ideally we should use some null buffer, but bzip2 doesn't seem to support that
88 // note that buffer size does affect performance and depend on the input data size
89 let mut buffer = [0u8; 1024];
90 let mut decoder = BzDecoder::new(data);
91 while let Ok(len) = decoder.read(&mut buffer) {
92 if len == 0 {
93 return true;
94 }
95 }
96 false
97 }
98}