1 files changed, 428 insertions, 0 deletions
diff --git a/src/db.rs b/src/db.rs
new file mode 100644
index 0000000..457921a
--- /dev/null
+++ b/src/db.rs
@@ -0,0 +1,428 @@
+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, last_checked: 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 update_last_checked(&self, image_id: i64, last_checked: 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();
+        }
+
+        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, last_checked 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, last_checked: DateTime<Utc>) -> i64 {
+        let _timer = crate::metrics::get().db_query_duration.start_timer();
+        self.conn.execute("
+            INSERT INTO images(registry, image, last_checked) VALUES (?1, ?2, ?3)
+        ", (registry, image, last_checked)).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 update_last_checked(&self, image_id: i64, last_checked: DateTime<Utc>) {
+        let _timer = crate::metrics::get().db_query_duration.start_timer();
+        self.conn.execute("UPDATE images SET last_checked = ?1 WHERE id = ?2", (last_checked, 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, last_checked) in self.images.borrow().iter() {
+            if r == registry && i == image {
+                return Some((*id, *last_checked));
+            }
+        }
+        return None;
+    }
+
+    fn insert_image(&self, registry: &str, image: &str, last_checked: 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, last_checked));
+        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));
+    }
+
+    fn update_last_checked(&self, image_id: i64, last_checked: DateTime<Utc>) {
+        for (_, _, id, lc) in self.images.borrow_mut().iter_mut() {
+            if *id == image_id {
+                *lc = last_checked;
+                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;
+                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_update_last_checked(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.update_last_checked(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_update_last_checked() {
+        test_update_last_checked(&StubDb::default());
+    }
+
+    #[test]
+    fn conformance_sqlite_update_last_checked() {
+        crate::metrics::init();
+        let dir = tempfile::tempdir().unwrap();
+        test_update_last_checked(&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")));
+    }
+}