From 284165f2000ec9ad5ac0c47128ce12eae1c734ef Mon Sep 17 00:00:00 2001 From: Jesper Jensen Date: Mon, 9 Feb 2026 17:50:36 +0100 Subject: Nicer database interface Right now I'm just doing loops, we we can add some batch stuff here later --- src/db.rs | 187 +++++++++++++++++++++++++++++++-------------------------- src/updater.rs | 53 +++++++++------- 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, } +#[derive(Debug, Clone)] +pub struct Tag { + pub id: i64, + pub image_id: i64, + pub tag: String, + pub digest: Option, +} + pub trait Db { fn insert_image(&self, image: &mut Image); fn get_image(&self, registry: &str, image: &str) -> Option; fn get_expired_images(&self, now: &DateTime) -> Vec; - fn get_tags_sorted(&self, image_id: i64) -> Vec; - 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; + 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); - fn get_tag_digest(&self, image_id: i64, tag: &str) -> Option; - 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 { + fn get_tags_sorted(&self, image_id: i64) -> Vec { 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) { + 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 { + 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) { 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, images: std::cell::RefCell>, - tags: std::cell::RefCell)>>, + tags: std::cell::RefCell>, } #[cfg(test)] @@ -265,47 +276,46 @@ impl Db for StubDb { return result; } - fn get_tags_sorted(&self, image_id: i64) -> Vec { + fn get_tags_sorted(&self, image_id: i64) -> Vec { 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) { - 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 { - 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) { + 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::::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 = 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 = 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, 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 = vec![]; + let mut to_insert: Vec = 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, 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, 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, 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, ®, &files, &mut outcomes); - assert_eq!(db.get_tags_sorted(img.id), vec!["1.21", "1.22"]); + let tags: Vec = 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 = db.get_tags_sorted(2).into_iter().map(|t| t.tag).collect(); + assert_eq!(tags, vec!["6.0", "6.2"]); } #[test] -- cgit v1.2.3