path: root/src/main.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/>.

#[macro_use]
extern crate anyhow;
extern crate dirs;

use anyhow::Result;
use blake2::{Blake2b, Digest};
use rusqlite::{params, Connection};
use std::env;
use std::io;
use std::path::Path;

/// Return a hex-encoded string representing the contents of `hash`.
fn hexlify(hash: &[u8]) -> String {
    hash.iter()
        .map(|b| format!("{:x}", b))
        .collect::<Vec<String>>()
        .join("")
}

#[derive(Debug)]
struct TagBase {
    conn: Connection,
}

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

struct Image {
    id: i64,
    blake2: String,
    filename: String,
    orig_dir: String,
}

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

    fn make_temporary() -> Result<Self> {
        let tb = TagBase {
            conn: Connection::open_in_memory()?,
        };
        tb.initialize_tables()?;
        Ok(tb)
    }

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

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

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

    /// Return the internal id for `tag`.
    fn tag_id(&self, tag: &str) -> Result<i64> {
        Ok(self
            .conn
            .prepare("SELECT * FROM tags WHERE name = ?")?
            .query_row(params![tag], |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.conn.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.conn.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<Vec<String>> {
        Ok(self
            .conn
            .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<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.len() == 0 {
            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
            .conn
            .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())
    }
}

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

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

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

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

    #[test]
    fn tag_image() {
        let tb = TagBase::make_temporary().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 = TagBase::make_temporary().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);
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: {} ACTION [ARGS ...]", args[0]);
        std::process::exit(1);
    }

    let tb = if let Ok(val) = env::var("TB_PATH") {
        TagBase::new(&val)
    } else {
        if let Some(dir) = dirs::config_dir() {
            TagBase::new(dir.join("pometka.db").to_str().unwrap())
        } else {
            TagBase::new(Path::new("/tmp").join("pometka.db").to_str().unwrap())
        }
    }
    .unwrap();

    let action = &args[1];
    match action.to_lowercase().as_str() {
        "add" => {
            if args.len() < 3 {
                eprintln!("usage: {} add PATH [TAGS ...]", args[0]);
                std::process::exit(1);
            }
            let path = Path::new(&args[2]);
            let hash = hash_file(path).unwrap();
            let id = tb.import_image(path, &hash).unwrap();
            for arg in args[3..].iter() {
                tb.add_tag(arg).unwrap();
                tb.tag_image(id, arg).unwrap();
            }
        }
        "add_tags" => {
            if args.len() < 3 {
                eprintln!("usage: {} add_tags ID [TAGS ...]", args[0]);
                std::process::exit(1);
            }
            let id = args[2].parse::<i64>().unwrap();
            for arg in args[3..].iter() {
                tb.add_tag(arg).unwrap();
                tb.tag_image(id, arg).unwrap();
            }
        }
        "remove_tags" => {
            if args.len() < 3 {
                eprintln!("usage: {} remove_tags ID [TAGS ...]", args[0]);
                std::process::exit(1);
            }
            let id = args[2].parse::<i64>().unwrap();
            for arg in args[3..].iter() {
                tb.untag_image(id, arg).unwrap();
            }
        }
        "query" => {
            let tags = args[2..].iter().map(|s| s.as_str()).collect::<Vec<&str>>();
            for tag in tags.iter() {
                if let Err(_) = tb.tag_id(tag) {
                    eprintln!("unknown tag: {}", tag);
                    std::process::exit(1);
                }
            }
            for image in tb.images_by_tags(&tags[..]).unwrap() {
                println!(
                    "{},{},{}{}",
                    image.id, image.blake2, image.orig_dir, image.filename
                );
            }
        }
        _ => {
            eprintln!("Unknown action '{}'", action);
            std::process::exit(1);
        }
    }
}