path: root/src/db.rs

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

#[derive(Debug, PartialEq)]
pub struct Image {
    pub id: i64,
    pub registry: String,
    pub image: String,
    pub expires_at: DateTime<Utc>,
}

pub trait Db {
    fn get_image(&self, registry: &str, image: &str) -> Option<Image>;
    fn insert_image(&self, image: &mut Image);
    fn get_expired_images(&self, now: &DateTime<Utc>) -> Vec<Image>;
    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<Image> {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        return self.conn.query_row("
            SELECT id, registry, image, expires_at FROM images
                WHERE registry = ?1 AND image = ?2
        ", (registry, image), |row| Ok(Image {
            id: row.get(0)?,
            registry: row.get(1)?,
            image: row.get(2)?,
            expires_at: row.get(3)?,
        })).optional().unwrap();
    }

    fn insert_image(&self, image: &mut Image) {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        self.conn.execute("
            INSERT INTO images(registry, image, expires_at) VALUES (?1, ?2, ?3)
        ", (&image.registry, &image.image, &image.expires_at)).unwrap();
        image.id = self.conn.last_insert_rowid();
    }

    fn get_expired_images(&self, now: &DateTime<Utc>) -> Vec<Image> {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        let mut stmt = self.conn.prepare("
            SELECT id, registry, image, expires_at FROM images WHERE expires_at <= ?1
        ").unwrap();
        let mut rows = stmt.query((now,)).unwrap();
        let mut images = vec![];
        while let Some(row) = rows.next().unwrap() {
            images.push(Image {
                id: row.get(0).unwrap(),
                registry: row.get(1).unwrap(),
                image: row.get(2).unwrap(),
                expires_at: row.get(3).unwrap(),
            });
        }
        return images;
    }

    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<Image>>,
    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<Image> {
        for img in self.images.borrow().iter() {
            if img.registry == registry && img.image == image {
                return Some(Image {
                    id: img.id,
                    registry: img.registry.clone(),
                    image: img.image.clone(),
                    expires_at: img.expires_at,
                });
            }
        }
        return None;
    }

    fn insert_image(&self, image: &mut Image) {
        let id = *self.next_id.borrow();
        *self.next_id.borrow_mut() += 1;
        image.id = id;
        self.images.borrow_mut().push(Image {
            id,
            registry: image.registry.clone(),
            image: image.image.clone(),
            expires_at: image.expires_at,
        });
    }

    fn get_expired_images(&self, now: &DateTime<Utc>) -> Vec<Image> {
        let mut result = vec![];
        for img in self.images.borrow().iter() {
            if img.expires_at <= *now {
                result.push(Image {
                    id: img.id,
                    registry: img.registry.clone(),
                    image: img.image.clone(),
                    expires_at: img.expires_at,
                });
            }
        }
        return result;
    }

    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 img in self.images.borrow_mut().iter_mut() {
            if img.id == image_id {
                img.expires_at = *expires_at;
                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 mut img1 = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        let mut img2 = Image { id: 0, registry: "docker.io".into(), image: "redis".into(), expires_at: now };
        db.insert_image(&mut img1);
        db.insert_image(&mut img2);
        assert!(img2.id > img1.id);
    }

    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 mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        let got = db.get_image("docker.io", "nginx").unwrap();
        assert_eq!(got.id, img.id);
        assert_eq!(got.expires_at, 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 mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        assert_eq!(db.get_tags_sorted(img.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 mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        db.insert_tag(img.id, "2.0", &now);
        db.insert_tag(img.id, "1.0", &now);
        db.insert_tag(img.id, "latest", &now);
        assert_eq!(db.get_tags_sorted(img.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 mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        db.insert_tag(img.id, "1.0", &now);
        db.insert_tag(img.id, "2.0", &now);
        db.delete_tag(img.id, "1.0");
        assert_eq!(db.get_tags_sorted(img.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 mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: t1 };
        db.insert_image(&mut img);
        db.set_expires_at(img.id, &t2);
        assert_eq!(db.get_image("docker.io", "nginx").unwrap().expires_at, 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 mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        db.insert_tag(img.id, "1.0", &now);
        assert_eq!(db.get_tag_digest(img.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 mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        db.insert_tag(img.id, "1.0", &now);
        db.update_tag_digest(img.id, "1.0", "sha256:abc", &now);
        assert_eq!(db.get_tag_digest(img.id, "1.0"), Some("sha256:abc".to_string()));
    }

    fn test_get_expired_images(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let past = now - chrono::Duration::days(1);
        let future = now + chrono::Duration::days(1);
        let mut expired = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: past };
        let mut fresh = Image { id: 0, registry: "docker.io".into(), image: "redis".into(), expires_at: future };
        db.insert_image(&mut expired);
        db.insert_image(&mut fresh);
        let result = db.get_expired_images(&now);
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].id, expired.id);
        assert_eq!(result[0].image, "nginx");
    }

    #[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")));
    }

    #[test]
    fn conformance_stub_get_expired_images() {
        test_get_expired_images(&StubDb::default());
    }

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