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::Connection;

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

/// Type of constraint on a column.
#[derive(Debug, PartialEq)]
pub enum Atom {
    Is(String),
    IsNot(String),
}

/// One of the addressable columns of the 'files' table.
#[derive(Debug, PartialEq)]
pub enum Field {
    Tag(Atom),
    Filename(Atom),
}

use Atom::{Is, IsNot};
use Field::{Filename, Tag};

/// A logical conjunction of `Field` contraints.
///
/// This is the general interface for queries to the tag database.
pub type Query = Vec<Field>;

/// Parse `query_string` into a `Query` object usable with the tag database.
pub fn parse_query<T: AsRef<str>>(query_string: T) -> Result<Query> {
    Ok(query_string
        .as_ref()
        .split_whitespace()
        .map(|part| {
            if part.starts_with("-") {
                if part.len() >= 9 && &part[1..10] == "filename:" {
                    Filename(IsNot(String::from(&part[10..])))
                } else {
                    Tag(IsNot(String::from(&part[1..])))
                }
            } else {
                if part.len() >= 9 && &part[..9] == "filename:" {
                    Filename(Is(String::from(&part[9..])))
                } else {
                    Tag(Is(String::from(part)))
                }
            }
        })
        .collect())
}

/// State object representing a point in a paginated query.
pub struct Keyset {
    pub last_id: i64,
    pub how_many: i64,
}

/// Information about an file in the database.
///
/// The result of any queries made to the tag database. `id` is the canonical
/// representation of the file, `path` is the path of the file in the
/// filesystem, and `tags` is a vector containing all of the strings the file is
/// tagged with.
pub struct File {
    pub id: i64,
    pub path: String,
    pub tags: Vec<String>,
}

/// Public interface to the tag database.
///
/// Essentially, a wrapper around `rusqlite::Connection` that allows only
/// queries which are sensical given the purpose of the tag databse.
pub struct TagDatabase(Connection);

impl TagDatabase {
    /// Return a `TagDatabase` connected to `path`.
    ///
    /// This constructor will create the file if it does not exist.
    pub fn new<T: AsRef<Path>>(path: T) -> Result<TagDatabase> {
        let tb = TagDatabase(Connection::open(path)?);
        tb.initialize_tables()?;
        Ok(tb)
    }

    /// Initialize the tag database's schema.
    fn initialize_tables(&self) -> Result<()> {
        self.0.execute(
            "CREATE TABLE IF NOT EXISTS files (
                 id        INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
                 path 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 (
                 file INTEGER NOT NULL,
                 tag  INTEGER NOT NULL
             )",
            params![],
        )?;
        Ok(())
    }

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

    /// Return the id of the file located at `path`.
    fn id_for_path<T: AsRef<Path>>(&self, path: T) -> Result<i64> {
        let path = String::from(path.as_ref().to_str().unwrap());
        Ok(self
            .0
            .prepare("SELECT id FROM files WHERE path = ?")?
            .query_row(params![path], |row| Ok(row.get::<_, i64>(0)))??)
    }

    /// Index the file at `path` in the tag database.
    pub fn add_file<T, S>(&self, path: T, tags: &Vec<S>) -> Result<File>
    where
        T: AsRef<Path>,
        S: AsRef<str>,
    {
        if self.id_for_path(&path).is_ok() {
            bail!("'{:?}' already indexed by the tag database.", path.as_ref())
        }

        // Insert the file row.
        self.0
            .prepare("INSERT INTO files (path) VALUES(?)")?
            .insert(params![path.as_ref().to_str().unwrap()])?;

        // Associate each tag with the file, creating a row for the tag if it
        // does not yet exist in the database.

        let path = String::from(path.as_ref().to_str().unwrap());
        let id = self
            .0
            .prepare("SELECT id FROM files WHERE path = ?")?
            .query_row(params![path], |row| Ok(row.get::<_, i64>(0)))??;
        self.tag_file(id, tags)?;

        let tags = self.tags_for_file(id)?;
        Ok(File { id, path, tags })
    }

