From c3764c125a4c4ae89692f54f42d4c0951ec73389 Mon Sep 17 00:00:00 2001 From: Marc Plano-Lesay Date: Sun, 4 May 2025 15:41:02 +1000 Subject: [PATCH] feat: add the concept of tags Tags are attached to all magnet links provided by a given source, and passed to actions. This allows for e.g. better categorization in Bitmagnet. --- README.md | 8 + .../2025-05-04-035709_add_tags/down.sql | 2 + migrations/2025-05-04-035709_add_tags/up.sql | 14 ++ src/actions/bitmagnet/action.rs | 7 + src/actions/bitmagnet/client.rs | 6 +- src/actions/transmission/action.rs | 8 +- src/actions/transmission/client.rs | 3 +- src/app.rs | 1 + src/config.rs | 2 + src/db.rs | 146 +++++++++++++++++- src/main.rs | 1 + src/models.rs | 29 +++- src/schema.rs | 24 ++- 13 files changed, 240 insertions(+), 11 deletions(-) create mode 100644 migrations/2025-05-04-035709_add_tags/down.sql create mode 100644 migrations/2025-05-04-035709_add_tags/up.sql diff --git a/README.md b/README.md index 0722e84..2807d05 100644 --- a/README.md +++ b/README.md @@ -67,11 +67,19 @@ password = "password" # Optional username = "reddituser1" title_filter = "720p|1080p" # Optional regex pattern imdb_id = "tt1234567" # Optional IMDB ID for better metadata when sending to Bitmagnet +tags = ["movie", "hd"] # Optional tags [sources.example2] username = "reddituser2" ``` +### Tags + +Tags are optional and can be assigned to sources. They provide additional metadata that actions can use in different ways: + +- **Transmission**: Tags are passed as labels. +- **Bitmagnet**: The first tag is used as the content type (e.g., `movie`, `tv_show`), which helps Bitmagnet categorize the content properly. + ## Potential future features - Support for additional BitTorrent clients (Deluge, qBittorrent, etc.) diff --git a/migrations/2025-05-04-035709_add_tags/down.sql b/migrations/2025-05-04-035709_add_tags/down.sql new file mode 100644 index 0000000..e1f1ed8 --- /dev/null +++ b/migrations/2025-05-04-035709_add_tags/down.sql @@ -0,0 +1,2 @@ +DROP TABLE magnet_tags; +DROP TABLE tags; diff --git a/migrations/2025-05-04-035709_add_tags/up.sql b/migrations/2025-05-04-035709_add_tags/up.sql new file mode 100644 index 0000000..775a72b --- /dev/null +++ b/migrations/2025-05-04-035709_add_tags/up.sql @@ -0,0 +1,14 @@ +CREATE TABLE tags +( + id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, + name TEXT UNIQUE NOT NULL +); + +CREATE TABLE magnet_tags +( + magnet_id INTEGER NOT NULL, + tag_id INTEGER NOT NULL, + FOREIGN KEY (magnet_id) REFERENCES magnets (id), + FOREIGN KEY (tag_id) REFERENCES tags (id), + PRIMARY KEY (magnet_id, tag_id) +); diff --git a/src/actions/bitmagnet/action.rs b/src/actions/bitmagnet/action.rs index be8d6db..b121b2d 100644 --- a/src/actions/bitmagnet/action.rs +++ b/src/actions/bitmagnet/action.rs @@ -37,12 +37,16 @@ impl Action for BitmagnetAction { let mut failed_magnets = Vec::new(); for magnet in unprocessed_magnets { + let tags = db.get_tags_for_magnet(magnet.id)?; + let tag_refs: Vec<&str> = tags.iter().map(|s| s.as_str()).collect(); + match self .client .submit_magnet( &magnet.link, &magnet.published_at.and_utc(), &magnet.imdb_id, + tag_refs, ) .await { @@ -53,6 +57,9 @@ impl Action for BitmagnetAction { magnet.title ); debug!("Magnet link: {}", magnet.link); + if !tags.is_empty() { + debug!("Tags: {:?}", tags); + } db.mark_magnet_processed_for_table::(magnet.id)?; processed_magnets.push(magnet); } diff --git a/src/actions/bitmagnet/client.rs b/src/actions/bitmagnet/client.rs index cec2955..33c77c1 100644 --- a/src/actions/bitmagnet/client.rs +++ b/src/actions/bitmagnet/client.rs @@ -34,6 +34,7 @@ impl BitmagnetClient { magnet: &str, published_at: &DateTime, imdb_id: &Option, + tags: Vec<&str>, ) -> Result<()> { let url = self .base_url @@ -43,7 +44,6 @@ impl BitmagnetClient { let structured_magnet = Magnet::new(magnet).map_err(|e| eyre!("Invalid magnet link: {:?}", e))?; - // Create a JSON object with the required fields let mut json_body = json!({ "publishedAt": published_at.to_rfc3339(), "source": "reddit-magnet", @@ -71,6 +71,10 @@ impl BitmagnetClient { json_body["size"] = json!(size); } + if let Some(tag) = tags.first() { + json_body["contentType"] = json!(tag); + } + let response = self .client .post(url) diff --git a/src/actions/transmission/action.rs b/src/actions/transmission/action.rs index ab653be..554dc2b 100644 --- a/src/actions/transmission/action.rs +++ b/src/actions/transmission/action.rs @@ -37,7 +37,10 @@ impl Action for TransmissionAction { let mut failed_magnets = Vec::new(); for magnet in unprocessed_magnets { - match self.client.submit_magnet(&magnet.link).await { + let tags = db.get_tags_for_magnet(magnet.id)?; + let tag_refs: Vec<&str> = tags.iter().map(|s| s.as_str()).collect(); + + match self.client.submit_magnet(&magnet.link, tag_refs).await { Ok(_) => { debug!( "Successfully submitted magnet link to {}: {}", @@ -45,6 +48,9 @@ impl Action for TransmissionAction { magnet.title ); debug!("Magnet link: {}", magnet.link); + if !tags.is_empty() { + debug!("Tags: {:?}", tags); + } db.mark_magnet_processed_for_table::(magnet.id)?; processed_magnets.push(magnet); } diff --git a/src/actions/transmission/client.rs b/src/actions/transmission/client.rs index 5f1f9a5..8466286 100644 --- a/src/actions/transmission/client.rs +++ b/src/actions/transmission/client.rs @@ -36,10 +36,11 @@ impl TransmissionClient { } /// Submit a magnet link to Transmission - pub async fn submit_magnet(&mut self, magnet: &str) -> Result<()> { + pub async fn submit_magnet(&mut self, magnet: &str, tags: Vec<&str>) -> Result<()> { let args = TorrentAddArgs { filename: Some(magnet.to_string()), download_dir: Some(self.download_dir.clone()), + labels: Some(tags.iter().map(|t| t.to_string()).collect()), ..Default::default() }; diff --git a/src/app.rs b/src/app.rs index 0c0a9c7..b4e4509 100644 --- a/src/app.rs +++ b/src/app.rs @@ -120,6 +120,7 @@ impl App { magnet_links, timestamp: post.created, imdb_id: source_config.imdb_id.clone(), + tags: source_config.tags.clone(), }; // Store the post info in the database diff --git a/src/config.rs b/src/config.rs index 7bda111..fdcd861 100644 --- a/src/config.rs +++ b/src/config.rs @@ -21,6 +21,8 @@ pub struct SourceConfig { pub username: String, pub title_filter: Option, pub imdb_id: Option, + #[serde(default)] + pub tags: Vec, } /// Main application configuration diff --git a/src/db.rs b/src/db.rs index 56cd851..f21c800 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1,5 +1,5 @@ -use crate::models::{Magnet, NewBitmagnetProcessed, NewMagnet, NewTransmissionProcessed}; -use crate::schema::{bitmagnet_processed, magnets, transmission_processed}; +use crate::models::{Magnet, NewBitmagnetProcessed, NewMagnet, NewTag, NewTransmissionProcessed}; +use crate::schema::{bitmagnet_processed, magnet_tags, magnets, tags, transmission_processed}; use crate::PostInfo; use color_eyre::eyre::{eyre, Result, WrapErr}; use diesel::prelude::*; @@ -129,10 +129,69 @@ impl Database { }) .collect::>(); - diesel::insert_into(magnets::table) - .values(&links) - .execute(&mut self.conn) - .wrap_err("Failed to save new magnet") + self.conn.transaction(|connection| { + // Insert magnets + let insert = diesel::insert_into(magnets::table) + .values(&links) + .execute(connection) + .wrap_err("Failed to save new magnet")?; + + // If we have tags, insert them and associate with the new magnets + // TODO: Can we get the list of inserted IDs above instead of re-querying? + if !post.tags.is_empty() && insert > 0 { + // Get the IDs of the newly inserted magnets + let new_magnet_ids: Vec = magnets::table + .select(magnets::id) + .filter(magnets::link.eq_any(&post.magnet_links)) + .filter(magnets::title.eq(&post.title)) + .load(connection) + .wrap_err("Failed to get IDs of newly inserted magnets")?; + + // For each tag, ensure it exists and get its ID + for tag_name in &post.tags { + // Insert tag if it doesn't exist + let new_tag = NewTag { name: tag_name }; + diesel::insert_into(tags::table) + .values(&new_tag) + .on_conflict_do_nothing() + .execute(connection) + .wrap_err("Failed to insert tag")?; + + // Get the tag ID + let tag_id = tags::table + .select(tags::id) + .filter(tags::name.eq(tag_name)) + .first::(connection) + .wrap_err("Failed to get tag ID")?; + + // Associate the tag with each new magnet + for magnet_id in &new_magnet_ids { + diesel::insert_into(magnet_tags::table) + .values(( + magnet_tags::magnet_id.eq(magnet_id), + magnet_tags::tag_id.eq(tag_id), + )) + .on_conflict_do_nothing() + .execute(connection) + .wrap_err("Failed to associate tag with magnet")?; + } + } + } + + Ok(insert) + }) + } + + /// Get tags for a specific magnet + pub fn get_tags_for_magnet(&mut self, magnet_id: i32) -> Result> { + let tag_names = tags::table + .inner_join(magnet_tags::table.on(tags::id.eq(magnet_tags::tag_id))) + .filter(magnet_tags::magnet_id.eq(magnet_id)) + .select(tags::name) + .load::(&mut self.conn) + .wrap_err("Failed to load tags for magnet")?; + + Ok(tag_names) } /// Get all magnets that have not been processed by a specific table @@ -195,6 +254,7 @@ mod tests { ], timestamp: Utc::now(), imdb_id: None, + tags: vec![], }; let expected_timestamp = post_info.timestamp.naive_utc(); @@ -234,6 +294,7 @@ mod tests { ], timestamp: Utc::now(), imdb_id: None, + tags: vec![], }; // First insertion should succeed and insert 2 links @@ -252,6 +313,7 @@ mod tests { ], timestamp: Utc::now(), imdb_id: Some("tt1234567".to_string()), + tags: vec!["test".to_string()], }; // Second insertion should succeed but only insert the new link @@ -274,6 +336,7 @@ mod tests { ], timestamp: Utc::now(), imdb_id: None, + tags: vec!["another_test".to_string()], }; // Third insertion should succeed but insert 0 links @@ -285,4 +348,75 @@ mod tests { let magnets = db.get_all_magnets().unwrap(); assert_eq!(magnets.len(), 3); } + + #[test] + fn test_store_and_retrieve_tags() { + let temp_dir = tempdir().unwrap(); + let db_path = temp_dir.path().join("test.db"); + + let mut db = Database::new(&db_path).unwrap(); + + // Create a post with tags + let post_info = PostInfo { + title: "Test Title with Tags".to_string(), + submitter: "test_user".to_string(), + subreddit: "test_subreddit".to_string(), + magnet_links: vec![ + "magnet:?xt=urn:btih:test_tag1".to_string(), + "magnet:?xt=urn:btih:test_tag2".to_string(), + ], + timestamp: Utc::now(), + imdb_id: None, + tags: vec!["action".to_string(), "comedy".to_string()], + }; + + // Insert the post with tags + let result = db.store_magnets(&post_info); + assert!(result.is_ok()); + assert_eq!(result.unwrap(), 2); // 2 links inserted + + // Get all magnets to find their IDs + let magnets = db.get_all_magnets().unwrap(); + assert_eq!(magnets.len(), 2); + + // Verify tags for each magnet + for magnet in magnets { + let tags = db.get_tags_for_magnet(magnet.id).unwrap(); + assert_eq!(tags.len(), 2); + assert!(tags.contains(&"action".to_string())); + assert!(tags.contains(&"comedy".to_string())); + } + + // Test adding a new magnet with some duplicate and some new tags + let post_info2 = PostInfo { + title: "Test Title with More Tags".to_string(), + submitter: "test_user".to_string(), + subreddit: "test_subreddit".to_string(), + magnet_links: vec!["magnet:?xt=urn:btih:test_tag3".to_string()], + timestamp: Utc::now(), + imdb_id: None, + tags: vec!["comedy".to_string(), "drama".to_string()], // One duplicate, one new + }; + + // Insert the second post + let result2 = db.store_magnets(&post_info2); + assert!(result2.is_ok()); + assert_eq!(result2.unwrap(), 1); // 1 link inserted + + // Get all magnets again + let magnets = db.get_all_magnets().unwrap(); + assert_eq!(magnets.len(), 3); + + // Find the new magnet (the one with test_tag3) + let new_magnet = magnets + .iter() + .find(|m| m.link == "magnet:?xt=urn:btih:test_tag3") + .unwrap(); + + // Verify tags for the new magnet + let tags = db.get_tags_for_magnet(new_magnet.id).unwrap(); + assert_eq!(tags.len(), 2); + assert!(tags.contains(&"comedy".to_string())); + assert!(tags.contains(&"drama".to_string())); + } } diff --git a/src/main.rs b/src/main.rs index 87dbfb5..c0d4709 100644 --- a/src/main.rs +++ b/src/main.rs @@ -28,6 +28,7 @@ pub struct PostInfo { pub subreddit: String, pub timestamp: DateTime, pub imdb_id: Option, + pub tags: Vec, } #[tokio::main] diff --git a/src/models.rs b/src/models.rs index 7c258e3..b255cb6 100644 --- a/src/models.rs +++ b/src/models.rs @@ -1,4 +1,4 @@ -use crate::schema::{bitmagnet_processed, magnets, transmission_processed}; +use crate::schema::{bitmagnet_processed, magnet_tags, magnets, tags, transmission_processed}; use chrono::NaiveDateTime; use diesel::prelude::*; @@ -28,6 +28,33 @@ pub struct NewMagnet<'a> { pub imdb_id: Option<&'a str>, } +#[allow(dead_code)] +#[derive(Queryable, Selectable, Identifiable)] +#[diesel(table_name = tags)] +#[diesel(check_for_backend(diesel::sqlite::Sqlite))] +pub struct Tag { + pub id: i32, + pub name: String, +} + +#[derive(Insertable)] +#[diesel(table_name = tags)] +#[diesel(check_for_backend(diesel::sqlite::Sqlite))] +pub struct NewTag<'a> { + pub name: &'a str, +} + +#[derive(Identifiable, Selectable, Queryable, Associations)] +#[diesel(belongs_to(Magnet))] +#[diesel(belongs_to(Tag))] +#[diesel(table_name = magnet_tags)] +#[diesel(primary_key(magnet_id, tag_id))] +#[diesel(check_for_backend(diesel::sqlite::Sqlite))] +pub struct MagnetTag { + pub magnet_id: i32, + pub tag_id: i32, +} + #[allow(dead_code)] #[derive(Queryable, Selectable)] #[diesel(table_name = bitmagnet_processed)] diff --git a/src/schema.rs b/src/schema.rs index f993232..d1014fd 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -8,6 +8,13 @@ diesel::table! { } } +diesel::table! { + magnet_tags (magnet_id, tag_id) { + magnet_id -> Integer, + tag_id -> Integer, + } +} + diesel::table! { magnets (id) { id -> Integer, @@ -20,6 +27,13 @@ diesel::table! { } } +diesel::table! { + tags (id) { + id -> Integer, + name -> Text, + } +} + diesel::table! { transmission_processed (id) { id -> Integer, @@ -29,6 +43,14 @@ diesel::table! { } diesel::joinable!(bitmagnet_processed -> magnets (magnet_id)); +diesel::joinable!(magnet_tags -> magnets (magnet_id)); +diesel::joinable!(magnet_tags -> tags (tag_id)); diesel::joinable!(transmission_processed -> magnets (magnet_id)); -diesel::allow_tables_to_appear_in_same_query!(bitmagnet_processed, magnets, transmission_processed,); +diesel::allow_tables_to_appear_in_same_query!( + bitmagnet_processed, + magnet_tags, + magnets, + tags, + transmission_processed, +);