use rusqlite::{Connection, Result};
use serde_json::Value;
use std::path::Path;
use std::sync::Mutex;
use std::time::{SystemTime, UNIX_EPOCH};
pub fn extract_dedup_uri(status: &Value) -> Option<&str> {
if let Some(reblog) = status.get("reblog") {
if !reblog.is_null() {
return reblog.get("uri")?.as_str();
}
}
status.get("uri")?.as_str()
}
pub struct SeenUriStore {
conn: Mutex<Connection>,
}
impl SeenUriStore {
pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
let conn = Connection::open(path)?;
conn.pragma_update(None, "journal_mode", "WAL")?;
conn.execute(
"CREATE TABLE IF NOT EXISTS seen_uris (
uri TEXT PRIMARY KEY,
first_seen INTEGER NOT NULL
)",
[],
)?;
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_first_seen ON seen_uris(first_seen)",
[],
)?;
Ok(Self {
conn: Mutex::new(conn),
})
}
pub fn is_seen(&self, uri: &str) -> Result<bool> {
let conn = self.conn.lock().expect(
"SeenUriStore mutex poisoned. This indicates a panic occurred \
while holding the lock. The application should be restarted.",
);
let mut stmt = conn.prepare_cached("SELECT 1 FROM seen_uris WHERE uri = ?")?;
let exists = stmt.exists([uri])?;
Ok(exists)
}
pub fn mark_seen(&self, uri: &str) -> Result<()> {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("Time went backwards")
.as_secs() as i64;
let conn = self.conn.lock().expect(
"SeenUriStore mutex poisoned. This indicates a panic occurred \
while holding the lock. The application should be restarted.",
);
conn.execute(
"INSERT OR IGNORE INTO seen_uris (uri, first_seen) VALUES (?, ?)",
(uri, now),
)?;
Ok(())
}
pub fn check_and_mark(&self, uri: &str) -> Result<bool> {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("Time went backwards")
.as_secs() as i64;
let conn = self.conn.lock().expect(
"SeenUriStore mutex poisoned. This indicates a panic occurred \
while holding the lock. The application should be restarted.",
);
let rows_changed = conn.execute(
"INSERT OR IGNORE INTO seen_uris (uri, first_seen) VALUES (?, ?)",
(uri, now),
)?;
Ok(rows_changed == 0)
}
pub fn cleanup(&self, max_age_secs: u64) -> Result<usize> {
let conn = self.conn.lock().expect(
"SeenUriStore mutex poisoned. This indicates a panic occurred \
while holding the lock. The application should be restarted.",
);
let removed = if max_age_secs == 0 {
conn.execute("DELETE FROM seen_uris", [])?
} else {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("Time went backwards")
.as_secs() as i64;
let cutoff = now - (max_age_secs as i64);
conn.execute("DELETE FROM seen_uris WHERE first_seen < ?", [cutoff])?
};
Ok(removed)
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_in_memory_store() {
let store = SeenUriStore::open(":memory:").unwrap();
let uri = "https://example.com/status/123";
assert!(!store.is_seen(uri).unwrap());
store.mark_seen(uri).unwrap();
assert!(store.is_seen(uri).unwrap());
}
#[test]
fn test_check_and_mark_atomic() {
let store = SeenUriStore::open(":memory:").unwrap();
let uri = "https://example.com/status/456";
assert!(!store.check_and_mark(uri).unwrap());
assert!(store.check_and_mark(uri).unwrap());
assert!(store.is_seen(uri).unwrap());
}
#[test]
fn test_extract_uri_from_regular_status() {
let status = json!({
"id": "123456",
"uri": "https://mastodon.social/users/testuser/statuses/123456",
"content": "<p>Hello, world!</p>"
});
let uri = extract_dedup_uri(&status);
assert_eq!(
uri,
Some("https://mastodon.social/users/testuser/statuses/123456")
);
}
#[test]
fn test_extract_uri_from_reblog() {
let status = json!({
"id": "789012",
"uri": "https://mastodon.social/users/booster/statuses/789012",
"reblog": {
"id": "123456",
"uri": "https://fosstodon.org/users/original/statuses/123456",
"content": "<p>Original post</p>"
}
});
let uri = extract_dedup_uri(&status);
assert_eq!(
uri,
Some("https://fosstodon.org/users/original/statuses/123456")
);
}
#[test]
fn test_extract_uri_with_null_reblog() {
let status = json!({
"id": "123456",
"uri": "https://mastodon.social/users/testuser/statuses/123456",
"reblog": null
});
let uri = extract_dedup_uri(&status);
assert_eq!(
uri,
Some("https://mastodon.social/users/testuser/statuses/123456")
);
}
#[test]
fn test_extract_uri_missing() {
let status = json!({
"id": "123456",
"content": "<p>No URI field</p>"
});
let uri = extract_dedup_uri(&status);
assert_eq!(uri, None);
}
}