gitlore/src/ingestion/discussion_queue.rs

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<Self> {
        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<Vec<PendingFetch>> {
    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::<std::result::Result<Vec<_>, _>>()?;

    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<String>) = 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<i64> = 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);
    }
}