db: Add default_versions.num_versions column

Author: eth3lbert

crates/crates_io_database/src/schema.rs

@@ -408,6 +408,8 @@ diesel::table! {
         crate_id -> Int4,
         /// Reference to the version in the `versions` table.
         version_id -> Int4,
+        /// The total number of versions.
+        num_versions -> Nullable<Int4>,
     }
 }
 

crates/crates_io_database_dump/src/dump-db.toml

@@ -105,6 +105,7 @@ dependencies = ["crates", "versions"]
 [default_versions.columns]
 crate_id = "public"
 version_id = "public"
+num_versions = "public"
 
 [deleted_crates]
 dependencies = ["users"]

crates/crates_io_database_dump/src/snapshots/[email protected]

@@ -18,7 +18,7 @@ BEGIN ISOLATION LEVEL REPEATABLE READ, READ ONLY;
     \copy (SELECT "crate_id", "created_at", "created_by", "owner_id", "owner_kind" FROM "crate_owners" WHERE NOT deleted) TO 'data/crate_owners.csv' WITH CSV HEADER
 
     \copy "versions" ("bin_names", "categories", "checksum", "crate_id", "crate_size", "created_at", "description", "documentation", "downloads", "edition", "features", "has_lib", "homepage", "id", "keywords", "license", "links", "num", "num_no_build", "published_by", "repository", "rust_version", "updated_at", "yanked") TO 'data/versions.csv' WITH CSV HEADER
-    \copy "default_versions" ("crate_id", "version_id") TO 'data/default_versions.csv' WITH CSV HEADER
+    \copy "default_versions" ("crate_id", "num_versions", "version_id") TO 'data/default_versions.csv' WITH CSV HEADER
     \copy "dependencies" ("crate_id", "default_features", "explicit_name", "features", "id", "kind", "optional", "req", "target", "version_id") TO 'data/dependencies.csv' WITH CSV HEADER
     \copy "version_downloads" ("date", "downloads", "version_id") TO 'data/version_downloads.csv' WITH CSV HEADER
 COMMIT;

crates/crates_io_database_dump/src/snapshots/[email protected]

@@ -60,7 +60,7 @@ BEGIN;
     \copy "crates_keywords" ("crate_id", "keyword_id") FROM 'data/crates_keywords.csv' WITH CSV HEADER
     \copy "crate_owners" ("crate_id", "created_at", "created_by", "owner_id", "owner_kind") FROM 'data/crate_owners.csv' WITH CSV HEADER
     \copy "versions" ("bin_names", "categories", "checksum", "crate_id", "crate_size", "created_at", "description", "documentation", "downloads", "edition", "features", "has_lib", "homepage", "id", "keywords", "license", "links", "num", "num_no_build", "published_by", "repository", "rust_version", "updated_at", "yanked") FROM 'data/versions.csv' WITH CSV HEADER
-    \copy "default_versions" ("crate_id", "version_id") FROM 'data/default_versions.csv' WITH CSV HEADER
+    \copy "default_versions" ("crate_id", "num_versions", "version_id") FROM 'data/default_versions.csv' WITH CSV HEADER
     \copy "dependencies" ("crate_id", "default_features", "explicit_name", "features", "id", "kind", "optional", "req", "target", "version_id") FROM 'data/dependencies.csv' WITH CSV HEADER
     \copy "version_downloads" ("date", "downloads", "version_id") FROM 'data/version_downloads.csv' WITH CSV HEADER
 

migrations/2025-02-05-083109_add-num-versions-column/down.sql

@@ -0,0 +1,4 @@
+ALTER TABLE default_versions
+    DROP COLUMN num_versions;
+
+DROP FUNCTION IF EXISTS update_num_versions_from_versions CASCADE;

migrations/2025-02-05-083109_add-num-versions-column/up.sql

