Add struct to represent image row

1 files changed, 140 insertions, 51 deletions

diff --git a/src/db.rs b/src/db.rs
index c914769..fc00621 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -2,9 +2,18 @@ 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<(i64, DateTime<Utc>)>;
-    fn insert_image(&self, registry: &str, image: &str, expires_at: &DateTime<Utc>) -> i64;
+    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>);
@@ -108,23 +117,43 @@ impl SqliteDb {
 }
 
 impl Db for SqliteDb {
-    fn get_image(&self, registry: &str, image: &str) -> Option<(i64, DateTime<Utc>)> {
+    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, expires_at FROM images
+            SELECT id, registry, image, 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();
+        ", (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, registry: &str, image: &str, expires_at: &DateTime<Utc>) -> i64 {
+    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)
-        ", (registry, image, expires_at)).unwrap();
-        return self.conn.last_insert_rowid();
+        ", (&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> {
@@ -172,7 +201,7 @@ impl Db for SqliteDb {
 #[cfg(test)]
 pub struct StubDb {
     next_id: std::cell::RefCell<i64>,
-    images: std::cell::RefCell<Vec<(String, String, i64, DateTime<Utc>)>>,
+    images: std::cell::RefCell<Vec<Image>>,
     tags: std::cell::RefCell<Vec<(i64, String, Option<String>, DateTime<Utc>)>>,
 }
 
@@ -189,20 +218,45 @@ impl Default for StubDb {
 
 #[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));
+    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, registry: &str, image: &str, expires_at: &DateTime<Utc>) -> i64 {
+    fn insert_image(&self, image: &mut Image) {
         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;
+        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> {
@@ -225,9 +279,9 @@ impl Db for StubDb {
     }
 
     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();
+        for img in self.images.borrow_mut().iter_mut() {
+            if img.id == image_id {
+                img.expires_at = *expires_at;
                 return;
             }
         }
@@ -260,9 +314,11 @@ mod tests {
 
     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);
+        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) {
@@ -271,57 +327,78 @@ mod tests {
 
     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);
+        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 id = db.insert_image("docker.io", "nginx", &now);
-        assert_eq!(db.get_tags_sorted(id), Vec::<String>::new());
+        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 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"]);
+        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 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"]);
+        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 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);
+        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 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);
+        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 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()));
+        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]
@@ -431,4 +508,16 @@ mod tests {
         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")));
+    }
 }