Merge pull request #5066 from arlosi/db-index

Generate index metadata files from the database

Author: Turbo87

Cargo.lock

@@ -10,15 +10,18 @@ edition = "2021"
 path = "lib.rs"
 
 [features]
-testing = ["serde_json"]
+testing = []
 
 [dependencies]
 anyhow = "=1.0.70"
 base64 = "=0.13.1"
 dotenv = "=0.15.0"
 git2 = "=0.17.1"
 serde = { version = "=1.0.160", features = ["derive"] }
+serde_json = "=1.0.96"
 tempfile = "=3.5.0"
 tracing = "=0.1.37"
 url = "=2.3.1"
-serde_json = { version = "=1.0.96", optional = true }
+
+[dev-dependencies]
+claims = "=0.7.1"

cargo-registry-index/Cargo.toml

@@ -148,6 +148,21 @@ pub struct Crate {
     pub v: Option<u32>,
 }
 
+impl Crate {
+    pub fn write_to<W: Write>(&self, mut writer: W) -> anyhow::Result<()> {
+        serde_json::to_writer(&mut writer, self)?;
+        writer.write_all(b"\n")?;
+        Ok(())
+    }
+}
+
+pub fn write_crates<W: Write>(crates: &[Crate], mut writer: W) -> anyhow::Result<()> {
+    for krate in crates {
+        krate.write_to(&mut writer)?;
+    }
+    Ok(())
+}
+
 #[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
 pub struct Dependency {
     pub name: String,
@@ -595,3 +610,57 @@ pub fn run_via_cli(command: &mut Command, credentials: &Credentials) -> anyhow::
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use claims::*;
+
+    #[test]
+    fn crate_writer() {
+        let krate = Crate {
+            name: "foo".to_string(),
+            vers: "1.2.3".to_string(),
+            deps: vec![],
+            cksum: "0123456789asbcdef".to_string(),
+            features: Default::default(),
+            features2: None,
+            yanked: None,
+            links: None,
+            v: None,
+        };
+        let mut buffer = Vec::new();
+        assert_ok!(krate.write_to(&mut buffer));
+        assert_ok_eq!(String::from_utf8(buffer), "\
+            {\"name\":\"foo\",\"vers\":\"1.2.3\",\"deps\":[],\"cksum\":\"0123456789asbcdef\",\"features\":{},\"yanked\":null}\n\
+        ");
+    }
+
+    #[test]
+    fn test_write_crates() {
+        let versions = vec!["0.1.0", "1.0.0-beta.1", "1.0.0", "1.2.3"];
+        let crates = versions
+            .into_iter()
+            .map(|vers| Crate {
+                name: "foo".to_string(),
+                vers: vers.to_string(),
+                deps: vec![],
+                cksum: "0123456789asbcdef".to_string(),
+                features: Default::default(),
+                features2: None,
+                yanked: None,
+                links: None,
+                v: None,
+            })
+            .collect::<Vec<_>>();
+
+        let mut buffer = Vec::new();
+        assert_ok!(write_crates(&crates, &mut buffer));
+        assert_ok_eq!(String::from_utf8(buffer), "\
+            {\"name\":\"foo\",\"vers\":\"0.1.0\",\"deps\":[],\"cksum\":\"0123456789asbcdef\",\"features\":{},\"yanked\":null}\n\
+            {\"name\":\"foo\",\"vers\":\"1.0.0-beta.1\",\"deps\":[],\"cksum\":\"0123456789asbcdef\",\"features\":{},\"yanked\":null}\n\
+            {\"name\":\"foo\",\"vers\":\"1.0.0\",\"deps\":[],\"cksum\":\"0123456789asbcdef\",\"features\":{},\"yanked\":null}\n\
+            {\"name\":\"foo\",\"vers\":\"1.2.3\",\"deps\":[],\"cksum\":\"0123456789asbcdef\",\"features\":{},\"yanked\":null}\n\
+        ");
+    }
+}

cargo-registry-index/lib.rs

@@ -1,3 +1,4 @@
+use crate::background_jobs::Job;
 use crate::{admin::dialoguer, config, db, models::Crate, schema::crates};
 
 use diesel::prelude::*;
@@ -54,5 +55,9 @@ fn delete(opts: Opts, conn: &mut PgConnection) {
         panic!("aborting transaction");
     }
 
-    uploader.delete_index(&client, &krate.name).unwrap();
+    if dotenv::var("FEATURE_INDEX_SYNC").is_ok() {
+        Job::enqueue_sync_to_index(&krate.name, conn).unwrap();
+    } else {
+        uploader.delete_index(&client, &krate.name).unwrap();
+    }
 }

src/admin/delete_crate.rs

