use rusqlite::Connection; use crate::core::backoff::compute_next_attempt_at; use crate::core::error::Result; use crate::core::time::now_ms; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum NoteableType { Issue, MergeRequest, } impl NoteableType { pub fn as_str(&self) -> &'static str { match self { Self::Issue => "Issue", Self::MergeRequest => "MergeRequest", } } pub fn parse(s: &str) -> Option { match s { "Issue" => Some(Self::Issue), "MergeRequest" => Some(Self::MergeRequest), _ => None, } } } pub struct PendingFetch { pub project_id: i64, pub noteable_type: NoteableType, pub noteable_iid: i64, pub attempt_count: i32, } pub fn queue_discussion_fetch( conn: &Connection, project_id: i64, noteable_type: NoteableType, noteable_iid: i64, ) -> Result<()> { conn.execute( "INSERT INTO pending_discussion_fetches (project_id, noteable_type, noteable_iid, queued_at) VALUES (?1, ?2, ?3, ?4) ON CONFLICT(project_id, noteable_type, noteable_iid) DO UPDATE SET queued_at = excluded.queued_at, attempt_count = 0, last_attempt_at = NULL, last_error = NULL, next_attempt_at = NULL", rusqlite::params![project_id, noteable_type.as_str(), noteable_iid, now_ms()], )?; Ok(()) } pub fn get_pending_fetches(conn: &Connection, limit: usize) -> Result> { let now = now_ms(); let mut stmt = conn.prepare( "SELECT project_id, noteable_type, noteable_iid, attempt_count FROM pending_discussion_fetches WHERE next_attempt_at IS NULL OR next_attempt_at <= ?1 ORDER BY queued_at ASC LIMIT ?2", )?; let rows = stmt .query_map(rusqlite::params![now, limit as i64], |row| { Ok(( row.get::<_, i64>(0)?, row.get::<_, String>(1)?, row.get::<_, i64>(2)?, row.get::<_, i32>(3)?, )) })? .collect::, _>>()?; let mut results = Vec::with_capacity(rows.len()); for (project_id, nt_str, noteable_iid, attempt_count) in rows { let noteable_type = NoteableType::parse(&nt_str).ok_or_else(|| { crate::core::error::LoreError::Other(format!( "Invalid noteable_type in pending_discussion_fetches: {}", nt_str )) })?; results.push(PendingFetch { project_id, noteable_type, noteable_iid, attempt_count, }); } Ok(results) } pub fn complete_fetch( conn: &Connection, project_id: i64, noteable_type: NoteableType, noteable_iid: i64, ) -> Result<()> { conn.execute( "DELETE FROM pending_discussion_fetches WHERE project_id = ?1 AND noteable_type = ?2 AND noteable_iid = ?3", rusqlite::params![project_id, noteable_type.as_str(), noteable_iid], )?; Ok(()) } pub fn record_fetch_error( conn: &Connection, project_id: i64, noteable_type: NoteableType, noteable_iid: i64, error: &str, ) -> Result<()> { let now = now_ms(); let attempt_count: i64 = conn.query_row( "SELECT attempt_count FROM pending_discussion_fetches WHERE project_id = ?1 AND noteable_type = ?2 AND noteable_iid = ?3", rusqlite::params![project_id, noteable_type.as_str(), noteable_iid], |row| row.get(0), )?; let new_attempt = attempt_count + 1; let next_at = compute_next_attempt_at(now, new_attempt); conn.execute( "UPDATE pending_discussion_fetches SET attempt_count = ?1, last_attempt_at = ?2, last_error = ?3, next_attempt_at = ?4 WHERE project_id = ?5 AND noteable_type = ?6 AND noteable_iid = ?7", rusqlite::params![ new_attempt, now, error, next_at, project_id, noteable_type.as_str(), noteable_iid ], )?; Ok(()) } #[cfg(test)] mod tests { use super::*; fn setup_db() -> Connection { let conn = Connection::open_in_memory().unwrap(); conn.execute_batch(" CREATE TABLE projects ( id INTEGER PRIMARY KEY, gitlab_project_id INTEGER UNIQUE NOT NULL, path_with_namespace TEXT NOT NULL, default_branch TEXT, web_url TEXT, created_at INTEGER, updated_at INTEGER, raw_payload_id INTEGER ); INSERT INTO projects (id, gitlab_project_id, path_with_namespace) VALUES (1, 100, 'group/project'); CREATE TABLE pending_discussion_fetches ( project_id INTEGER NOT NULL REFERENCES projects(id), noteable_type TEXT NOT NULL, noteable_iid INTEGER NOT NULL, queued_at INTEGER NOT NULL, attempt_count INTEGER NOT NULL DEFAULT 0, last_attempt_at INTEGER, last_error TEXT, next_attempt_at INTEGER, PRIMARY KEY(project_id, noteable_type, noteable_iid) ); CREATE INDEX idx_pending_discussions_next_attempt ON pending_discussion_fetches(next_attempt_at); ").unwrap(); conn } #[test] fn test_queue_and_get() { let conn = setup_db(); queue_discussion_fetch(&conn, 1, NoteableType::Issue, 42).unwrap(); let fetches = get_pending_fetches(&conn, 100).unwrap(); assert_eq!(fetches.len(), 1); assert_eq!(fetches[0].project_id, 1); assert_eq!(fetches[0].noteable_type, NoteableType::Issue); assert_eq!(fetches[0].noteable_iid, 42); assert_eq!(fetches[0].attempt_count, 0); } #[test] fn test_requeue_resets_backoff() { let conn = setup_db(); queue_discussion_fetch(&conn, 1, NoteableType::Issue, 42).unwrap(); record_fetch_error(&conn, 1, NoteableType::Issue, 42, "network error").unwrap(); let attempt: i32 = conn .query_row( "SELECT attempt_count FROM pending_discussion_fetches WHERE noteable_iid = 42", [], |r| r.get(0), ) .unwrap(); assert_eq!(attempt, 1); queue_discussion_fetch(&conn, 1, NoteableType::Issue, 42).unwrap(); let attempt: i32 = conn .query_row( "SELECT attempt_count FROM pending_discussion_fetches WHERE noteable_iid = 42", [], |r| r.get(0), ) .unwrap(); assert_eq!(attempt, 0); } #[test] fn test_backoff_respected() { let conn = setup_db(); queue_discussion_fetch(&conn, 1, NoteableType::Issue, 42).unwrap(); conn.execute( "UPDATE pending_discussion_fetches SET next_attempt_at = 9999999999999 WHERE noteable_iid = 42", [], ).unwrap(); let fetches = get_pending_fetches(&conn, 100).unwrap(); assert!(fetches.is_empty()); } #[test] fn test_complete_removes() { let conn = setup_db(); queue_discussion_fetch(&conn, 1, NoteableType::Issue, 42).unwrap(); complete_fetch(&conn, 1, NoteableType::Issue, 42).unwrap(); let count: i64 = conn .query_row("SELECT COUNT(*) FROM pending_discussion_fetches", [], |r| { r.get(0) }) .unwrap(); assert_eq!(count, 0); } #[test] fn test_error_increments_attempts() { let conn = setup_db(); queue_discussion_fetch(&conn, 1, NoteableType::MergeRequest, 10).unwrap(); record_fetch_error(&conn, 1, NoteableType::MergeRequest, 10, "timeout").unwrap(); let (attempt, error): (i32, Option) = conn.query_row( "SELECT attempt_count, last_error FROM pending_discussion_fetches WHERE noteable_iid = 10", [], |r| Ok((r.get(0)?, r.get(1)?)), ).unwrap(); assert_eq!(attempt, 1); assert_eq!(error, Some("timeout".to_string())); let next_at: Option = conn .query_row( "SELECT next_attempt_at FROM pending_discussion_fetches WHERE noteable_iid = 10", [], |r| r.get(0), ) .unwrap(); assert!(next_at.is_some()); } #[test] fn test_noteable_type_parse() { assert_eq!(NoteableType::parse("Issue"), Some(NoteableType::Issue)); assert_eq!( NoteableType::parse("MergeRequest"), Some(NoteableType::MergeRequest) ); assert_eq!(NoteableType::parse("invalid"), None); } }