310 lines
6.8 KiB
Rust
310 lines
6.8 KiB
Rust
use std::fs;
|
|
use std::io::Write;
|
|
use std::str::FromStr;
|
|
use std::{io, path::PathBuf};
|
|
|
|
use anyhow::{Result, anyhow};
|
|
use rusqlite::{Connection, OpenFlags, params_from_iter, types::Value};
|
|
|
|
const DB_PATH: &str = ".tags";
|
|
|
|
pub struct Files {
|
|
pub files: Vec<(String, String)>,
|
|
pub paths: Vec<String>,
|
|
pub show_tags: bool,
|
|
}
|
|
|
|
pub fn init_db() -> anyhow::Result<()> {
|
|
let conn = Connection::open(DB_PATH)?;
|
|
|
|
conn.execute(
|
|
r#"CREATE TABLE tag(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name VARCHAR(255) NOT NULL UNIQUE
|
|
);"#,
|
|
(),
|
|
)?;
|
|
|
|
conn.execute(
|
|
r#"CREATE TABLE file(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
path VARCHAR(255) NOT NULL UNIQUE
|
|
);"#,
|
|
(),
|
|
)?;
|
|
|
|
conn.execute(
|
|
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)
|
|
);"#,
|
|
(),
|
|
)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
pub 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(())
|
|
}
|
|
|
|
pub 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)
|
|
}
|
|
|
|
pub 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(())
|
|
}
|
|
|
|
pub 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(())
|
|
}
|
|
|
|
pub 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(())
|
|
}
|
|
|
|
pub 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(())
|
|
}
|
|
|
|
pub fn list_files(conn: &Connection, all_files: bool, query: Option<String>) -> Result<Files> {
|
|
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 all_files {
|
|
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) = 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 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 files: Vec<(String, String)> = stmt
|
|
.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?
|
|
.flatten()
|
|
.collect();
|
|
|
|
Ok(Files {
|
|
files,
|
|
paths,
|
|
show_tags,
|
|
})
|
|
}
|