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.
This commit is contained in:
Marc Plano-Lesay 2025-05-04 15:41:02 +10:00
parent f19e02988f
commit c3764c125a
Signed by: kernald
GPG key ID: 66A41B08CC62A6CF
13 changed files with 240 additions and 11 deletions

View file

@ -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.)

View file

@ -0,0 +1,2 @@
DROP TABLE magnet_tags;
DROP TABLE tags;

View file

@ -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)
);

View file

@ -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::<BitmagnetProcessedTable>(magnet.id)?;
processed_magnets.push(magnet);
}

View file

@ -34,6 +34,7 @@ impl BitmagnetClient {
magnet: &str,
published_at: &DateTime<Utc>,
imdb_id: &Option<String>,
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)

View file

@ -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::<TransmissionProcessedTable>(magnet.id)?;
processed_magnets.push(magnet);
}

View file

@ -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()
};

View file

@ -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

View file

@ -21,6 +21,8 @@ pub struct SourceConfig {
pub username: String,
pub title_filter: Option<String>,
pub imdb_id: Option<String>,
#[serde(default)]
pub tags: Vec<String>,
}
/// Main application configuration

146
src/db.rs
View file

@ -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::<Vec<NewMagnet>>();
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<i32> = 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::<i32>(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<Vec<String>> {
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::<String>(&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()));
}
}

View file

@ -28,6 +28,7 @@ pub struct PostInfo {
pub subreddit: String,
pub timestamp: DateTime<Utc>,
pub imdb_id: Option<String>,
pub tags: Vec<String>,
}
#[tokio::main]

View file

@ -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)]

View file

@ -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,
);