path: root/src/database.rs

// Copyright © 2020 Jakob L. Kreuze <zerodaysfordays@sdf.org>
//
// This file is part of бирка-тян.
//
// бирка-тян is free software; you can redistribute it and/or modify it
// under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation; either version 3 of the License, or
// (at your option) any later version.
//
// бирка-тян is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with бирка-тян. If not, see <http://www.gnu.org/licenses/>.

use anyhow::Result;
use blake2::{Blake2b, Digest};
use rusqlite::Connection;

use std::convert::AsRef;
use std::io;
use std::path::Path;

/// A connection to the tag database.
#[derive(Debug)]
pub struct TagDatabase(Connection);

/// Tag query with a sense of "state" for pagination.
#[derive(Debug)]
pub struct Query<'a> {
    tags: &'a [&'a str],
    last_id: Option<i64>,
    limit: Option<i64>,
}

impl<'a> Query<'a> {
    /// Construct a new query for images tagged with `tags`.
    ///
    /// When given to the database, this query will return at most `limit`
    /// results, if specified, and will return only images newer than `last_id`,
    /// if specified.
    pub fn new(tags: &'a [&'a str], last_id: Option<i64>, limit: Option<i64>) -> Self {
        Query {
            tags,
            last_id,
            limit,
        }
    }
}

/// The result of an image query in бирка-тян's tag database.
#[derive(Debug)]
pub struct Image {
    pub id: i64,
    pub blake2: String,
    pub filename: String,
    pub orig_dir: String,
}

/// Return a base64-encoded BLAKE2b hash identifying the file at `path`.
pub fn hash_file(path: &Path) -> Result<String> {
    let mut file = std::fs::File::open(path)?;
    let mut hasher = Blake2b::new();
    io::copy(&mut file, &mut hasher)?;
    Ok(base64::encode(&hasher.result()))
}

/// Split the path into a tuple containing the parent directory and the filename.
pub fn split_path(path: &Path) -> Result<(String, String)> {
    if path.is_dir() {
        bail!("`path` does not name a file.");
    }

    let file_name = path
        .file_name()
        .and_then(|os| os.to_str())
        .ok_or(anyhow!("Couldn't parse file name."))?;

    // Convert `path` to a string and strip off the `file_name` we've just
    // obtained to yield the name of the parent directory.
    let path = path.to_str().ok_or(anyhow!("Couldn't parse path."))?;
    let split_index = path.len() - file_name.len();
    let parent_directory = &path[0..split_index];

    Ok((file_name.to_string(), parent_directory.to_string()))
}

impl TagDatabase {
    /// Instantiate a new `TagDatabase` whose backing database is at `path`.
    pub fn new<T: AsRef<Path>>(path: &T) -> Result<Self> {
        let tb = TagDatabase(Connection::open(path)?);
        tb.initialize_tables()?;
        Ok(tb)
    }

