From d59a92b04a398b274fbeda0de6d8e64c718837dd Mon Sep 17 00:00:00 2001 From: "Jakob L. Kreuze" Date: Mon, 25 May 2020 13:05:09 -0400 Subject: Split off into two files. --- src/database.rs | 318 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 300 +--------------------------------------------------- 2 files changed, 321 insertions(+), 297 deletions(-) create mode 100644 src/database.rs (limited to 'src') 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 +// +// 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 . + +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 { + 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 { + 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 { + 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 { + 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> { + let ids: Vec = 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::>() + .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 { + 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> { + 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 = (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); + } +} diff --git a/src/main.rs b/src/main.rs index 2dc0f2e..d2caba3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,222 +19,15 @@ extern crate anyhow; extern crate dirs; +mod database; + use anyhow::Result; use blake2::{Blake2b, Digest}; -use rusqlite::{params, Connection}; +use database::TagDatabase; use std::env; use std::io; use std::path::Path; -/// A connection to бирка-тян's tag database. -#[derive(Debug)] -struct TagDatabase(Connection); - -impl TagDatabase { - /// Instantiate a new `TagDatabase` whose backing database is at `path`. - pub fn new(path: &str) -> Result { - let tb = TagDatabase(Connection::open(path)?); - tb.initialize_tables()?; - Ok(tb) - } - - fn new_mem() -> Result { - let tb = TagDatabase(Connection::open_in_memory()?); - 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. - 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`. - fn tag_id(&self, tag: &str) -> Result { - 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. - fn import_image(&self, path: &Path, hash: &str) -> Result { - 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`. - fn image_id(&self, path: &Path) -> Result { - 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`. - 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`. - 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`. - fn tags_for_image(&self, image_id: i64) -> Result> { - 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()) - } - - /// Return the intersection of the sets of images matching each of `tags`. - fn images_by_tags(&self, tags: &[&str]) -> Result> { - let ids: Vec = 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::>() - .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()) - } -} - /// Return a base64-encoded BLAKE2b hash identifying the file at `path`. fn hash_file(path: &Path) -> Result { let mut file = std::fs::File::open(path)?; @@ -243,15 +36,6 @@ fn hash_file(path: &Path) -> Result { Ok(base64::encode(&hasher.result())) } -/// The result of an image query in бирка-тян's tag database. -#[derive(Debug)] -struct Image { - id: i64, - blake2: String, - filename: String, - orig_dir: String, -} - fn main() { let args: Vec = env::args().collect(); if args.len() < 2 { @@ -333,81 +117,3 @@ fn main() { } } } - -#[cfg(test)] -mod tests { - use super::*; - - #[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 = (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); - } -} -- cgit v1.3