path: root/src/db.rs

use chrono::{DateTime, Utc};
use rusqlite::Connection;
use rusqlite::OptionalExtension;

pub trait Db {
    fn get_image(&self, registry: &str, image: &str) -> Option<(i64, DateTime<Utc>)>;
    fn insert_image(&self, registry: &str, image: &str, expires_at: &DateTime<Utc>) -> i64;
    fn get_tags_sorted(&self, image_id: i64) -> Vec<String>;
    fn delete_tag(&self, image_id: i64, tag: &str);
    fn insert_tag(&self, image_id: i64, tag: &str, fetched_at: &DateTime<Utc>);
    fn set_expires_at(&self, image_id: i64, expires_at: &DateTime<Utc>);
    fn get_tag_digest(&self, image_id: i64, tag: &str) -> Option<String>;
    fn update_tag_digest(&self, image_id: i64, tag: &str, digest: &str, fetched_at: &DateTime<Utc>);
}

pub struct SqliteDb {
    conn: Connection,
}

impl SqliteDb {
    pub fn new(db_path: &std::path::PathBuf) -> Self {
        let conn = Connection::open(db_path).unwrap();
        conn.execute("
            CREATE TABLE IF NOT EXISTS migrations (
                id INTEGER PRIMARY KEY NOT NULL
            )
        ", ()).unwrap();

        let newest_migration: u32 = conn.query_row("
            SELECT MAX(id) FROM migrations
        ", [], |row| row.get::<_, Option<u32>>(0)).unwrap().unwrap_or(0);

        if newest_migration < 1 {
            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (1, )).unwrap();
        }

        if newest_migration < 2 {
            conn.execute("
                CREATE TABLE known_images (
                    id INTEGER PRIMARY KEY NOT NULL,
                    registry TEXT NOT NULL,
                    image TEXT NOT NULL,
                    tag TEXT NOT NULL
                )
            ", ()).unwrap();

            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (2, )).unwrap();
        }

        if newest_migration < 3 {
            conn.execute("
                ALTER TABLE known_images ADD COLUMN
                    discovered DATETIME NOT NULL
            ", ()).unwrap();

            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (3, )).unwrap();
        }

        if newest_migration < 4 {
            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (4, )).unwrap();
        }

        if newest_migration < 5 {
            conn.execute("
                CREATE TABLE images (
                    id INTEGER PRIMARY KEY NOT NULL,
                    registry TEXT NOT NULL,
                    image TEXT NOT NULL,
                    last_checked DATETIME NOT NULL
                )
            ", ()).unwrap();

            conn.execute("
                CREATE UNIQUE INDEX images__registry_image
                    ON images(registry, image)
            ", ()).unwrap();

            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (5, )).unwrap();
        }

        if newest_migration < 6 {
            conn.execute("
                CREATE TABLE tags (
                    id INTEGER PRIMARY KEY,
                    image_id INTEGER NOT NULL REFERENCES images(id),
                    tag TEXT NOT NULL,
                    digest TEXT,
                    fetched_at DATETIME NOT NULL
                )
            ", ()).unwrap();

            conn.execute("
                CREATE UNIQUE INDEX tags__image_id_tag
                    ON tags(image_id, tag)
            ", ()).unwrap();

            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (6, )).unwrap();
        }

        if newest_migration < 7 {
            conn.execute("ALTER TABLE images RENAME COLUMN last_checked TO expires_at", ()).unwrap();
            conn.execute("UPDATE images SET expires_at = datetime(expires_at, '+1440 minutes')", ()).unwrap();
            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (7, )).unwrap();
        }

        return SqliteDb { conn };
    }
}

impl Db for SqliteDb {
    fn get_image(&self, registry: &str, image: &str) -> Option<(i64, DateTime<Utc>)> {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        return self.conn.query_row("
            SELECT id, expires_at FROM images
                WHERE registry = ?1 AND image = ?2
        ", (registry, image), |row| Ok((
            row.get::<_, i64>(0)?,
            row.get::<_, DateTime<Utc>>(1)?,
        ))).optional().unwrap();
    }

    fn insert_image(&self, registry: &str, image: &str, expires_at: &DateTime<Utc>) -> i64 {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        self.conn.execute("
            INSERT INTO images(registry, image, expires_at) VALUES (?1, ?2, ?3)
        ", (registry, image, expires_at)).unwrap();
        return self.conn.last_insert_rowid();
    }

    fn get_tags_sorted(&self, image_id: i64) -> Vec<String> {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        let mut stmt = self.conn.prepare("SELECT tag FROM tags WHERE image_id = ?1 ORDER BY tag").unwrap();
        let mut rows = stmt.query((image_id,)).unwrap();
        let mut tags = vec![];
        while let Some(row) = rows.next().unwrap() {
            tags.push(row.get(0).unwrap());
        }
        return tags;
    }

    fn delete_tag(&self, image_id: i64, tag: &str) {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        self.conn.execute("DELETE FROM tags WHERE image_id = ?1 AND tag = ?2",
            (image_id, tag)).unwrap();
    }

    fn insert_tag(&self, image_id: i64, tag: &str, fetched_at: &DateTime<Utc>) {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        self.conn.execute("INSERT INTO tags(image_id, tag, fetched_at) VALUES (?1, ?2, ?3)",
            (image_id, tag, fetched_at)).unwrap();
    }

    fn set_expires_at(&self, image_id: i64, expires_at: &DateTime<Utc>) {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        self.conn.execute("UPDATE images SET expires_at = ?1 WHERE id = ?2", (expires_at, image_id)).unwrap();
    }

    fn get_tag_digest(&self, image_id: i64, tag: &str) -> Option<String> {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        return self.conn.query_row("
            SELECT digest FROM tags WHERE image_id = ?1 AND tag = ?2 AND digest IS NOT NULL
        ", (image_id, tag), |row| row.get::<_, String>(0)).optional().unwrap();
    }

    fn update_tag_digest(&self, image_id: i64, tag: &str, digest: &str, fetched_at: &DateTime<Utc>) {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        self.conn.execute("UPDATE tags SET digest = ?1, fetched_at = ?2 WHERE image_id = ?3 AND tag = ?4",
            (digest, fetched_at, image_id, tag)).unwrap();
    }
}

#[cfg(test)]
pub struct StubDb {
    next_id: std::cell::RefCell<i64>,
    images: std::cell::RefCell<Vec<(String, String, i64, DateTime<Utc>)>>,
    tags: std::cell::RefCell<Vec<(i64, String, Option<String>, DateTime<Utc>)>>,
}

#[cfg(test)]
impl Default for StubDb {
    fn default() -> Self {
        return StubDb {
            next_id: std::cell::RefCell::new(1),
            images: std::cell::RefCell::new(vec![]),
            tags: std::cell::RefCell::new(vec![]),
        };
    }
}

#[cfg(test)]
impl Db for StubDb {
    fn get_image(&self, registry: &str, image: &str) -> Option<(i64, DateTime<Utc>)> {
        for (r, i, id, expires_at) in self.images.borrow().iter() {
            if r == registry && i == image {
                return Some((*id, *expires_at));
            }
        }
        return None;
    }

    fn insert_image(&self, registry: &str, image: &str, expires_at: &DateTime<Utc>) -> i64 {
        let id = *self.next_id.borrow();
        *self.next_id.borrow_mut() += 1;
        self.images.borrow_mut().push((registry.to_string(), image.to_string(), id, expires_at.clone()));
        return id;
    }

    fn get_tags_sorted(&self, image_id: i64) -> Vec<String> {
        let mut result = vec![];
        for (id, tag, _, _) in self.tags.borrow().iter() {
            if *id == image_id {
                result.push(tag.clone());
            }
        }
        result.sort();
        return result;
    }

    fn delete_tag(&self, image_id: i64, tag: &str) {
        self.tags.borrow_mut().retain(|(id, t, _, _)| !(*id == image_id && t == tag));
    }

    fn insert_tag(&self, image_id: i64, tag: &str, fetched_at: &DateTime<Utc>) {
        self.tags.borrow_mut().push((image_id, tag.to_string(), None, fetched_at.clone()));
    }

    fn set_expires_at(&self, image_id: i64, expires_at: &DateTime<Utc>) {
        for (_, _, id, ea) in self.images.borrow_mut().iter_mut() {
            if *id == image_id {
                *ea = expires_at.clone();
                return;
            }
        }
    }

    fn get_tag_digest(&self, image_id: i64, tag: &str) -> Option<String> {
        for (id, t, digest, _) in self.tags.borrow().iter() {
            if *id == image_id && t == tag {
                return digest.clone();
            }
        }
        return None;
    }

    fn update_tag_digest(&self, image_id: i64, tag: &str, digest: &str, fetched_at: &DateTime<Utc>) {
        for (id, t, d, fa) in self.tags.borrow_mut().iter_mut() {
            if *id == image_id && t == tag {
                *d = Some(digest.to_string());
                *fa = fetched_at.clone();
                return;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::TimeZone;

    fn test_insert_image_returns_incrementing_ids(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let id1 = db.insert_image("docker.io", "nginx", &now);
        let id2 = db.insert_image("docker.io", "redis", &now);
        assert!(id2 > id1);
    }

    fn test_get_image_returns_none_for_unknown(db: &dyn Db) {
        assert_eq!(db.get_image("docker.io", "unknown"), None);
    }

    fn test_get_image_returns_inserted(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let id = db.insert_image("docker.io", "nginx", &now);
        let (got_id, got_time) = db.get_image("docker.io", "nginx").unwrap();
        assert_eq!(got_id, id);
        assert_eq!(got_time, now);
    }

    fn test_get_tags_returns_empty_for_no_tags(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let id = db.insert_image("docker.io", "nginx", &now);
        assert_eq!(db.get_tags_sorted(id), Vec::<String>::new());
    }

    fn test_get_tags_returns_sorted(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let id = db.insert_image("docker.io", "nginx", &now);
        db.insert_tag(id, "2.0", &now);
        db.insert_tag(id, "1.0", &now);
        db.insert_tag(id, "latest", &now);
        assert_eq!(db.get_tags_sorted(id), vec!["1.0", "2.0", "latest"]);
    }

    fn test_delete_tag_removes_tag(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let id = db.insert_image("docker.io", "nginx", &now);
        db.insert_tag(id, "1.0", &now);
        db.insert_tag(id, "2.0", &now);
        db.delete_tag(id, "1.0");
        assert_eq!(db.get_tags_sorted(id), vec!["2.0"]);
    }

    fn test_set_expires_at(db: &dyn Db) {
        let t1 = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let t2 = Utc.with_ymd_and_hms(2000, 1, 1, 1, 0, 0).unwrap();
        let id = db.insert_image("docker.io", "nginx", &t1);
        db.set_expires_at(id, &t2);
        assert_eq!(db.get_image("docker.io", "nginx").unwrap().1, t2);
    }

    fn test_get_tag_digest_returns_none_when_unset(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let id = db.insert_image("docker.io", "nginx", &now);
        db.insert_tag(id, "1.0", &now);
        assert_eq!(db.get_tag_digest(id, "1.0"), None);
    }

    fn test_update_tag_digest(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let id = db.insert_image("docker.io", "nginx", &now);
        db.insert_tag(id, "1.0", &now);
        db.update_tag_digest(id, "1.0", "sha256:abc", &now);
        assert_eq!(db.get_tag_digest(id, "1.0"), Some("sha256:abc".to_string()));
    }

    #[test]
    fn conformance_stub_insert_image_ids() {
        test_insert_image_returns_incrementing_ids(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_insert_image_ids() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_insert_image_returns_incrementing_ids(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_get_image_unknown() {
        test_get_image_returns_none_for_unknown(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_get_image_unknown() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_get_image_returns_none_for_unknown(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_get_image_inserted() {
        test_get_image_returns_inserted(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_get_image_inserted() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_get_image_returns_inserted(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_get_tags_empty() {
        test_get_tags_returns_empty_for_no_tags(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_get_tags_empty() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_get_tags_returns_empty_for_no_tags(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_get_tags_sorted() {
        test_get_tags_returns_sorted(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_get_tags_sorted() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_get_tags_returns_sorted(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_delete_tag() {
        test_delete_tag_removes_tag(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_delete_tag() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_delete_tag_removes_tag(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_set_expires_at() {
        test_set_expires_at(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_set_expires_at() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_set_expires_at(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_get_tag_digest_none() {
        test_get_tag_digest_returns_none_when_unset(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_get_tag_digest_none() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_get_tag_digest_returns_none_when_unset(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_update_tag_digest() {
        test_update_tag_digest(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_update_tag_digest() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_update_tag_digest(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }
}