Nicer database interface

Right now I'm just doing loops, we we can add some batch stuff here
later

2 files changed, 134 insertions, 106 deletions

diff --git a/src/db.rs b/src/db.rs
index 26ea78a..f1a65fd 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -2,7 +2,7 @@ use chrono::{DateTime, Utc};
 use rusqlite::Connection;
 use rusqlite::OptionalExtension;
 
-#[derive(Debug, PartialEq)]
+#[derive(Debug, Clone)]
 pub struct Image {
     pub id: i64,
     pub registry: String,
@@ -10,17 +10,24 @@ pub struct Image {
     pub expires_at: DateTime<Utc>,
 }
 
+#[derive(Debug, Clone)]
+pub struct Tag {
+    pub id: i64,
+    pub image_id: i64,
+    pub tag: String,
+    pub digest: Option<String>,
+}
+
 pub trait Db {
     fn insert_image(&self, image: &mut Image);
 
     fn get_image(&self, registry: &str, image: &str) -> Option<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);
+    fn get_tags_sorted(&self, image_id: i64) -> Vec<Tag>;
+    fn insert_tags(&self, tags: &mut [Tag]);
+    fn delete_tags(&self, tag_ids: &[i64]);
+    fn update_tags(&self, tags: &[Tag]);
     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);
 }
 
 pub struct SqliteDb {
@@ -162,45 +169,49 @@ impl Db for SqliteDb {
         return images;
     }
 
-    fn get_tags_sorted(&self, image_id: i64) -> Vec<String> {
+    fn get_tags_sorted(&self, image_id: i64) -> Vec<Tag> {
         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 stmt = self.conn.prepare("SELECT id, image_id, tag, digest 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());
+            tags.push(Tag {
+                id: row.get(0).unwrap(),
+                image_id: row.get(1).unwrap(),
+                tag: row.get(2).unwrap(),
+                digest: row.get(3).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) {
+    fn insert_tags(&self, tags: &mut [Tag]) {
         let _timer = crate::metrics::get().db_query_duration.start_timer();
-        self.conn.execute("INSERT INTO tags(image_id, tag) VALUES (?1, ?2)",
-            (image_id, tag)).unwrap();
+        for tag in tags {
+            self.conn.execute("INSERT INTO tags(image_id, tag, digest) VALUES (?1, ?2, ?3)",
+                (&tag.image_id, &tag.tag, &tag.digest)).unwrap();
+            tag.id = self.conn.last_insert_rowid();
+        }
     }
 
-    fn set_expires_at(&self, image_id: i64, expires_at: &DateTime<Utc>) {
+    fn delete_tags(&self, tag_ids: &[i64]) {
         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();
+        for id in tag_ids {
+            self.conn.execute("DELETE FROM tags WHERE id = ?1", (id,)).unwrap();
+        }
     }
 
-    fn get_tag_digest(&self, image_id: i64, tag: &str) -> Option<String> {
+    fn update_tags(&self, tags: &[Tag]) {
         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();
+        for tag in tags {
+            self.conn.execute("UPDATE tags SET tag = ?1, digest = ?2 WHERE id = ?3",
+                (&tag.tag, &tag.digest, &tag.id)).unwrap();
+        }
     }
 
-    fn update_tag_digest(&self, image_id: i64, tag: &str, digest: &str) {
+    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 tags SET digest = ?1 WHERE image_id = ?2 AND tag = ?3",
-            (digest, image_id, tag)).unwrap();
+        self.conn.execute("UPDATE images SET expires_at = ?1 WHERE id = ?2", (expires_at, image_id)).unwrap();
     }
 }
 
@@ -208,7 +219,7 @@ impl Db for SqliteDb {
 pub struct StubDb {
     next_id: std::cell::RefCell<i64>,
     images: std::cell::RefCell<Vec<Image>>,
-    tags: std::cell::RefCell<Vec<(i64, String, Option<String>)>>,
+    tags: std::cell::RefCell<Vec<Tag>>,
 }
 
 #[cfg(test)]
@@ -265,47 +276,46 @@ impl Db for StubDb {
         return result;
     }
 
-    fn get_tags_sorted(&self, image_id: i64) -> Vec<String> {
+    fn get_tags_sorted(&self, image_id: i64) -> Vec<Tag> {
         let mut result = vec![];
-        for (id, tag, _) in self.tags.borrow().iter() {
-            if *id == image_id {
+        for tag in self.tags.borrow().iter() {
+            if tag.image_id == image_id {
                 result.push(tag.clone());
             }
         }
-        result.sort();
+        result.sort_by(|a, b| a.tag.cmp(&b.tag));
         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_tags(&self, tags: &mut [Tag]) {
+        for tag in tags {
+            let id = *self.next_id.borrow();
+            *self.next_id.borrow_mut() += 1;
+            tag.id = id;
+            self.tags.borrow_mut().push(tag.clone());
+        }
     }
 
-    fn insert_tag(&self, image_id: i64, tag: &str) {
-        self.tags.borrow_mut().push((image_id, tag.to_string(), None));
+    fn delete_tags(&self, tag_ids: &[i64]) {
+        self.tags.borrow_mut().retain(|t| !tag_ids.contains(&t.id));
     }
 
-    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 update_tags(&self, tags: &[Tag]) {
+        for tag in tags {
+            for t in self.tags.borrow_mut().iter_mut() {
+                if t.id == tag.id {
+                    t.tag = tag.tag.clone();
+                    t.digest = tag.digest.clone();
+                    break;
+                }
             }
         }
     }
 
-    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) {
-        for (id, t, d) in self.tags.borrow_mut().iter_mut() {
-            if *id == image_id && t == tag {
-                *d = Some(digest.to_string());
+    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;
             }
         }
@@ -327,7 +337,7 @@ mod tests {
     }
 
     fn test_get_image_returns_none_for_unknown(db: &dyn Db) {
-        assert_eq!(db.get_image("docker.io", "unknown"), None);
+        assert!(db.get_image("docker.io", "unknown").is_none());
     }
 
     fn test_get_image_returns_inserted(db: &dyn Db) {
@@ -343,27 +353,34 @@ mod tests {
         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());
+        assert!(db.get_tags_sorted(img.id).is_empty());
     }
 
     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");
-        db.insert_tag(img.id, "1.0");
-        db.insert_tag(img.id, "latest");
-        assert_eq!(db.get_tags_sorted(img.id), vec!["1.0", "2.0", "latest"]);
+        db.insert_tags(&mut [
+            Tag { id: 0, image_id: img.id, tag: "2.0".into(), digest: None },
+            Tag { id: 0, image_id: img.id, tag: "1.0".into(), digest: None },
+            Tag { id: 0, image_id: img.id, tag: "latest".into(), digest: None },
+        ]);
+        let tags: Vec<String> = db.get_tags_sorted(img.id).into_iter().map(|t| t.tag).collect();
+        assert_eq!(tags, vec!["1.0", "2.0", "latest"]);
     }
 
-    fn test_delete_tag_removes_tag(db: &dyn Db) {
+    fn test_delete_tags_removes_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);
-        db.insert_tag(img.id, "1.0");
-        db.insert_tag(img.id, "2.0");
-        db.delete_tag(img.id, "1.0");
-        assert_eq!(db.get_tags_sorted(img.id), vec!["2.0"]);
+        let mut tags = [
+            Tag { id: 0, image_id: img.id, tag: "1.0".into(), digest: None },
+            Tag { id: 0, image_id: img.id, tag: "2.0".into(), digest: None },
+        ];
+        db.insert_tags(&mut tags);
+        db.delete_tags(&[tags[0].id]);
+        let result: Vec<String> = db.get_tags_sorted(img.id).into_iter().map(|t| t.tag).collect();
+        assert_eq!(result, vec!["2.0"]);
     }
 
     fn test_set_expires_at(db: &dyn Db) {
@@ -375,21 +392,23 @@ mod tests {
         assert_eq!(db.get_image("docker.io", "nginx").unwrap().expires_at, t2);
     }
 
-    fn test_get_tag_digest_returns_none_when_unset(db: &dyn Db) {
+    fn test_tag_digest_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");
-        assert_eq!(db.get_tag_digest(img.id, "1.0"), None);
+        db.insert_tags(&mut [Tag { id: 0, image_id: img.id, tag: "1.0".into(), digest: None }]);
+        assert_eq!(db.get_tags_sorted(img.id)[0].digest, None);
     }
 
-    fn test_update_tag_digest(db: &dyn Db) {
+    fn test_update_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);
-        db.insert_tag(img.id, "1.0");
-        db.update_tag_digest(img.id, "1.0", "sha256:abc");
-        assert_eq!(db.get_tag_digest(img.id, "1.0"), Some("sha256:abc".to_string()));
+        let mut tags = [Tag { id: 0, image_id: img.id, tag: "1.0".into(), digest: None }];
+        db.insert_tags(&mut tags);
+        tags[0].digest = Some("sha256:abc".into());
+        db.update_tags(&tags);
+        assert_eq!(db.get_tags_sorted(img.id)[0].digest, Some("sha256:abc".to_string()));
     }
 
     fn test_get_expired_images(db: &dyn Db) {
@@ -467,15 +486,15 @@ mod tests {
     }
 
     #[test]
-    fn conformance_stub_delete_tag() {
-        test_delete_tag_removes_tag(&StubDb::default());
+    fn conformance_stub_delete_tags() {
+        test_delete_tags_removes_tags(&StubDb::default());
     }
 
     #[test]
-    fn conformance_sqlite_delete_tag() {
+    fn conformance_sqlite_delete_tags() {
         crate::metrics::init();
         let dir = tempfile::tempdir().unwrap();
-        test_delete_tag_removes_tag(&SqliteDb::new(&dir.path().join("db.sqlite")));
+        test_delete_tags_removes_tags(&SqliteDb::new(&dir.path().join("db.sqlite")));
     }
 
     #[test]
@@ -491,27 +510,27 @@ mod tests {
     }
 
     #[test]
-    fn conformance_stub_get_tag_digest_none() {
-        test_get_tag_digest_returns_none_when_unset(&StubDb::default());
+    fn conformance_stub_tag_digest_none() {
+        test_tag_digest_none_when_unset(&StubDb::default());
     }
 
     #[test]
-    fn conformance_sqlite_get_tag_digest_none() {
+    fn conformance_sqlite_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_tag_digest_none_when_unset(&SqliteDb::new(&dir.path().join("db.sqlite")));
     }
 
     #[test]
-    fn conformance_stub_update_tag_digest() {
-        test_update_tag_digest(&StubDb::default());
+    fn conformance_stub_update_tags() {
+        test_update_tags(&StubDb::default());
     }
 
     #[test]
-    fn conformance_sqlite_update_tag_digest() {
+    fn conformance_sqlite_update_tags() {
         crate::metrics::init();
         let dir = tempfile::tempdir().unwrap();
-        test_update_tag_digest(&SqliteDb::new(&dir.path().join("db.sqlite")));
+        test_update_tags(&SqliteDb::new(&dir.path().join("db.sqlite")));
     }
 
     #[test]
diff --git a/src/updater.rs b/src/updater.rs
index 9f04710..7ab6483 100644
--- a/src/updater.rs
+++ b/src/updater.rs
@@ -1,6 +1,6 @@
 use crate::version::{VersionPattern, CompareOutcome};
 use crate::docker::DockerRef;
-use crate::db::{Db, Image};
+use crate::db::{Db, Image, Tag};
 use crate::registry::Registries;
 use crate::metrics;
 use chrono::{DateTime, Utc};
@@ -59,21 +59,25 @@ pub fn update_images(now: &DateTime<Utc>, db: &dyn Db, reg: &dyn Registries, fil
         match reg.get_tags(&expired.registry, &expired.image) {
             Some(mut fetched_tags) => {
                 fetched_tags.sort();
+                let mut to_delete: Vec<i64> = vec![];
+                let mut to_insert: Vec<Tag> = vec![];
                 let mut existing = cached_tags.iter().peekable();
                 for t in &fetched_tags {
-                    while existing.peek().is_some_and(|e| *e < t) {
+                    while existing.peek().is_some_and(|e| e.tag.as_str() < t.as_str()) {
                         let gone = existing.next().unwrap();
-                        db.delete_tag(expired.id, gone);
+                        to_delete.push(gone.id);
                     }
-                    if existing.peek() == Some(&t) {
+                    if existing.peek().is_some_and(|e| e.tag.as_str() == t.as_str()) {
                         existing.next();
                     } else {
-                        db.insert_tag(expired.id, t);
+                        to_insert.push(Tag { id: 0, image_id: expired.id, tag: t.clone(), digest: None });
                     }
                 }
                 for gone in existing {
-                    db.delete_tag(expired.id, gone);
+                    to_delete.push(gone.id);
                 }
+                db.delete_tags(&to_delete);
+                db.insert_tags(&mut to_insert);
                 let new_expires_at = *now + cache_max_age;
                 db.set_expires_at(expired.id, &new_expires_at);
             }
@@ -117,12 +121,12 @@ pub fn update_images(now: &DateTime<Utc>, db: &dyn Db, reg: &dyn Registries, fil
             let mut current = VersionPattern::parse(&tag_str);
             let mut new_tag = None;
 
-            for candidate_str in &tags {
-                let candidate = VersionPattern::parse(&candidate_str);
-                match current.compare(&candidate) {
+            for candidate in &tags {
+                let candidate_pattern = VersionPattern::parse(&candidate.tag);
+                match current.compare(&candidate_pattern) {
                     CompareOutcome::Higher => {
-                        new_tag = Some(candidate_str.clone());
-                        current = candidate;
+                        new_tag = Some(candidate.tag.clone());
+                        current = candidate_pattern;
                     },
                     CompareOutcome::Lower => {},
                     CompareOutcome::Incompatible => {},
@@ -143,7 +147,7 @@ pub fn update_images(now: &DateTime<Utc>, db: &dyn Db, reg: &dyn Registries, fil
         if let Some(ref digest) = img.digest {
             let tag_for_digest = tag.as_deref().unwrap_or("latest");
             let cached_digest = if image_record.expires_at > *now {
-                db.get_tag_digest(image_record.id, tag_for_digest)
+                tags.iter().find(|t| t.tag == tag_for_digest).and_then(|t| t.digest.clone())
             } else {
                 None
             };
@@ -153,7 +157,10 @@ pub fn update_images(now: &DateTime<Utc>, db: &dyn Db, reg: &dyn Registries, fil
                 None => {
                     match reg.get_digest(&image_record.registry, &image_record.image, tag_for_digest) {
                         Some(fetched) => {
-                            db.update_tag_digest(image_record.id, tag_for_digest, &fetched);
+                            if let Some(mut tag_record) = tags.iter().find(|t| t.tag == tag_for_digest).cloned() {
+                                tag_record.digest = Some(fetched.clone());
+                                db.update_tags(&[tag_record]);
+                            }
                             Ok(fetched)
                         }
                         None => Err(format!("digest not found for {}:{}", image_record.image, tag_for_digest)),
@@ -210,8 +217,10 @@ mod tests {
         let past_expiry = now - chrono::Duration::days(1);
         let mut img = Image { id: 0, registry: "registry.hub.docker.com".into(), image: "nginx".into(), expires_at: past_expiry };
         db.insert_image(&mut img);
-        db.insert_tag(img.id, "1.20");
-        db.insert_tag(img.id, "1.21");
+        db.insert_tags(&mut [
+            Tag { id: 0, image_id: img.id, tag: "1.20".into(), digest: None },
+            Tag { id: 0, image_id: img.id, tag: "1.21".into(), digest: None },
+        ]);
 
         let files = vec![FileInput {
             path: "/test".into(),
@@ -221,7 +230,8 @@ mod tests {
         let mut outcomes = vec![];
         update_images(&now, &db, &reg, &files, &mut outcomes);
 
-        assert_eq!(db.get_tags_sorted(img.id), vec!["1.21", "1.22"]);
+        let tags: Vec<String> = db.get_tags_sorted(img.id).into_iter().map(|t| t.tag).collect();
+        assert_eq!(tags, vec!["1.21", "1.22"]);
         let got = db.get_image("registry.hub.docker.com", "nginx").unwrap();
         assert!(got.expires_at > now);
     }
@@ -236,7 +246,7 @@ mod tests {
         let future_expiry = now + chrono::Duration::minutes(30);
         let mut img = Image { id: 0, registry: "registry.hub.docker.com".into(), image: "nginx".into(), expires_at: future_expiry };
         db.insert_image(&mut img);
-        db.insert_tag(img.id, "1.21");
+        db.insert_tags(&mut [Tag { id: 0, image_id: img.id, tag: "1.21".into(), digest: None }]);
 
         let files = vec![FileInput {
             path: "/test".into(),
@@ -344,8 +354,7 @@ mod tests {
         let future_expiry = now + chrono::Duration::minutes(30);
         let mut img = Image { id: 0, registry: "registry.hub.docker.com".into(), image: "nginx".into(), expires_at: future_expiry };
         db.insert_image(&mut img);
-        db.insert_tag(img.id, "1.21");
-        db.update_tag_digest(img.id, "1.21", "sha256:0000000000000000000000000000000000000000000000000000000000000003");
+        db.insert_tags(&mut [Tag { id: 0, image_id: img.id, tag: "1.21".into(), digest: Some("sha256:0000000000000000000000000000000000000000000000000000000000000003".into()) }]);
 
         let files = vec![FileInput {
             path: "/test".into(),
@@ -578,8 +587,7 @@ mod tests {
         let past_expiry = now - chrono::Duration::days(1);
         let mut img = Image { id: 0, registry: "registry.hub.docker.com".into(), image: "nginx".into(), expires_at: past_expiry };
         db.insert_image(&mut img);
-        db.insert_tag(img.id, "1.21");
-        db.update_tag_digest(img.id, "1.21", "sha256:0000000000000000000000000000000000000000000000000000000000000003");
+        db.insert_tags(&mut [Tag { id: 0, image_id: img.id, tag: "1.21".into(), digest: Some("sha256:0000000000000000000000000000000000000000000000000000000000000003".into()) }]);
 
         let files = vec![FileInput {
             path: "/test".into(),
@@ -614,7 +622,8 @@ mod tests {
         assert!(outcomes[0].is_err());
         assert!(db.get_image("registry.hub.docker.com", "redis").is_some());
         // Redis tags ARE populated (cache refresh happens before processing), but no patches are produced
-        assert_eq!(db.get_tags_sorted(2), vec!["6.0", "6.2"]);
+        let tags: Vec<String> = db.get_tags_sorted(2).into_iter().map(|t| t.tag).collect();
+        assert_eq!(tags, vec!["6.0", "6.2"]);
     }
 
     #[test]