@@ -0,0 +1,27 @@
+ALTER TABLE default_versions
+    ADD COLUMN num_versions INTEGER;
+
+COMMENT ON COLUMN default_versions.num_versions IS 'The total number of versions.';
+
+CREATE OR REPLACE FUNCTION update_num_versions_from_versions() RETURNS TRIGGER AS $$
+BEGIN
+    IF (TG_OP = 'INSERT') THEN
+        INSERT INTO default_versions (crate_id, version_id, num_versions)
+        VALUES (NEW.crate_id, NEW.id, 1)
+        ON CONFLICT (crate_id) DO UPDATE
+        SET num_versions = EXCLUDED.num_versions + 1;
+        RETURN NEW;
+    ELSIF (TG_OP = 'DELETE') THEN
+        UPDATE default_versions
+        SET num_versions = num_versions - 1
+        WHERE crate_id = OLD.crate_id;
+        RETURN OLD;
+    END IF;
+END
+$$ LANGUAGE plpgsql;
+
+DROP TRIGGER IF EXISTS trigger_update_num_versions_from_versions ON versions;
+CREATE TRIGGER trigger_update_num_versions_from_versions
+     AFTER INSERT OR DELETE ON versions
+     FOR EACH ROW
+     EXECUTE PROCEDURE update_num_versions_from_versions();

src/controllers/krate/publish.rs

@@ -445,7 +445,7 @@ pub async fn publish(app: AppState, req: Parts, body: Body) -> AppResult<Json<Go
         // Upsert the `default_value` determined by the existing `default_value` and the
         // published version. Note that this could potentially write an outdated version
         // (although this should not happen regularly), as we might be comparing to an
-        // outdated value.
+        // outdated value. The initial record will be handled by the trigger function.
         //
         // Compared to only using a background job, this prevents us from getting into a
         // situation where a crate exists in the `crates` table but doesn't have a default
@@ -470,14 +470,6 @@ pub async fn publish(app: AppState, req: Parts, body: Body) -> AppResult<Json<Go
             // Update the default version asynchronously in a background job
             // to ensure correctness and eventual consistency.
             UpdateDefaultVersion::new(krate.id).enqueue(conn).await?;
-        } else {
-            diesel::insert_into(default_versions::table)
-                .values((
-                    default_versions::crate_id.eq(krate.id),
-                    default_versions::version_id.eq(version.id),
-                ))
-                .execute(conn)
-                .await?;
         }
 
         // Update all keywords for this crate

src/tests/krate/publish/basics.rs

@@ -116,12 +116,23 @@ async fn new_krate_twice() {
 // The primary purpose is to verify that the `default_version` we provide is as expected.
 #[tokio::test(flavor = "multi_thread")]
 async fn new_krate_twice_alt() {
+    use crate::schema::default_versions;
+
     let (app, _, _, token) = TestApp::full().with_token().await;
+    let mut conn = app.db_conn().await;
 
     let crate_to_publish =
         PublishBuilder::new("foo_twice", "2.0.0").description("2.0.0 description");
     token.publish_crate(crate_to_publish).await.good();
 
+    let num_versions = default_versions::table
+        .select(default_versions::num_versions)
+        .load::<Option<i32>>(&mut conn)
+        .await
+        .unwrap();
+    assert_eq!(num_versions.len(), 1);
+    assert_eq!(num_versions[0], Some(1));
+
     let crate_to_publish = PublishBuilder::new("foo_twice", "0.99.0");
     let response = token.publish_crate(crate_to_publish).await;
     assert_eq!(response.status(), StatusCode::OK);
@@ -130,6 +141,14 @@ async fn new_krate_twice_alt() {
         ".crate.updated_at" => "[datetime]",
     });
 
+    let num_versions = default_versions::table
+        .select(default_versions::num_versions)
+        .load::<Option<i32>>(&mut conn)
+        .await
+        .unwrap();
+    assert_eq!(num_versions.len(), 1);
+    assert_eq!(num_versions[0], Some(2));
+
     let crates = app.crates_from_index_head("foo_twice");
     assert_json_snapshot!(crates);