diff options
Diffstat (limited to 'src/database.rs')
| -rw-r--r-- | src/database.rs | 318 |
1 files changed, 318 insertions, 0 deletions
diff --git a/src/database.rs b/src/database.rs new file mode 100644 index 0000000..77397a5 --- /dev/null +++ b/src/database.rs @@ -0,0 +1,318 @@ +// 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); + +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 images_by_tags(&self, tags: &[&str]) -> Result<Vec<Image>> { + let ids: Vec<i64> = tags + .iter() + .filter_map(|tag| self.tag_id(tag).ok()) + .collect(); + let mut tags = Vec::from(tags); + tags.dedup(); + + if ids.len() != tags.len() { + return Ok(vec![]); + } + + let union = if tags.is_empty() { + String::from( + "INNER JOIN mapping + ON images.id = mapping.image", + ) + } else { + (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") + }; + // WHERE id > {lastId} + // LIMIT 50 + let query = format!( + "SELECT id, blake2, filename, orig_dir + FROM images + {union} + GROUP BY id + ORDER BY id DESC", + union = union + ); + + 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()) + } +} + +/// 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); + } +} |