Favor registry for queued crates

l4l · Joshua Nelson · commit 8687b7b9e889 · 2020-10-30T11:30:20.000-04:00
diff --git a/src/bin/cratesfyi.rs b/src/bin/cratesfyi.rs
@@ -170,9 +170,12 @@ impl QueueSubcommand {
                 crate_name,
                 crate_version,
                 build_priority,
-            } => ctx
-                .build_queue()?
-                .add_crate(&crate_name, &crate_version, build_priority)?,
+            } => ctx.build_queue()?.add_crate(
+                &crate_name,
+                &crate_version,
+                build_priority,
+                ctx.config()?.registry_url.as_deref(),
+            )?,
 
             Self::DefaultPriority { subcommand } => subcommand.handle_args(ctx)?,
         }
diff --git a/src/build_queue.rs b/src/build_queue.rs
@@ -11,6 +11,7 @@ pub(crate) struct QueuedCrate {
     pub(crate) name: String,
     pub(crate) version: String,
     pub(crate) priority: i32,
+    pub(crate) registry: Option<String>,
 }
 
 #[derive(Debug)]
@@ -29,10 +30,16 @@ impl BuildQueue {
         }
     }
 
-    pub fn add_crate(&self, name: &str, version: &str, priority: i32) -> Result<()> {
+    pub fn add_crate(
+        &self,
+        name: &str,
+        version: &str,
+        priority: i32,
+        registry: Option<&str>,
+    ) -> Result<()> {
         self.db.get()?.execute(
-            "INSERT INTO queue (name, version, priority) VALUES ($1, $2, $3);",
-            &[&name, &version, &priority],
+            "INSERT INTO queue (name, version, priority, registry) VALUES ($1, $2, $3, $4);",
+            &[&name, &version, &priority, &registry],
         )?;
         Ok(())
     }
@@ -63,7 +70,7 @@ impl BuildQueue {
 
     pub(crate) fn queued_crates(&self) -> Result<Vec<QueuedCrate>> {
         let query = self.db.get()?.query(
-            "SELECT id, name, version, priority
+            "SELECT id, name, version, priority, registry
              FROM queue
              WHERE attempt < $1
              ORDER BY priority ASC, attempt ASC, id ASC",
@@ -77,6 +84,7 @@ impl BuildQueue {
                 name: row.get("name"),
                 version: row.get("version"),
                 priority: row.get("priority"),
+                registry: row.get("registry"),
             })
             .collect())
     }
@@ -149,7 +157,7 @@ mod tests {
                 ("high-priority-baz", "1.0.0", -1000),
             ];
             for krate in &test_crates {
-                queue.add_crate(krate.0, krate.1, krate.2)?;
+                queue.add_crate(krate.0, krate.1, krate.2, None)?;
             }
 
             let assert_next = |name| -> Result<()> {
@@ -214,9 +222,9 @@ mod tests {
             let queue = env.build_queue();
 
             assert_eq!(queue.pending_count()?, 0);
-            queue.add_crate("foo", "1.0.0", 0)?;
+            queue.add_crate("foo", "1.0.0", 0, None)?;
             assert_eq!(queue.pending_count()?, 1);
-            queue.add_crate("bar", "1.0.0", 0)?;
+            queue.add_crate("bar", "1.0.0", 0, None)?;
             assert_eq!(queue.pending_count()?, 2);
 
             queue.process_next_crate(|krate| {
@@ -235,11 +243,11 @@ mod tests {
             let queue = env.build_queue();
 
             assert_eq!(queue.prioritized_count()?, 0);
-            queue.add_crate("foo", "1.0.0", 0)?;
+            queue.add_crate("foo", "1.0.0", 0, None)?;
             assert_eq!(queue.prioritized_count()?, 1);
-            queue.add_crate("bar", "1.0.0", -100)?;
+            queue.add_crate("bar", "1.0.0", -100, None)?;
             assert_eq!(queue.prioritized_count()?, 2);
-            queue.add_crate("baz", "1.0.0", 100)?;
+            queue.add_crate("baz", "1.0.0", 100, None)?;
             assert_eq!(queue.prioritized_count()?, 2);
 
             queue.process_next_crate(|krate| {
@@ -262,9 +270,9 @@ mod tests {
             let queue = env.build_queue();
 
             assert_eq!(queue.failed_count()?, 0);
-            queue.add_crate("foo", "1.0.0", -100)?;
+            queue.add_crate("foo", "1.0.0", -100, None)?;
             assert_eq!(queue.failed_count()?, 0);
-            queue.add_crate("bar", "1.0.0", 0)?;
+            queue.add_crate("bar", "1.0.0", 0, None)?;
 
             for _ in 0..MAX_ATTEMPTS {
                 assert_eq!(queue.failed_count()?, 0);
@@ -296,7 +304,7 @@ mod tests {
                 ("baz", "1.0.0", 10),
             ];
             for krate in &test_crates {
-                queue.add_crate(krate.0, krate.1, krate.2)?;
+                queue.add_crate(krate.0, krate.1, krate.2, None)?;
             }
 
             assert_eq!(
diff --git a/src/db/migrate.rs b/src/db/migrate.rs
@@ -475,6 +475,20 @@ pub fn migrate(version: Option<Version>, conn: &mut Client) -> CratesfyiResult<(
                 DROP TYPE feature;                         
             "
         ),
+        migration!(
+            context,
+            20,
+            // description
+            "Support alternative registries",
+            // upgrade query
+            "
+                ALTER TABLE queue ADD COLUMN registry TEXT DEFAULT NULL;
+            ",
+            // downgrade query
+            "
+                ALTER TABLE queue DROP COLUMN registry;
+            "
+        )
     ];
 
     for migration in migrations {
diff --git a/src/docbuilder/queue.rs b/src/docbuilder/queue.rs
@@ -45,10 +45,12 @@ impl DocBuilder {
                 ChangeKind::Added => {
                     let priority = get_crate_priority(&mut conn, &krate.name)?;
 
-                    match self
-                        .build_queue
-                        .add_crate(&krate.name, &krate.version, priority)
-                    {
+                    match self.build_queue.add_crate(
+                        &krate.name,
+                        &krate.version,
+                        priority,
+                        index.repository_url(),
+                    ) {
                         Ok(()) => {
                             debug!("{}-{} added into build queue", krate.name, krate.version);
                             crates_added += 1;
@@ -79,7 +81,13 @@ impl DocBuilder {
         queue.process_next_crate(|krate| {
             processed = true;
 
-            builder.build_package(&krate.name, &krate.version, PackageKind::CratesIo)?;
+            let kind = krate
+                .registry
+                .as_ref()
+                .map(|r| PackageKind::Registry(r.as_str()))
+                .unwrap_or(PackageKind::CratesIo);
+
+            builder.build_package(&krate.name, &krate.version, kind)?;
             Ok(())
         })?;
 
diff --git a/src/index/mod.rs b/src/index/mod.rs
@@ -113,4 +113,8 @@ impl Index {
             );
         }
     }
+
+    pub fn repository_url(&self) -> Option<&str> {
+        self.repository_url.as_deref()
+    }
 }
diff --git a/src/web/releases.rs b/src/web/releases.rs
@@ -1126,9 +1126,9 @@ mod tests {
                 .expect("missing heading")
                 .any(|el| el.text_contents().contains("nothing")));
 
-            queue.add_crate("foo", "1.0.0", 0)?;
-            queue.add_crate("bar", "0.1.0", -10)?;
-            queue.add_crate("baz", "0.0.1", 10)?;
+            queue.add_crate("foo", "1.0.0", 0, None)?;
+            queue.add_crate("bar", "0.1.0", -10, None)?;
+            queue.add_crate("baz", "0.0.1", 10, None)?;
 
             let full = kuchiki::parse_html().one(web.get("/releases/queue").send()?.text()?);
             let items = full

Original file line number	Diff line number	Diff line change
`@@ -113,4 +113,8 @@ impl Index {`
`113`	`113`	`);`
`114`	`114`	`}`
`115`	`115`	`}`
	`116`	`+`
	`117`	`+ pub fn repository_url(&self) -> Option<&str> {`
	`118`	`+ self.repository_url.as_deref()`
	`119`	`+ }`
`116`	`120`	`}`