    /// Remove the file named by `id` from the tag database.
    pub fn remove_file(&self, id: i64) -> Result<()> {
        // Remove the file from the `files` table,
        self.0
            .execute("DELETE FROM files WHERE id = ?", params![id])?;
        // but also remove it from the table for mapping tags.
        self.0
            .execute("DELETE FROM mapping WHERE file = ?", params![id])?;
        Ok(())
    }

    /// Return the identifier for `tag`, or `Err` if no such tag exists.
    fn tag_id<T: AsRef<str>>(&self, tag: T) -> Result<i64> {
        Ok(self
            .0
            .prepare("SELECT * FROM tags WHERE name = ?")?
            .query_row(params![tag.as_ref()], |row| Ok(row.get::<_, i64>(0)))??)
    }

    /// Return the identifiers for all tags named similarly to `tag`.
    ///
    /// This function treats `tag` as a wildcard, and therefore only makes sense
    /// in the context of a query. For inserting images to the database, or any
    /// other action where a "canonical" tag is desired, reach for `tag_id`
    /// instead.
    fn tag_ids<T: AsRef<str>>(&self, tag: T) -> Result<Vec<i64>> {
        Ok(self
            .0
            .prepare("SELECT id FROM tags WHERE name LIKE ?")?
            .query_map(params![tag.as_ref().replace("*", "%")], |row| row.get(0))?
            .filter_map(|id| id.ok())
            .collect())
    }

    /// Associate all of `tags` with the file named by `id`.
    ///
    /// If any tag lacks a row in the database, it will be created. The contents
    /// of `tags` should not contain wildcardcard expressions.
    pub fn tag_file<T: AsRef<str>>(&self, id: i64, tags: &Vec<T>) -> Result<()> {
        for tag in tags.iter() {
            // Ensure that a row for `tag` into the database.
            if self.tag_id(tag.as_ref()).is_err() {
                self.0
                    .execute("INSERT INTO tags (name) VALUES(?)", params![tag.as_ref()])?;
            }
            self.0.execute(
                "INSERT INTO mapping (file, tag) VALUES(?, ?)",
                params![id, self.tag_id(tag.as_ref()).unwrap()],
            )?;
        }
        Ok(())
    }

    /// Remove all of `tags` from the file named by `id`.
    ///
    /// The contents of `tags` can contain wildcardcard expressions.
    pub fn untag_file<T: AsRef<str>>(&self, id: i64, tags: &Vec<T>) -> Result<()> {
        for tag in tags.iter() {
            // Silently skip any tags which do not exist in the database.
            if let Ok(tags) = self.tag_ids(tag.as_ref()) {
                for tag in tags.iter() {
                    self.0
                        .prepare("DELETE FROM mapping WHERE file = ? AND tag = ?")?
                        .insert(params![id, tag])?;
                }
            }
        }
        Ok(())
    }

    /// Return the `File` object for the file named by `id`.
    pub fn file_by_id(&self, id: i64) -> Result<File> {
        let tags = self.tags_for_file(id)?;
        Ok(self
            .0
            .prepare("SELECT id, path FROM files WHERE id = ?")?
            .query_row(params![id], |row| {
                Ok(File {
                    id: row.get(0)?,
                    path: row.get(1)?,
                    tags,
                })
            })?)
    }

    /// Return the `File` object for the file at `path`.
    pub fn file_by_path<T: AsRef<Path>>(&self, path: T) -> Result<File> {
        self.file_by_id(self.id_for_path(path)?)
    }

    /// Return the files in the database satisfying `q` and `k`.
    pub fn query(&self, q: &Query, k: Option<Keyset>) -> Result<Vec<File>> {
        // Return an empty vector if the query contains tags that aren't in the
        // database.
        for part in q.iter() {
            if let Tag(Is(name)) = part {
                if self.tag_id(name).is_err() {
                    return Ok(vec![]);
                }
            }
        }

        let predicates = q
            .iter()
            .filter_map(|field| match field {
                Tag(Is(expression)) => Some(
                    self.tag_ids(expression)
                        .ok()?
                        .iter()
                        .map(|tag| format!("(mapping.tag = {})", tag))
                        .collect::<Vec<_>>(),
                ),
                Tag(IsNot(expression)) => Some(
                    self.tag_ids(expression)
                        .ok()?
                        .iter()
                        .map(|tag| {
                            format!(
                                "(files.id NOT IN
                                      (SELECT files.id FROM files
                                       JOIN mapping ON files.id = mapping.file
                                       WHERE mapping.tag = {}))",
                                tag
                            )
                        })
                        .collect::<Vec<_>>(),
                ),
                Filename(Is(expression)) => Some(vec![format!(
                    "(file.path LIKE %{})",
                    expression.replace("*", "%")
                )]),
                Filename(IsNot(expression)) => Some(vec![format!(
                    "(file.path NOT LIKE %{})",
                    expression.replace("*", "%")
                )]),
            })
            .flat_map(|s| s)
            .collect::<Vec<_>>();

