Mass renames and tagging

Cargo.lock

@@ -52,6 +52,12 @@ dependencies = [
  "windows-sys",
 ]
 
+[[package]]
+name = "anyhow"
+version = "1.0.98"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487"
+
 [[package]]
 name = "bitflags"
 version = "2.9.0"
@@ -249,6 +255,7 @@ dependencies = [
 name = "taggo"
 version = "0.1.0"
 dependencies = [
+ "anyhow",
  "clap",
  "rusqlite",
 ]

Cargo.toml

@@ -17,3 +17,4 @@ unsafe_code = { level = "forbid" }
 [dependencies]
 clap = { version = "4.5", features = ["derive"] }
 rusqlite = { version = "0.35", features = ["bundled"] }
+anyhow = { version = "1.0" }

src/cli.rs

@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+
 use clap::{Args, Parser, Subcommand};
 
 #[derive(Debug, Parser, PartialEq, Eq)]
@@ -10,13 +12,24 @@ pub struct CliArgs {
 #[derive(Debug, Subcommand, PartialEq, Eq)]
 pub enum Commands {
 	Init,
+	Untag(UntagArgs),
+	Tag {
+		#[clap(flatten)]
+		args: TagArgs,
+
+		#[arg(short, long, requires = "file_tags", value_delimiter = ' ', num_args = 1..)]
+		files: Vec<PathBuf>,
+
+		#[arg(short('t'), long("tags"), requires = "files", value_delimiter = ' ', num_args = 1..)]
+		file_tags: Vec<String>,
+	},
 	Tags(TagsArgs),
+	Files(FilesArgs),
 }
 
 #[derive(Debug, Args, PartialEq, Eq)]
 pub struct TagsArgs {
-	#[arg(long("list"), short)]
-	pub list: bool,
+	pub query: Option<String>,
 
 	#[command(subcommand)]
 	pub commands: Option<TagsCommands>,
@@ -24,6 +37,32 @@ pub struct TagsArgs {
 
 #[derive(Debug, Subcommand, PartialEq, Eq)]
 pub enum TagsCommands {
+	List,
 	Add { add: Vec<String> },
 	Remove { remove: Vec<String> },
+	Rename { renames: Vec<String> },
+}
+
+#[derive(Debug, Args, PartialEq, Eq)]
+pub struct TagArgs {
+	// #[clap(conflicts_with = "files")]
+	pub file: Option<PathBuf>,
+
+	// #[clap(conflicts_with = "files")]
+	pub tags: Vec<String>,
+}
+
+#[derive(Debug, Args, PartialEq, Eq)]
+pub struct UntagArgs {
+	pub file: PathBuf,
+
+	pub tags: Option<Vec<String>>,
+}
+
+#[derive(Debug, Args, PartialEq, Eq)]
+pub struct FilesArgs {
+	pub query: Option<String>,
+
+	#[arg(long, short, exclusive = true, default_value_t = false)]
+	pub all: bool,
 }

src/main.rs

@@ -1,119 +1,70 @@
-use std::{io, path, process::exit};
-
-use io::Write;
-
+use anyhow::anyhow;
 use clap::Parser;
 use cli::CliArgs;
-use rusqlite::{Connection, params_from_iter};
+use rusqlite::{Connection, OpenFlags};
 
 const DB_PATH: &str = ".tags";
 
 mod cli;
+mod tags;
 
-fn main() {
+fn main() -> anyhow::Result<()> {
 	let args = CliArgs::parse();
 
 	match args.commands {
-		cli::Commands::Init => {
-			if !has_database() {
-				let conn = Connection::open(DB_PATH).unwrap();
-				init_db(&conn);
-			} else {
-				panic!("Database is already initialized in this folder");
-			}
-			exit(0);
-		}
-		cli::Commands::Tags(args) => {
-			let conn = Connection::open(DB_PATH).unwrap();
-
-			if args.list {
-				let mut w = io::stdout();
-				let tags = list_tags(&conn);
-				for tag in tags {
-					writeln!(&mut w, "{}", tag).unwrap();
-				}
-				w.flush().unwrap();
-				return;
-			}
-
-			match args.commands {
-				Some(cli::TagsCommands::Add { add }) => add_tags(&conn, add),
-				Some(cli::TagsCommands::Remove { remove }) => remove_tags(&conn, remove),
-				_ => (),
-			};
-		}
+		cli::Commands::Init => init_db().map_err(|_| anyhow!("Database already exists!")),
+		cli::Commands::Tag {
+			args,
+			files,
+			file_tags,
+		} => tags::handle_tag(args, files, file_tags),
+		cli::Commands::Untag(args) => tags::handle_untag(args),
+		cli::Commands::Tags(args) => tags::handle_tags(args),
+		cli::Commands::Files(args) => tags::handle_files(args),
 	}
 }
 
-fn list_tags(conn: &Connection) -> Vec<String> {
-	let mut stmt = conn.prepare("SELECT name FROM tag").unwrap();
-	let result = stmt.query_map([], |row| row.get(0)).unwrap();
+fn init_db() -> anyhow::Result<()> {
+	let conn = Connection::open(DB_PATH)?;
 
-	let mut tags = Vec::new();
-	for name in result {
-		tags.push(name.unwrap());
-	}
-
-	tags
-}
-
-fn remove_tags(conn: &Connection, tags: Vec<String>) {
-	let mut query = r#"DELETE FROM tag WHERE name IN ("#.to_string();
-
-	for (i, _tag) in tags.iter().enumerate() {
-		query.push('?');
-		if i < tags.len() - 1 {
-			query.push(',');
-		}
-	}
-
-	query.push(')');
-
-	conn.execute(&query, params_from_iter(tags)).unwrap();
-}
-
-fn add_tags(conn: &Connection, tags: Vec<String>) {
-	let mut query = r#"INSERT INTO tag(name) VALUES"#.to_string();
-
-	for (i, _tag) in tags.iter().enumerate() {
-		query.push_str("(?)");
-		if i < tags.len() - 1 {
-			query.push(',');
-		}
-	}
-
-	conn.execute(&query, params_from_iter(tags)).unwrap();
-}
-
-fn init_db(conn: &Connection) {
 	conn.execute(
-		r#"CREATE TABLE IF NOT EXISTS tag(
+		r#"CREATE TABLE tag(
 			id INTEGER PRIMARY KEY AUTOINCREMENT, 
 			name VARCHAR(255) NOT NULL UNIQUE
 		);"#,
 		(),
-	)
-	.unwrap();
+	)?;
 
 	conn.execute(
-		r#"CREATE TABLE IF NOT EXISTS file(
-			id INT NOT NULL PRIMARY KEY, 
+		r#"CREATE TABLE file(
+			id INTEGER PRIMARY KEY AUTOINCREMENT, 
 			path VARCHAR(255) NOT NULL UNIQUE
 		);"#,
 		(),
-	)
-	.unwrap();
+	)?;
 
 	conn.execute(
-		r#"CREATE TABLE IF NOT EXISTS file_tag(
-			tag_id INT REFERENCES tag(id), 
-			file_id INT REFERENCES file(id)
+		r#"CREATE TABLE file_tag(
+			file_id INTEGER REFERENCES file(id) ON DELETE CASCADE,
+			tag_id INTEGER REFERENCES tag(id) ON DELETE CASCADE,
+			PRIMARY KEY (file_id, tag_id)
 		);"#,
 		(),
-	)
-	.unwrap();
+	)?;
+
+	Ok(())
 }
 
