store_sqlite/
claim_store.rs1use std::collections::HashSet;
4use std::sync::Mutex;
5
6use chrono::Utc;
7use rusqlite::{params, Connection, OptionalExtension, TransactionBehavior};
8use strsim::levenshtein;
9use uuid::Uuid;
10
11use substrate_core::claim_port::{ClaimPort, WorkItem, WorkItemState};
12
13use crate::error::StoreError;
14use crate::schema;
15
16const JACCARD_THRESHOLD: f64 = 0.75;
18const LEVENSHTEIN_THRESHOLD: f64 = 0.85;
20
21pub struct SqliteClaimStore {
23 conn: Mutex<Connection>,
24}
25
26impl SqliteClaimStore {
27 pub fn open(path: &str) -> Result<Self, StoreError> {
29 let conn = Connection::open(path)?;
30 schema::init(&conn)?;
31 Ok(Self {
32 conn: Mutex::new(conn),
33 })
34 }
35
36 pub fn open_in_memory() -> Result<Self, StoreError> {
38 let conn = Connection::open_in_memory()?;
39 schema::init(&conn)?;
40 Ok(Self {
41 conn: Mutex::new(conn),
42 })
43 }
44}
45
46fn token_jaccard(a: &str, b: &str) -> f64 {
47 let a_lower = a.to_lowercase();
48 let b_lower = b.to_lowercase();
49 let ta: HashSet<&str> = a_lower.split_whitespace().collect();
50 let tb: HashSet<&str> = b_lower.split_whitespace().collect();
51 if ta.is_empty() && tb.is_empty() {
52 return 1.0;
53 }
54 let inter = ta.intersection(&tb).count();
55 let union = ta.union(&tb).count();
56 if union == 0 {
57 return 0.0;
58 }
59 inter as f64 / union as f64
60}
61
62fn levenshtein_similarity(a: &str, b: &str) -> f64 {
63 let dist = levenshtein(a, b);
64 let max_len = a.chars().count().max(b.chars().count()).max(1);
65 1.0 - (dist as f64 / max_len as f64)
66}
67
68pub fn bodies_are_near_duplicate(a: &str, b: &str) -> bool {
70 token_jaccard(a, b) >= JACCARD_THRESHOLD
71 || levenshtein_similarity(a, b) >= LEVENSHTEIN_THRESHOLD
72}
73
74impl ClaimPort for SqliteClaimStore {
75 type Error = StoreError;
76
77 fn enqueue(&self, queue: &str, body: &str) -> Result<Uuid, StoreError> {
78 if self.is_near_duplicate(queue, body)? {
79 return Err(StoreError::Duplicate(format!(
80 "near-duplicate in queue {queue}"
81 )));
82 }
83 let id = Uuid::new_v4();
84 let conn = self.conn.lock().unwrap();
85 conn.execute(
86 "INSERT INTO work_queue (id, queue, body, state, claimed_by, created_at) \
87 VALUES (?1, ?2, ?3, 'pending', NULL, ?4)",
88 params![id.to_string(), queue, body, Utc::now().to_rfc3339()],
89 )?;
90 Ok(id)
91 }
92
93 fn claim_next(&self, queue: &str, worker_id: &str) -> Result<Option<WorkItem>, StoreError> {
94 let mut conn = self.conn.lock().unwrap();
95 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
96
97 let row: Option<(String, String, String)> = tx
98 .query_row(
99 "SELECT id, queue, body FROM work_queue \
100 WHERE queue=?1 AND state='pending' \
101 ORDER BY created_at ASC LIMIT 1",
102 params![queue],
103 |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
104 )
105 .optional()?;
106
107 let Some((id, q, body)) = row else {
108 tx.rollback()?;
109 return Ok(None);
110 };
111
112 let updated = tx.execute(
113 "UPDATE work_queue SET state='claimed', claimed_by=?2 \
114 WHERE id=?1 AND state='pending'",
115 params![id, worker_id],
116 )?;
117
118 if updated != 1 {
119 tx.rollback()?;
120 return Ok(None);
121 }
122
123 tx.commit()?;
124
125 Ok(Some(WorkItem {
126 id: Uuid::parse_str(&id).unwrap_or_else(|_| Uuid::new_v4()),
127 queue: q,
128 body,
129 state: WorkItemState::Claimed,
130 claimed_by: Some(worker_id.to_string()),
131 }))
132 }
133
134 fn is_near_duplicate(&self, queue: &str, body: &str) -> Result<bool, StoreError> {
135 let conn = self.conn.lock().unwrap();
136 let mut stmt = conn.prepare(
137 "SELECT body FROM work_queue \
138 WHERE queue=?1 AND state IN ('pending', 'claimed')",
139 )?;
140 let rows = stmt.query_map(params![queue], |row| row.get::<_, String>(0))?;
141 for row in rows {
142 let existing = row?;
143 if bodies_are_near_duplicate(body, &existing) {
144 return Ok(true);
145 }
146 }
147 Ok(false)
148 }
149}
150
151#[cfg(test)]
152mod tests {
153 use super::*;
154
155 #[test]
156 fn near_duplicate_detection() {
157 assert!(bodies_are_near_duplicate(
158 "refactor auth module login flow",
159 "refactor auth module login flo"
160 ));
161 assert!(bodies_are_near_duplicate(
162 "fix the login bug in auth module",
163 "fix the login bug in auth modul"
164 ));
165 assert!(!bodies_are_near_duplicate(
166 "implement payment gateway",
167 "write unit tests for scheduler"
168 ));
169 }
170}