    /// Create the table structure for the tag database.
    fn initialize_tables(&self) -> Result<()> {
        self.0.execute(
            "CREATE TABLE IF NOT EXISTS images (
                 id       INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
                 blake2   TEXT NOT NULL,
                 filename TEXT NOT NULL,
                 orig_dir TEXT NOT NULL
             )",
            params![],
        )?;
        self.0.execute(
            "CREATE TABLE IF NOT EXISTS tags (
                 id   INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
                 name TEXT NOT NULL
             )",
            params![],
        )?;
        self.0.execute(
            "CREATE TABLE IF NOT EXISTS mapping (
                 image  INTEGER NOT NULL,
                 tag    INTEGER NOT NULL
             )",
            params![],
        )?;
        Ok(())
    }

    /// Insert `tag` into the database, if it does not already exist.
    pub fn add_tag(&self, tag: &str) -> Result<()> {
        let exists = self
            .0
            .prepare("SELECT * FROM tags WHERE name = ?")?
            .exists(params![tag])?;
        if !exists {
            self.0
                .execute("INSERT INTO tags (name) VALUES(?)", params![tag])?;
        }
        Ok(())
    }

    /// Return the internal id for `tag`.
    pub fn tag_id(&self, tag: &str) -> Result<i64> {
        Ok(self
            .0
            .prepare("SELECT * FROM tags WHERE name = ?")?
            .query_row(params![tag], |row| Ok(row.get::<_, i64>(0)))??)
    }

    /// Insert the image at `path` into the data.
    pub fn import_image(&self, path: &Path, hash: &str) -> Result<i64> {
        let (file_name, parent_directory) = split_path(path)?;
        Ok(self
            .0
            .prepare("INSERT INTO images (blake2, filename, orig_dir) VALUES(?,?,?)")?
            .insert(params![hash, file_name, parent_directory])?)
    }

    /// Return the internal id for the image at `path`.
    pub fn image_id(&self, path: &Path) -> Result<i64> {
        let (file_name, path) = split_path(path)?;
        Ok(self
            .0
            .prepare("SELECT id FROM images WHERE orig_dir = ? AND filename = ?")?
            .query_row(params![path, file_name], |row| Ok(row.get::<_, i64>(0)))??)
    }

    /// Return the `Image` struct for the image with `id`.
    pub fn image_by_id(&self, id: i64) -> Result<Image> {
        Ok(self
            .0
            .prepare(
                "SELECT id, blake2, filename, orig_dir
                 FROM images
                 WHERE id = ?",
            )?
            .query_row(params![id], |row| {
                Ok(Image {
                    id: row.get(0)?,
                    blake2: row.get(1)?,
                    filename: row.get(2)?,
                    orig_dir: row.get(3)?,
                })
            })?)
    }

    /// Associate `tag` with the image specified by `image_id`.
    pub fn tag_image(&self, image_id: i64, tag: &str) -> Result<()> {
        self.0.execute(
            "INSERT INTO mapping (image, tag) VALUES(?, ?)",
            params![image_id, self.tag_id(tag)?],
        )?;
        Ok(())
    }

    /// Disassociate `tag` with the image specified by `image_id`.
    pub fn untag_image(&self, image_id: i64, tag: &str) -> Result<()> {
        self.0.execute(
            "DELETE FROM mapping WHERE image = ? AND tag = ?",
            params![image_id, self.tag_id(tag)?],
        )?;
        Ok(())
    }

    /// Return the tags for the image specified by `image_id`.
    pub fn tags_for_image(&self, image_id: i64) -> Result<Vec<String>> {
        Ok(self
            .0
            .prepare(
                "SELECT tags.name
                 FROM mapping
                 INNER JOIN tags ON mapping.tag = tags.id
                 WHERE mapping.image = ?;",
            )?
            .query_map(params![image_id], |row| row.get(0))?
            .filter_map(|tag| tag.ok())
            .collect())
    }

    /// Return the intersection of the sets of images matching each of `tags`.
    pub fn query(&self, q: Query) -> Result<Vec<Image>> {
        let ids: Vec<i64> = q
            .tags
            .iter()
            .filter_map(|tag| self.tag_id(tag).ok())
            .collect();
        let mut tags = Vec::from(q.tags);
        tags.dedup();

        // Query contained a tag that's not in the database.
        if ids.len() != tags.len() {
            return Ok(vec![]);
        }

        let union = if tags.is_empty() {
            None
        } else {
            Some(
                (1..ids.len() + 1)
                    .zip(ids)
                    .map(|pair| {
                        let (n, id) = pair;
                        format!(
                            "INNER JOIN mapping as m{n}
                             ON images.id = m{n}.image
                             AND {id} = m{n}.tag",
                            n = n,
                            id = id
                        )
                    })
                    .collect::<Vec<String>>()
                    .join("\n"),
            )
        };

        let query = format!(
            "SELECT id, blake2, filename, orig_dir
            FROM images
            {union}
            {since}
            GROUP BY id
            ORDER BY id DESC
            {limit}",
            since = if let Some(n) = q.last_id {
                format!("WHERE id < {}", n)
            } else {
                String::from("")
            },
            limit = if let Some(n) = q.limit {
                format!("LIMIT {}", n)
            } else {
                String::from("")
            },
            union = union.or(Some("".into())).unwrap(),
        );

        Ok(self
            .0
            .prepare(&query)?
            .query_map(
                params![],
                |row| -> rusqlite::Result<(i64, String, String, String)> {
                    Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?))
                },
            )?
            .filter_map(|res| {
                res.map(|pair| {
                    let (id, blake2, filename, orig_dir) = pair;
                    Image {
                        id,
                        blake2,
                        filename,
                        orig_dir,
                    }
                })
                .ok()
            })
            .collect())
    }

    /// Return the intersection of the sets of images matching each of `tags`.
    pub fn images_by_tags(&self, tags: &[&str]) -> Result<Vec<Image>> {
        self.query(Query {
            tags,
            last_id: None,
            limit: None,
        })
    }

    /// Return the `n` greatest tags, ordered by number of associated images.
    pub fn top_tags(&self, n: i64) -> Result<Vec<(i64, String)>> {
        Ok(self
            .0
            .prepare(
                "SELECT tags.name, COUNT(mapping.tag) AS tag_count
                 FROM mapping
                 INNER JOIN tags ON tags.id = mapping.tag
                 GROUP BY tag
                 ORDER BY tag_count DESC
                 LIMIT ?",
            )?
            .query_map(params![n], |row| -> rusqlite::Result<(i64, String)> {
                Ok((row.get(1)?, row.get(0)?))
            })?
            .filter_map(|s| s.ok())
            .collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    impl TagDatabase {
        fn new_mem() -> Result<Self> {
            let tb = TagDatabase(Connection::open_in_memory()?);
            tb.initialize_tables()?;
            Ok(tb)
        }
    }

    #[test]
    fn tables_exist() {
        let tb = TagDatabase::new_mem().unwrap();
        for table in ["images", "tags", "mapping"].iter() {
            let exists =
                tb.0.prepare("SELECT name FROM sqlite_master WHERE type='table' AND name=?;")
                    .unwrap()
                    .exists(params![table])
                    .unwrap();
            assert!(exists);
        }
    }

    #[test]
    fn insert_file() {
        let tb = TagDatabase::new_mem().unwrap();
        let id = tb
            .import_image(Path::new("/fake/path"), "fakehash")
            .unwrap();
        let exists =
            tb.0.prepare("SELECT filename, blake2 FROM images WHERE id = ?;")
                .unwrap()
                .exists(params![id])
                .unwrap();
        assert!(exists);
        let exists =
            tb.0.prepare("SELECT filename, blake2 FROM images WHERE id = ?;")
                .unwrap()
                .exists(params![id + 1])
                .unwrap();
        assert!(!exists);
    }

    #[test]
    fn add_tag() {
        let tb = TagDatabase::new_mem().unwrap();
        tb.add_tag("test").unwrap();
        tb.tag_id("test").unwrap();
    }

    #[test]
    fn tag_image() {
        let tb = TagDatabase::new_mem().unwrap();
        let id = tb
            .import_image(Path::new("/fake/path"), "fakehash")
            .unwrap();
        tb.add_tag("test").unwrap();
        tb.tag_image(id, "test").unwrap();
        let tags = tb.tags_for_image(id).unwrap();
        assert_eq!(tags.len(), 1);
        assert!(tags.contains(&String::from("test")));
    }

    #[test]
    fn query_by_tag() {
        let tb = TagDatabase::new_mem().unwrap();
        let ids: Vec<i64> = (1..4)
            .map(|i| {
                let i = i.to_string();
                tb.import_image(Path::new(&i), &i).unwrap()
            })
            .collect();
        tb.add_tag("1").unwrap();
        tb.add_tag("2").unwrap();
        tb.tag_image(ids[0], "1").unwrap();
        tb.tag_image(ids[1], "2").unwrap();
        tb.tag_image(ids[2], "1").unwrap();
        tb.tag_image(ids[2], "2").unwrap();
        assert_eq!(tb.images_by_tags(&["1"]).unwrap().len(), 2);
        assert_eq!(tb.images_by_tags(&["2"]).unwrap().len(), 2);
        assert_eq!(tb.images_by_tags(&["1", "2"]).unwrap().len(), 1);
    }

    #[test]
    fn top_tags() {
        let tb = TagDatabase::new_mem().unwrap();
        let ids: Vec<i64> = (1..4)
            .map(|i| {
                let i = i.to_string();
                tb.import_image(Path::new(&i), &i).unwrap()
            })
            .collect();
        tb.add_tag("1").unwrap();
        tb.add_tag("2").unwrap();
        tb.tag_image(ids[0], "1").unwrap();
        tb.tag_image(ids[1], "2").unwrap();
        tb.tag_image(ids[2], "1").unwrap();
        tb.tag_image(ids[2], "2").unwrap();
        assert_eq!(tb.top_tags(2).unwrap().len(), 2);
        assert_eq!(tb.top_tags(2).unwrap()[0].0, 2);
        assert_eq!(tb.top_tags(2).unwrap()[1].0, 2);
    }
}