        let query = format!(
            "SELECT files.id FROM files
             JOIN mapping
             ON files.id = mapping.file
             WHERE {predicate}
             {since}
             GROUP BY files.id
             ORDER BY files.id DESC
             {limit}",
            predicate = if predicates.is_empty() {
                "1 == 1".into()
            } else {
                predicates.join(" AND ")
            },
            since = if let Some(ref keyset) = k {
                format!("AND files.id < {}", keyset.last_id)
            } else {
                "".into()
            },
            limit = if let Some(ref keyset) = k {
                format!("LIMIT {}", keyset.how_many)
            } else {
                "".into()
            }
        );

        Ok(self
            .0
            .prepare(&query)?
            .query_map(params![], |row| row.get::<_, i64>(0))?
            .filter_map(|id| id.ok().and_then(|id| self.file_by_id(id).ok()))
            .collect())
    }

    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 parse_query_empty() {
        let parsed = parse_query("").unwrap();
        assert!(parsed.len() == 0);
    }

    #[test]
    fn parse_query_tags_only() {
        let parsed = parse_query("tag1 tag2 tag3").unwrap();
        assert_eq!(parsed.len(), 3);
        assert!(parsed.contains(&Tag(Is("tag1".into()))));
        assert!(parsed.contains(&Tag(Is("tag2".into()))));
        assert!(parsed.contains(&Tag(Is("tag3".into()))));
    }

    #[test]
    fn parse_query_tag_negations_only() {
        let parsed = parse_query("-tag1 -tag2 -tag3").unwrap();
        assert_eq!(parsed.len(), 3);
        assert!(parsed.contains(&Tag(IsNot("tag1".into()))));
        assert!(parsed.contains(&Tag(IsNot("tag2".into()))));
        assert!(parsed.contains(&Tag(IsNot("tag3".into()))));
    }

    #[test]
    fn parse_query_filenames_only() {
        let parsed = parse_query("filename:1 filename:2 filename:3").unwrap();
        assert_eq!(parsed.len(), 3);
        assert!(parsed.contains(&Filename(Is("1".into()))));
        assert!(parsed.contains(&Filename(Is("2".into()))));
        assert!(parsed.contains(&Filename(Is("3".into()))));
    }

    #[test]
    fn parse_query_filename_negations_only() {
        let parsed = parse_query("-filename:1 -filename:2 -filename:3").unwrap();
        assert_eq!(parsed.len(), 3);
        assert!(parsed.contains(&Filename(IsNot("1".into()))));
        assert!(parsed.contains(&Filename(IsNot("2".into()))));
        assert!(parsed.contains(&Filename(IsNot("3".into()))));
    }

    #[test]
    fn parse_query_combined() {
        let parsed = parse_query("tag1 -tag2 filename:1 -filename:2").unwrap();
        assert_eq!(parsed.len(), 4);
        assert!(parsed.contains(&Tag(Is("tag1".into()))));
        assert!(parsed.contains(&Tag(IsNot("tag2".into()))));
        assert!(parsed.contains(&Filename(Is("1".into()))));
        assert!(parsed.contains(&Filename(IsNot("2".into()))));
    }

    #[test]
    fn test_add_file_fail_on_already_indexed() {
        let tb = TagDatabase::new_mem().unwrap();
        tb.add_file("test", &vec!["tag1"]).unwrap();
        assert!(tb.add_file("test", &vec!["tag2"]).is_err());
    }

    #[test]
    fn test_add_file_output_reflects_input() {
        let tb = TagDatabase::new_mem().unwrap();
        let img = tb.add_file("test", &vec!["tag1", "tag2"]).unwrap();
        assert_eq!(img.path, "test");
        assert_eq!(img.tags.len(), 2);
        assert!(img.tags.contains(&String::from("tag1")));
        assert!(img.tags.contains(&String::from("tag2")));
    }

    #[test]
    fn test_remove_file() {
        let tb = TagDatabase::new_mem().unwrap();
        let img = tb.add_file("test1", &vec!["tag1", "tag2"]).unwrap();
        tb.remove_file(img.id).unwrap();

        let query = parse_query("tag1").unwrap();
        let results = tb.query(&query, None).unwrap();
        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_query_tags_only() {
        let tb = TagDatabase::new_mem().unwrap();
        tb.add_file("test1", &vec!["tag1", "tag2"]).unwrap();
        tb.add_file("test2", &vec!["tag2"]).unwrap();

        let query = parse_query("tag1").unwrap();
        let results = tb.query(&query, None).unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].path, "test1");

        let query = parse_query("tag2").unwrap();
        let results = tb.query(&query, None).unwrap();
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn test_query_tags_not_indexed() {
        let tb = TagDatabase::new_mem().unwrap();
        tb.add_file("test1", &vec!["tag1", "tag2"]).unwrap();

        let query = parse_query("tag3").unwrap();
        let results = tb.query(&query, None).unwrap();
        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_query_tag_negations_only() {
        let tb = TagDatabase::new_mem().unwrap();
        tb.add_file("test1", &vec!["tag1", "tag2"]).unwrap();
        tb.add_file("test2", &vec!["tag2"]).unwrap();

        let query = parse_query("-tag1").unwrap();
        let results = tb.query(&query, None).unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].path, "test2");
    }

    #[test]
    fn test_query_tags_with_keyset() {
        let tb = TagDatabase::new_mem().unwrap();
        tb.add_file("test1", &vec!["tag1"]).unwrap();
        tb.add_file("test2", &vec!["tag1"]).unwrap();

        let query = parse_query("tag1").unwrap();
        let results = tb
            .query(
                &query,
                Some(Keyset {
                    last_id: 3,
                    how_many: 1,
                }),
            )
            .unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].path, "test2");

        let results = tb
            .query(
                &query,
                Some(Keyset {
                    last_id: 2,
                    how_many: 1,
                }),
            )
            .unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].path, "test1");
    }

    #[test]
    fn test_empty_query() {
        let tb = TagDatabase::new_mem().unwrap();
        tb.add_file("test1", &vec!["tag1"]).unwrap();
        tb.add_file("test2", &vec!["tag1"]).unwrap();

        let query = parse_query("").unwrap();
        let results = tb.query(&query, None).unwrap();
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn test_tag_file() {
        let tb = TagDatabase::new_mem().unwrap();
        let img = tb.add_file("test1", &vec!["tag1", "tag2"]).unwrap();
        tb.tag_file(img.id, &vec!["tag3"]).unwrap();

        let query = parse_query("tag3").unwrap();
        let results = tb.query(&query, None).unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].path, "test1");
    }

    #[test]
    fn test_untag_file() {
        let tb = TagDatabase::new_mem().unwrap();
        let img = tb.add_file("test1", &vec!["tag1", "tag2"]).unwrap();
        tb.untag_file(img.id, &vec!["tag2"]).unwrap();

        let query = parse_query("tag2").unwrap();
        let results = tb.query(&query, None).unwrap();
        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_id_resolution() {
        let tb = TagDatabase::new_mem().unwrap();
        let img = tb.add_file("test1", &vec!["tag1", "tag2"]).unwrap();
        assert_eq!(tb.file_by_id(img.id).unwrap().path, "test1");
    }

    #[test]
    fn test_path_resolution() {
        let tb = TagDatabase::new_mem().unwrap();
        let img = tb.add_file("test1", &vec!["tag1", "tag2"]).unwrap();
        assert_eq!(tb.file_by_path(img.path).unwrap().path, "test1");
    }

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