sqlar/
compress.rs

1//! Reimplementation of the `sqlar_compress` and `sqlar_uncompress` functions for SQLite.
2//!
3//! Not every sqlite build has these functions enabled, e.g. if Zlib wasn't available at build
4//! time.
5//! By reimplementing them we guarantee their availability.
6//!
7//! SQLite is Public Domain. See https://www.sqlite.org/copyright.html
8//!
9//! SPDX-FileCopyrightText: 2021 Jan-Erik Rediger <janerik@fnordig.de>
10//! SPDX-License-Identifier: CC0-1.0
11
12use rusqlite::functions::FunctionFlags;
13use rusqlite::types::{Value, ValueRef};
14use rusqlite::{Connection, Error, Result};
15
16use flate2::read::ZlibDecoder;
17use flate2::write::ZlibEncoder;
18use flate2::Compression;
19use std::io::prelude::*;
20
21/// Create the functions `rusty_sqlar_compress` and `rusty_sqlar_uncompress` in the database.
22pub fn init(db: &Connection) -> Result<()> {
23    sqlar_compress(db)?;
24    sqlar_uncompress(db)?;
25    Ok(())
26}
27
28/// Implementation of the `sqlar_compress(X)` SQL function.
29///
30/// If the type of X is SQLITE_BLOB, and compressing that blob using
31/// zlib utility function compress() yields a smaller blob, return the
32/// compressed blob. Otherwise, return a copy of X.
33///
34/// SQLar uses the "zlib format" for compressed content.  The zlib format
35/// contains a two-byte identification header and a four-byte checksum at
36/// the end.  This is different from ZIP which uses the raw deflate format.
37///
38/// Future enhancements to SQLar might add support for new compression formats.
39/// If so, those new formats will be identified by alternative headers in the
40/// compressed data.
41fn sqlar_compress(db: &Connection) -> Result<()> {
42    db.create_scalar_function(
43        "rusty_sqlar_compress",
44        1,
45        FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
46        |ctx| {
47            assert_eq!(ctx.len(), 1, "called with unexpected number of arguments");
48
49            let value = ctx.get_raw(0);
50            let value = match value {
51                // Try to compress a blob
52                ValueRef::Blob(blob) => {
53                    let mut enc = ZlibEncoder::new(Vec::new(), Compression::default());
54                    enc.write_all(blob)
55                        .map_err(|_| Error::UserFunctionError("error in compress()".into()))?;
56                    match enc.finish() {
57                        // If it is actually compressed, return the compressed data.
58                        Ok(compressed) if compressed.len() < blob.len() => Value::Blob(compressed),
59                        // Otherwise return a copy of the data
60                        Ok(_) => Value::from(ValueRef::Blob(blob)),
61                        // Or return an error
62                        Err(_) => {
63                            return Err(Error::UserFunctionError("error in compress()".into()))
64                        }
65                    }
66                }
67                // For other types return a copy
68                value => Value::from(value),
69            };
70
71            Ok(value)
72        },
73    )
74}
75
76/// Implementation of the `sqlar_uncompress(X,SZ)` SQL function
77///
78/// Parameter SZ is interpreted as an integer. If it is less than or
79/// equal to zero, then this function returns a copy of X. Or, if
80/// SZ is equal to the size of X when interpreted as a blob, also
81/// return a copy of X. Otherwise, decompress blob X using zlib
82/// utility function uncompress() and return the results (another
83/// blob).
84fn sqlar_uncompress(db: &Connection) -> Result<()> {
85    db.create_scalar_function(
86        "rusty_sqlar_uncompress",
87        2,
88        FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
89        |ctx| {
90            assert_eq!(ctx.len(), 2, "called with unexpected number of arguments");
91
92            let value = ctx.get_raw(0);
93            let size = ctx.get::<i32>(1)?;
94
95            if size <= 0 {
96                return Ok(Value::from(value));
97            }
98
99            let value = match value {
100                // Already uncompressed, return a copy of the data
101                ValueRef::Blob(blob) if blob.len() == size as usize => Value::from(value),
102                // Otherwise, try to decode the blob.
103                ValueRef::Blob(blob) => {
104                    let mut dec = ZlibDecoder::new(blob);
105                    let mut out = Vec::new();
106                    dec.read_to_end(&mut out)
107                        .map_err(|_| Error::UserFunctionError("error in compress()".into()))?;
108
109                    Value::Blob(out)
110                }
111                // For other types return a copy
112                value => Value::from(value),
113            };
114
115            Ok(value)
116        },
117    )
118}