@@ -1,3 +1,4 @@
+use crate::background_jobs::Job;
 use crate::{
     admin::dialoguer,
     db,
@@ -57,4 +58,8 @@ fn delete(opts: Opts, conn: &mut PgConnection) {
     if !opts.yes && !dialoguer::confirm("commit?") {
         panic!("aborting transaction");
     }
+
+    if dotenv::var("FEATURE_INDEX_SYNC").is_ok() {
+        Job::enqueue_sync_to_index(&krate.name, conn).unwrap();
+    }
 }

src/admin/delete_version.rs

@@ -5,6 +5,7 @@ use crate::{
     schema::versions,
 };
 
+use crate::background_jobs::Job;
 use diesel::prelude::*;
 
 #[derive(clap::Parser, Debug)]
@@ -64,7 +65,11 @@ fn yank(opts: Opts, conn: &mut PgConnection) {
         .execute(conn)
         .unwrap();
 
-    crate::worker::sync_yanked(krate.name, v.num)
-        .enqueue(conn)
-        .unwrap();
+    if dotenv::var("FEATURE_INDEX_SYNC").is_ok() {
+        Job::enqueue_sync_to_index(&krate.name, conn).unwrap();
+    } else {
+        crate::worker::sync_yanked(krate.name, v.num)
+            .enqueue(conn)
+            .unwrap();
+    }
 }

src/admin/yank_version.rs

@@ -18,6 +18,8 @@ pub enum Job {
     IndexAddCrate(IndexAddCrateJob),
     IndexSquash,
     IndexSyncToHttp(IndexSyncToHttpJob),
+    SyncToGitIndex(SyncToIndexJob),
+    SyncToSparseIndex(SyncToIndexJob),
     IndexUpdateYanked(IndexUpdateYankedJob),
     NormalizeIndex(NormalizeIndexJob),
     RenderAndUploadReadme(RenderAndUploadReadmeJob),
@@ -47,8 +49,47 @@ impl Job {
     const INDEX_UPDATE_YANKED: &str = "sync_yanked";
     const NORMALIZE_INDEX: &str = "normalize_index";
     const RENDER_AND_UPLOAD_README: &str = "render_and_upload_readme";
+    const SYNC_TO_GIT_INDEX: &str = "sync_to_git_index";
+    const SYNC_TO_SPARSE_INDEX: &str = "sync_to_sparse_index";
     const UPDATE_DOWNLOADS: &str = "update_downloads";
 
+    pub fn enqueue_sync_to_index<T: ToString>(
+        krate: T,
+        conn: &mut PgConnection,
+    ) -> Result<(), EnqueueError> {
+        use crate::schema::background_jobs::dsl::*;
+
+        let to_git = Self::sync_to_git_index(krate.to_string());
+        let to_git = (
+            job_type.eq(to_git.as_type_str()),
+            data.eq(to_git.to_value()?),
+        );
+
+        let to_sparse = Self::sync_to_sparse_index(krate.to_string());
+        let to_sparse = (
+            job_type.eq(to_sparse.as_type_str()),
+            data.eq(to_sparse.to_value()?),
+        );
+
+        diesel::insert_into(background_jobs)
+            .values(vec![to_git, to_sparse])
+            .execute(conn)?;
+
+        Ok(())
+    }
+
+    pub fn sync_to_git_index<T: ToString>(krate: T) -> Job {
+        Job::SyncToGitIndex(SyncToIndexJob {
+            krate: krate.to_string(),
+        })
+    }
+
+    pub fn sync_to_sparse_index<T: ToString>(krate: T) -> Job {
+        Job::SyncToSparseIndex(SyncToIndexJob {
+            krate: krate.to_string(),
+        })
+    }
+
     fn as_type_str(&self) -> &'static str {
         match self {
             Job::DailyDbMaintenance => Self::DAILY_DB_MAINTENANCE,
@@ -59,6 +100,8 @@ impl Job {
             Job::IndexUpdateYanked(_) => Self::INDEX_UPDATE_YANKED,
             Job::NormalizeIndex(_) => Self::NORMALIZE_INDEX,
             Job::RenderAndUploadReadme(_) => Self::RENDER_AND_UPLOAD_README,
+            Job::SyncToGitIndex(_) => Self::SYNC_TO_GIT_INDEX,
+            Job::SyncToSparseIndex(_) => Self::SYNC_TO_SPARSE_INDEX,
             Job::UpdateDownloads => Self::UPDATE_DOWNLOADS,
         }
     }
@@ -73,6 +116,8 @@ impl Job {
             Job::IndexUpdateYanked(inner) => serde_json::to_value(inner),
             Job::NormalizeIndex(inner) => serde_json::to_value(inner),
             Job::RenderAndUploadReadme(inner) => serde_json::to_value(inner),
+            Job::SyncToGitIndex(inner) => serde_json::to_value(inner),
+            Job::SyncToSparseIndex(inner) => serde_json::to_value(inner),
             Job::UpdateDownloads => Ok(serde_json::Value::Null),
         }
     }
@@ -101,6 +146,8 @@ impl Job {
             Self::INDEX_UPDATE_YANKED => Job::IndexUpdateYanked(from_value(value)?),
             Self::NORMALIZE_INDEX => Job::NormalizeIndex(from_value(value)?),
             Self::RENDER_AND_UPLOAD_README => Job::RenderAndUploadReadme(from_value(value)?),
+            Self::SYNC_TO_GIT_INDEX => Job::SyncToGitIndex(from_value(value)?),
+            Self::SYNC_TO_SPARSE_INDEX => Job::SyncToSparseIndex(from_value(value)?),
             Self::UPDATE_DOWNLOADS => Job::UpdateDownloads,
             job_type => Err(PerformError::from(format!("Unknown job type {job_type}")))?,
         })
@@ -136,6 +183,8 @@ impl Job {
                 args.base_url.as_deref(),
                 args.pkg_path_in_vcs.as_deref(),
             ),
+            Job::SyncToGitIndex(args) => worker::sync_to_git_index(env, conn, &args.krate),
+            Job::SyncToSparseIndex(args) => worker::sync_to_sparse_index(env, conn, &args.krate),
             Job::UpdateDownloads => worker::perform_update_downloads(&mut *fresh_connection(pool)?),
         }
     }
