sqlite_compressions/
gzip.rs

1use std::io::{Read, Write};
2
3use flate2::read::GzDecoder;
4use flate2::write::GzEncoder;
5use flate2::Compression;
6use rusqlite::Error::UserFunctionError;
7
8use crate::common::{register_compression, Encoder};
9use crate::rusqlite::{Connection, Result};
10
11/// Register the `gzip` SQL functions with the given `SQLite` connection.
12/// The function takes a single argument and returns the [GZIP compression](https://en.wikipedia.org/wiki/Gzip) (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_gzip_functions;
21/// # fn main() -> Result<()> {
22/// let db = Connection::open_in_memory()?;
23/// register_gzip_functions(&db)?;
24/// let result: Vec<u8> = db.query_row("SELECT gzip('hello')", [], |r| r.get(0))?;
25/// let expected = b"\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\xff\xcb\x48\xcd\xc9\xc9\x07\x00\x86\xa6\x10\x36\x05\x00\x00\x00";
26/// assert_eq!(result, expected);
27/// let result: String = db.query_row("SELECT CAST(gzip_decode(gzip('world')) AS TEXT)", [], |r| r.get(0))?;
28/// let expected = "world";
29/// assert_eq!(result, expected);
30/// let result: bool = db.query_row("SELECT gzip_test(gzip('world'))", [], |r| r.get(0))?;
31/// let expected = true;
32/// assert_eq!(result, expected);
33/// # Ok(())
34/// # }
35/// ```
36pub fn register_gzip_functions(conn: &Connection) -> Result<()> {
37    register_compression::<GzipEncoder>(conn)
38}
39
40pub struct GzipEncoder;
41
42impl Encoder for GzipEncoder {
43    fn enc_name() -> &'static str {
44        "gzip"
45    }
46    fn dec_name() -> &'static str {
47        "gzip_decode"
48    }
49    fn test_name() -> &'static str {
50        "gzip_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 > 9 {
56                return Err(UserFunctionError(
57                    "The optional second argument to gzip() must be between 0 and 9".into(),
58                ));
59            }
60            Compression::new(param)
61        } else {
62            Compression::default()
63        };
64
65        let mut encoder = GzEncoder::new(Vec::new(), quality);
66        encoder
67            .write_all(data)
68            .map_err(|e| UserFunctionError(e.into()))?;
69        encoder.finish().map_err(|e| UserFunctionError(e.into()))
70    }
71
72    fn decode(data: &[u8]) -> Result<Vec<u8>> {
73        let mut decompressed = Vec::new();
74        GzDecoder::new(data)
75            .read_to_end(&mut decompressed)
76            .map_err(|e| UserFunctionError(e.into()))?;
77        Ok(decompressed)
78    }
79
80    fn test(data: &[u8]) -> bool {
81        // reuse the same buffer when decompressing
82        // ideally we should use some null buffer, but flate2 doesn't seem to support that
83        // note that buffer size does affect performance and depend on the input data size
84        let mut buffer = [0u8; 1024];
85        let mut decoder = GzDecoder::new(data);
86        while let Ok(len) = decoder.read(&mut buffer) {
87            if len == 0 {
88                return true;
89            }
90        }
91        false
92    }
93}