sqlite_compressions/
brotli.rs

1use std::io::{Read, Write};
2
3use brotli::{CompressorWriter, Decompressor};
4use rusqlite::Error::UserFunctionError;
5
6use crate::common::{register_compression, Encoder};
7use crate::rusqlite::{Connection, Result};
8
9/// Register the `brotli` SQL function with the given `SQLite` connection.
10/// The function takes a single argument and returns the [Brotli compression](https://en.wikipedia.org/wiki/Brotli) (blob) of that argument.
11/// The argument can be either a string or a blob.
12/// If the argument is `NULL`, the result is `NULL`.
13///
14/// # Example
15///
16/// ```
17/// # use sqlite_compressions::rusqlite::{Connection, Result};
18/// # use sqlite_compressions::register_brotli_functions;
19/// # fn main() -> Result<()> {
20/// let db = Connection::open_in_memory()?;
21/// register_brotli_functions(&db)?;
22/// let result: Vec<u8> = db.query_row("SELECT brotli('hello')", [], |r| r.get(0))?;
23/// let expected = b"\x0b\x02\x80\x68\x65\x6c\x6c\x6f\x03";
24/// assert_eq!(result, expected);
25/// let result: String = db.query_row("SELECT CAST(brotli_decode(brotli('world')) AS TEXT)", [], |r| r.get(0))?;
26/// let expected = "world";
27/// assert_eq!(result, expected);
28/// let result: bool = db.query_row("SELECT brotli_test(brotli('world'))", [], |r| r.get(0))?;
29/// let expected = true;
30/// assert_eq!(result, expected);
31/// # Ok(())
32/// # }
33/// ```
34pub fn register_brotli_functions(conn: &Connection) -> Result<()> {
35    register_compression::<BrotliEncoder>(conn)
36}
37
38pub struct BrotliEncoder;
39
40impl Encoder for BrotliEncoder {
41    fn enc_name() -> &'static str {
42        "brotli"
43    }
44    fn dec_name() -> &'static str {
45        "brotli_decode"
46    }
47    fn test_name() -> &'static str {
48        "brotli_test"
49    }
50
51    fn encode(data: &[u8], quality: Option<u32>) -> Result<Vec<u8>> {
52        let mut encoder = CompressorWriter::new(Vec::new(), 4 * 1024, quality.unwrap_or(11), 22);
53        encoder
54            .write_all(data)
55            .map_err(|e| UserFunctionError(e.into()))?;
56        Ok(encoder.into_inner())
57    }
58
59    fn decode(data: &[u8]) -> Result<Vec<u8>> {
60        let mut decompressed = Vec::new();
61        Decompressor::new(data, 4 * 1024)
62            .read_to_end(&mut decompressed)
63            .map_err(|e| UserFunctionError(e.into()))?;
64        Ok(decompressed)
65    }
66
67    fn test(data: &[u8]) -> bool {
68        // reuse the same buffer when decompressing
69        // ideally we should use some null buffer, but flate2 doesn't seem to support that
70        // note that buffer size does affect performance and depend on the input data size
71        let mut buffer = [0u8; 4 * 1024];
72        let mut decoder = Decompressor::new(data, 4 * 1024);
73        while let Ok(len) = decoder.read(&mut buffer) {
74            if len == 0 {
75                return true;
76            }
77        }
78        false
79    }
80}