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 rusqlite::{params, Connection};
use std::path::Path;

/// A connection to бирка-тян's 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 TagDatabase {
    /// Instantiate a new `TagDatabase` whose backing database is at `path`.
    pub fn new(path: &str) -> 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)))??)
    }

    /// Split the path into a tuple containing the parent directory and the filename.
    fn split_path(&self, 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()))
    }

    /// Insert the image at `path` into the data.
    pub fn import_image(&self, path: &Path, hash: &str) -> Result<i64> {
        let (file_name, parent_directory) = self.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) = self.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)))??)
    }

    /// 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 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 = if let (Some(id), Some(n)) = (q.last_id, q.limit) {
            format!(
                "SELECT id, blake2, filename, orig_dir
                 FROM images
                 {union}
                 WHERE id > {last_id}
                 GROUP BY id
                 ORDER BY id DESC
                 LIMIT {limit}",
                last_id = id,
                union = union.or(Some("".into())).unwrap(),
                limit = n
            )
        } else {
            format!(
                "SELECT id, blake2, filename, orig_dir
                 FROM images
                 {union}
                 GROUP BY id
                 ORDER BY id DESC",
                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,
        })
    }
}

/// 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,
}

#[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)
        }

        /// 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, COUNT(mapping.tag) as tag_count
                     FROM mapping
                     INNER JOIN images ON images.id = ?
                     INNER JOIN tags
                     GROUP BY tags.name;",
                )?
                .query_map(params![image_id], |row| row.get(0))?
                .filter_map(|tag| tag.ok())
                .collect())
        }
    }

    #[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);
    }
}