-pub fn has_database() -> bool {
-	path::Path::new(DB_PATH).exists()
+pub fn try_database() -> anyhow::Result<Connection> {
+	let conn = Connection::open_with_flags(
+		DB_PATH,
+		OpenFlags::SQLITE_OPEN_READ_WRITE
+			| OpenFlags::SQLITE_OPEN_URI
+			| OpenFlags::SQLITE_OPEN_NO_MUTEX,
+	)
+	.map_err(|_| {
+		anyhow!("Database does not exist! Run `taggo init` first to generate a new database.")
+	})?;
+
+	Ok(conn)
 }

src/tags.rs

@@ -0,0 +1,342 @@
+use std::{
+	fs,
+	io::{self, Write},
+	path::PathBuf,
+	str::FromStr,
+};
+
+use anyhow::{Result, anyhow};
+
+use rusqlite::{Connection, params_from_iter, types::Value};
+
+pub const TERM_RESET: &str = "\x1b[0m";
+pub const TERM_BOLD: &str = "\x1b[1m";
+
+use crate::{
+	cli::{self, FilesArgs, TagArgs, TagsArgs, UntagArgs},
+	try_database,
+};
+
+pub fn handle_files(args: FilesArgs) -> Result<()> {
+	let conn = try_database()?;
+
+	let mut sql = r#"SELECT path, group_concat(name) as tags FROM file_tag
+			JOIN file ON file.id = file_tag.file_id
+			JOIN tag ON tag.id = file_tag.tag_id"#
+		.to_string();
+
+	let mut show_tags = true;
+
+	let mut paths: Vec<PathBuf> = vec![];
+	if args.all {
+		paths = fs::read_dir(".")?
+			.filter_map(|f| f.ok())
+			.filter(|f| f.metadata().is_ok() && f.metadata().unwrap().is_file())
+			.map(|f| f.path())
+			.collect();
+	} else if let Some(query) = args.query {
+		let path = PathBuf::from_str(&query).unwrap_or_default();
+		match path.exists() {
+			true => {
+				sql.push_str(&format!(" WHERE path IN ('{}')", query));
+				paths = vec![path];
+			}
+			false => {
+				sql.push_str(&format!(" WHERE tag.name LIKE '%{}%'", query));
+				show_tags = false;
+			}
+		};
+	}
+
+	let mut all_paths: Vec<String> = paths
+		.iter_mut()
+		.map(|f| f.to_string_lossy().replace("./", "").to_string())
+		.collect();
+
+	sql.push_str(" GROUP BY path");
+
+	let mut stmt = conn.prepare(&sql)?;
+	let tagged_files: Vec<(String, String)> = stmt
+		.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?
+		.flatten()
+		.collect();
+
+	let mut w = io::stdout();
+	if !tagged_files.is_empty() {
+		if args.all {
+			writeln!(&mut w, "{TERM_BOLD}Tagged Files: {TERM_RESET}")?;
+		}
+		for (file, tags) in tagged_files {
+			// Remove already tagged files
+			if let Some(idx) = all_paths.iter().position(|v| *v == file) {
+				all_paths.swap_remove(idx);
+			}
+
+			if show_tags {
+				writeln!(&mut w, "{}: {}", file, tags)?;
+			} else {
+				writeln!(&mut w, "{}", file)?;
+			}
+		}
+	} else {
+		writeln!(&mut w, "No files registered!")?;
+	}
+
+	if args.all && !all_paths.is_empty() {
+		writeln!(&mut w, "\n{TERM_BOLD}Untagged Files: {TERM_RESET}")?;
+		for path in all_paths {
+			let path_str = path;
+			writeln!(&mut w, "{}", path_str)?;
+		}
+	}
+
+	w.flush()?;
+
+	Ok(())
+}
+
+pub fn handle_tag(args: TagArgs, files: Vec<PathBuf>, file_tags: Vec<String>) -> Result<()> {
+	let mut conn = try_database()?;
+
+	if let Some(file) = args.file {
+		tag_file(&mut conn, &file, &args.tags)?;
+	} else {
+		// TODO: Maybe rework this at some point so its more optimized for batch insertion ?
+		for file in &files {
+			tag_file(&mut conn, file, &file_tags)?;
+		}
+	}
+
+	Ok(())
+}
+
+pub fn handle_untag(args: UntagArgs) -> Result<()> {
+	let mut conn = try_database()?;
+
+	untag_file(&mut conn, args.file, args.tags)?;
+
+	Ok(())
+}
+
+pub fn handle_tags(args: TagsArgs) -> Result<()> {
+	let mut conn = try_database()?;
+
+	match args.commands {
+		Some(cli::TagsCommands::List) | None => {
+			let mut w = io::stdout();
+			let tags = list_tags(&conn)?;
+			if !tags.is_empty() {
+				for tag in tags {
+					writeln!(&mut w, "{} - {}", tag.0, tag.1)?;
+				}
+			} else {
+				writeln!(&mut w, "No tags registered!")?;
+			}
+			w.flush()?;
+			Ok(())
+		}
+		Some(cli::TagsCommands::Add { add }) => add_tags(&mut conn, add),
+		Some(cli::TagsCommands::Remove { remove }) => remove_tags(&mut conn, remove),
+		Some(cli::TagsCommands::Rename { renames }) => rename_tag(&conn, renames),
+	}
+}
+
+fn rename_tag(conn: &Connection, renames: Vec<String>) -> Result<()> {
+	if renames.is_empty() {
+		return Err(anyhow!("No input found!"));
+	} else if renames.len() % 2 != 0 {
+		return Err(anyhow!(
+			"Attempting to rename an uneven amount of entries: {}",
+			renames.len()
+		));
+	}
+
+	let mut stmt = conn.prepare("UPDATE tag SET name = ? WHERE name = ?")?;
+
+	let mut i = 0;
+	while i < renames.len() {
+		// TODO: Merge this when let chaining is stabilized in 1.88
+		if let Some(old) = renames.get(i) {
+			if let Some(new) = renames.get(i + 1) {
+				stmt.execute([&new, &old])?;
+			}
+		}
+		i += 2;
+	}
+
+	Ok(())
+}
+
+fn list_tags(conn: &Connection) -> Result<Vec<(String, i32)>> {
+	let mut stmt = conn.prepare(
+		r#"SELECT name, COUNT(file_id) as count FROM tag
+		LEFT JOIN file_tag ON file_tag.tag_id = tag.id
+		GROUP BY name
+		ORDER BY count DESC"#,
+	)?;
+	let result = stmt
+		.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?
+		.flatten();
+
+	let mut tags: Vec<(String, i32)> = Vec::new();
+	for entry in result {
+		tags.push(entry);
+	}
+
+	Ok(tags)
+}
+
+fn remove_tags(conn: &mut Connection, tags: Vec<String>) -> Result<()> {
+	let tx = conn.transaction()?;
+	{
+		let mut query = r#"DELETE FROM tag WHERE name IN ("#.to_string();
+
+		for (i, _tag) in tags.iter().enumerate() {
+			query.push('?');
+			if i < tags.len() - 1 {
+				query.push(',');
+			}
+		}
+
+		query.push(')');
+
+		tx.execute(&query, params_from_iter(&tags))?;
+
+		if tx.changes() > 0 {
+			let mut w = io::stdout();
+			writeln!(&mut w, "Removed Tags:")?;
+			for tag in tags {
+				writeln!(&mut w, "  {tag}")?;
+			}
+		}
+	}
+	tx.commit()?;
+
+	Ok(())
+}
+
+fn add_tags(conn: &mut Connection, tags: Vec<String>) -> Result<()> {
+	let tx = conn.transaction()?;
+	{
+		let mut query = r#"INSERT INTO tag(name) VALUES"#.to_string();
+
+		for (i, _tag) in tags.iter().enumerate() {
+			query.push_str("(?)");
+			if i < tags.len() - 1 {
+				query.push(',');
+			}
+		}
+
+		tx.execute(&query, params_from_iter(&tags))?;
+
+		if tx.changes() > 0 {
+			let mut w = io::stdout();
+			writeln!(&mut w, "Added Tags:")?;
+			for tag in tags {
+				writeln!(&mut w, "  {tag}")?;
+			}
+		}
+	}
+	tx.commit()?;
+
+	Ok(())
+}
+
+fn tag_file(conn: &mut Connection, file: &PathBuf, tags: &Vec<String>) -> Result<()> {
+	let file = file
+		.to_str()
+		.ok_or(anyhow!("Unable to find path {:?}", file))?
+		.to_string();
+
+	let tx = conn.transaction()?;
+	{
+		let file_id: Option<i32> = tx
+			.query_row(r#"SELECT id FROM file WHERE path = ?"#, [&file], |row| {
+				row.get(0)
+			})
+			.ok();
+
+		let file_id = match file_id {
+			Some(id) => id,
+			None => {
+				let mut stmt = tx.prepare(r#"INSERT INTO file(path) VALUES (?) RETURNING id"#)?;
+				stmt.query_row([file], |row| -> Result<i32, rusqlite::Error> { row.get(0) })?
+			}
+		};
+
+		let mut sql = "SELECT id FROM tag WHERE name IN ".to_string();
+		sql.push('(');
+		let size = tags.len();
+		for (i, _tag) in tags.iter().enumerate() {
+			sql.push('?');
+			if i < size - 1 {
+				sql.push(',');
+			}
+		}
+		sql.push(')');
+
+		let mut stmt = tx.prepare(&sql)?;
+		let result = stmt
+			.query_map(
+				params_from_iter(tags),
+				|row| -> Result<i32, rusqlite::Error> { row.get(0) },
+			)?
+			.flatten();
+
+		let mut tags_ids = Vec::new();
+		for id in result {
+			tags_ids.push(id);
+		}
+
+		let mut stmt = tx.prepare(r#"INSERT INTO file_tag(file_id, tag_id) VALUES (?, ?)"#)?;
+
+		for tag_id in tags_ids {
+			stmt.execute((file_id, tag_id))?;
+		}
+	}
+	tx.commit()?;
+
+	Ok(())
+}
+
+fn untag_file(conn: &mut Connection, file: PathBuf, tags: Option<Vec<String>>) -> Result<()> {
+	let file = file
+		.to_str()
+		.ok_or(anyhow!("Unable to find path {:?}", file))?
+		.to_string();
+
+	let tx = conn.transaction()?;
+	{
+		let _file_id: i32 = tx
+			.query_row(r#"SELECT id FROM file WHERE path = ?"#, [&file], |row| {
+				row.get(0)
+			})
+			.map_err(|_| anyhow!("No {} file registered!", file))?;
+
+		let mut sql = r#"
+			DELETE FROM file_tag
+			WHERE ROWID IN (
+				SELECT file_tag.ROWID FROM file_tag
+				INNER JOIN file ON file_tag.file_id = file.id
+				INNER JOIN tag ON file_tag.tag_id = tag.id"#
+			.to_string();
+
+		let params = match tags {
+			Some(tags) => {
+				let tags = tags.join(",");
+				sql.push_str(" WHERE file.path = ?1 AND tag.name IN (?2))");
+				vec![Value::Text(file), Value::Text(tags)]
+			}
+			None => {
+				sql.push_str(" WHERE file.path = ?1)");
+				vec![Value::Text(file)]
+			}
+		};
+
+		let mut stmt = tx.prepare(&sql)?;
+		stmt.execute(params_from_iter(params))?;
+	}
+	tx.commit()?;
+
+	Ok(())
+}