@@ -172,6 +221,11 @@ pub struct IndexSyncToHttpJob {
     pub(super) crate_name: String,
 }
 
+#[derive(Serialize, Deserialize)]
+pub struct SyncToIndexJob {
+    pub(super) krate: String,
+}
+
 #[derive(Serialize, Deserialize)]
 pub struct IndexUpdateYankedJob {
     pub(super) krate: String,

src/background_jobs.rs

@@ -47,6 +47,7 @@ pub struct Server {
     pub version_id_cache_ttl: Duration,
     pub cdn_user_agent: String,
     pub balance_capacity: BalanceCapacityConfig,
+    pub feature_index_sync: bool,
 }
 
 impl Default for Server {
@@ -151,6 +152,7 @@ impl Default for Server {
             cdn_user_agent: dotenv::var("WEB_CDN_USER_AGENT")
                 .unwrap_or_else(|_| "Amazon CloudFront".into()),
             balance_capacity: BalanceCapacityConfig::from_environment(),
+            feature_index_sync: dotenv::var("FEATURE_INDEX_SYNC").is_ok(),
         }
     }
 }

src/config.rs

@@ -1,6 +1,7 @@
 //! Functionality related to publishing a new crate or version of a crate.
 
 use crate::auth::AuthCheck;
+use crate::background_jobs::Job;
 use axum::body::Bytes;
 use flate2::read::GzDecoder;
 use hex::ToHex;
@@ -271,7 +272,12 @@ pub async fn publish(app: AppState, req: BytesRequest) -> AppResult<Json<GoodCra
                 links,
                 v,
             };
-            worker::add_crate(git_crate).enqueue(conn)?;
+
+            if app.config.feature_index_sync {
+                Job::enqueue_sync_to_index(&krate.name, conn)?;
+            } else {
+                worker::add_crate(git_crate).enqueue(conn)?;
+            }
 
             // The `other` field on `PublishWarnings` was introduced to handle a temporary warning
             // that is no longer needed. As such, crates.io currently does not return any `other`

src/controllers/krate/publish.rs

@@ -1,6 +1,7 @@
 //! Endpoints for yanking and unyanking specific versions of crates
 
 use crate::auth::AuthCheck;
+use crate::background_jobs::Job;
 
 use super::version_and_crate;
 use crate::controllers::cargo_prelude::*;
@@ -84,7 +85,11 @@ fn modify_yank(
 
     insert_version_owner_action(conn, version.id, user.id, api_token_id, action)?;
 
-    worker::sync_yanked(krate.name, version.num).enqueue(conn)?;
+    if state.config.feature_index_sync {
+        Job::enqueue_sync_to_index(&krate.name, conn)?;
+    } else {
+        worker::sync_yanked(krate.name, version.num).enqueue(conn)?;
+    }
 
     ok_true()
 }

src/controllers/version/yank.rs

@@ -44,6 +44,16 @@ pub enum DependencyKind {
     // if you add a kind here, be sure to update `from_row` below.
 }
 
+impl From<IndexDependencyKind> for DependencyKind {
+    fn from(dk: IndexDependencyKind) -> Self {
+        match dk {
+            IndexDependencyKind::Normal => DependencyKind::Normal,
+            IndexDependencyKind::Build => DependencyKind::Build,
+            IndexDependencyKind::Dev => DependencyKind::Dev,
+        }
+    }
+}
+
 impl From<DependencyKind> for IndexDependencyKind {
     fn from(dk: DependencyKind) -> Self {
         match dk {