Support custom indexes

Author: l4l

Cargo.lock

@@ -20,7 +20,7 @@ log = "0.4"
 regex = "1"
 structopt = "0.3"
 crates-index = "0.15.1"
-crates-index-diff = "7"
+crates-index-diff = "7.1"
 reqwest = { version = "0.10.6", features = ["blocking", "json"] } # TODO: Remove blocking when async is ready
 semver = { version = "0.9", features = ["serde"] }
 slug = "=0.1.1"
@@ -40,7 +40,7 @@ schemamama = "0.3"
 schemamama_postgres = "0.3"
 systemstat = "0.1.4"
 prometheus = { version = "0.10.0", default-features = false }
-rustwide = "0.10.0"
+rustwide = { git = "https://github.com/rust-lang/rustwide.git" , rev = "bcf19b84128c983d02257a0ae9b91e5a6772122b"}
 mime_guess = "2"
 dotenv = "0.15"
 zstd = "0.5"

Cargo.toml

@@ -6,7 +6,8 @@ use std::sync::Arc;
 use docs_rs::db::{self, add_path_into_database, Pool, PoolClient};
 use docs_rs::utils::{remove_crate_priority, set_crate_priority};
 use docs_rs::{
-    BuildQueue, Config, Context, DocBuilder, Index, Metrics, RustwideBuilder, Server, Storage,
+    BuildQueue, Config, Context, DocBuilder, Index, Metrics, PackageKind, RustwideBuilder, Server,
+    Storage,
 };
 use failure::{err_msg, Error, ResultExt};
 use once_cell::sync::OnceCell;
@@ -257,8 +258,16 @@ enum BuildSubcommand {
         #[structopt(name = "CRATE_VERSION")]
         crate_version: Option<String>,
 
+        /// Url for registry different from cratesio
+        #[structopt(name = "CRATE_REGISTRY")]
+        crate_registry: Option<String>,
+
         /// Build a crate at a specific path
-        #[structopt(short = "l", long = "local", conflicts_with_all(&["CRATE_NAME", "CRATE_VERSION"]))]
+        #[structopt(
+            short = "l",
+            long = "local",
+            conflicts_with_all(&["CRATE_NAME", "CRATE_VERSION", "CRATE_REGISTRY"])
+        )]
         local: Option<PathBuf>,
     },
 
@@ -299,6 +308,7 @@ impl BuildSubcommand {
             Self::Crate {
                 crate_name,
                 crate_version,
+                crate_registry,
                 local,
             } => {
                 let mut builder = rustwide_builder()?;
@@ -313,7 +323,10 @@ impl BuildSubcommand {
                             &crate_name.ok_or_else(|| err_msg("must specify name if not local"))?,
                             &crate_version
                                 .ok_or_else(|| err_msg("must specify version if not local"))?,
-                            None,
+                            crate_registry
+                                .as_ref()
+                                .map(|s| PackageKind::Registry(s.as_str()))
+                                .unwrap_or(PackageKind::CratesIo),
                         )
                         .context("Building documentation failed")?;
                 }
@@ -593,7 +606,16 @@ impl Context for BinContext {
     fn index(&self) -> Result<Arc<Index>, Error> {
         Ok(self
             .index
-            .get_or_try_init::<_, Error>(|| Ok(Arc::new(Index::new(&*self.config()?)?)))?
+            .get_or_try_init::<_, Error>(|| {
+                let config = self.config()?;
+                Ok(Arc::new(
+                    if let Some(registry_url) = config.registry_url.clone() {
+                        Index::from_url(config.registry_index_path.clone(), registry_url)
+                    } else {
+                        Index::new(config.registry_index_path.clone())
+                    }?,
+                ))
+            })?
             .clone())
     }
 }

src/bin/cratesfyi.rs

@@ -12,6 +12,7 @@ pub struct Config {
 
     pub prefix: PathBuf,
     pub registry_index_path: PathBuf,
+    pub registry_url: Option<String>,
 
     // Database connection params
     pub(crate) database_url: String,
@@ -56,6 +57,7 @@ impl Config {
 
             prefix: prefix.clone(),
             registry_index_path: env("REGISTRY_INDEX_PATH", prefix.join("crates.io-index"))?,
+            registry_url: maybe_env("REGISTRY_URL")?,
 
             database_url: require_env("CRATESFYI_DATABASE_URL")?,
             max_pool_size: env("DOCSRS_MAX_POOL_SIZE", 90)?,

src/config.rs

@@ -4,8 +4,8 @@ mod queue;
 mod rustwide_builder;
 
 pub(crate) use self::limits::Limits;
-pub use self::rustwide_builder::RustwideBuilder;
 pub(crate) use self::rustwide_builder::{BuildResult, DocCoverage};
+pub use self::rustwide_builder::{PackageKind, RustwideBuilder};
 
 use crate::db::Pool;
 use crate::error::Result;

src/docbuilder/mod.rs

@@ -1,6 +1,6 @@
 //! Updates registry index and builds new packages
 
-use super::{DocBuilder, RustwideBuilder};
+use super::{DocBuilder, PackageKind, RustwideBuilder};
 use crate::error::Result;
 use crate::utils::get_crate_priority;
 use crate::Index;
@@ -79,7 +79,7 @@ impl DocBuilder {
         queue.process_next_crate(|krate| {
             processed = true;
 
-            builder.build_package(&krate.name, &krate.version, None)?;
+            builder.build_package(&krate.name, &krate.version, PackageKind::CratesIo)?;
             Ok(())
         })?;
 

src/docbuilder/queue.rs

@@ -56,6 +56,12 @@ const ESSENTIAL_FILES_UNVERSIONED: &[&str] = &[
 const DUMMY_CRATE_NAME: &str = "empty-library";
 const DUMMY_CRATE_VERSION: &str = "1.0.0";
 
+pub enum PackageKind<'a> {
+    Local(&'a Path),
+    CratesIo,
+    Registry(&'a str),
+}
+
 pub struct RustwideBuilder {
     workspace: Workspace,
     toolchain: Toolchain,
@@ -259,7 +265,12 @@ impl RustwideBuilder {
         crates_from_path(
             &self.config.registry_index_path.clone(),
             &mut |name, version| {
-                if let Err(err) = self.build_package(name, version, None) {
+                let registry_url = self.config.registry_url.clone();
+                let package_kind = registry_url
+                    .as_ref()
+                    .map(|r| PackageKind::Registry(r.as_str()))
+                    .unwrap_or(PackageKind::CratesIo);
+                if let Err(err) = self.build_package(name, version, package_kind) {
                     warn!("failed to build package {} {}: {}", name, version, err);
                 }
             },
@@ -273,14 +284,14 @@ impl RustwideBuilder {
                 err.context(format!("failed to load local package {}", path.display()))
             })?;
         let package = metadata.root();
-        self.build_package(&package.name, &package.version, Some(path))
+        self.build_package(&package.name, &package.version, PackageKind::Local(path))
     }
 
     pub fn build_package(
         &mut self,
         name: &str,
         version: &str,
-        local: Option<&Path>,
+        kind: PackageKind<'_>,
     ) -> Result<bool> {
         let mut conn = self.db.get()?;
 
@@ -302,10 +313,10 @@ impl RustwideBuilder {
         let mut build_dir = self.workspace.build_dir(&format!("{}-{}", name, version));
         build_dir.purge()?;
 
-        let krate = if let Some(path) = local {
-            Crate::local(path)
-        } else {
-            Crate::crates_io(name, version)
+        let krate = match kind {
+            PackageKind::Local(path) => Crate::local(path),
+            PackageKind::CratesIo => Crate::crates_io(name, version),
+            PackageKind::Registry(registry) => Crate::registry(registry, name, version),
         };
         krate.fetch(&self.workspace)?;
 

src/docbuilder/rustwide_builder.rs

@@ -3,7 +3,7 @@ use std::{path::PathBuf, process::Command};
 use url::Url;
 
 use self::{api::Api, crates::Crates};
-use crate::{error::Result, Config};
+use crate::error::Result;
 use failure::ResultExt;
 
 pub(crate) mod api;
@@ -12,9 +12,10 @@ mod crates;
 pub struct Index {
     path: PathBuf,
     api: Api,
+    repository_url: Option<String>,
 }
 
-#[derive(serde::Deserialize, Clone)]
+#[derive(Debug, serde::Deserialize, Clone)]
 #[serde(rename_all = "kebab-case")]
 struct IndexConfig {
     dl: String,
@@ -40,19 +41,44 @@ fn load_config(repo: &git2::Repository) -> Result<IndexConfig> {
 }
 
 impl Index {
-    pub fn new(app_config: &Config) -> Result<Self> {
-        let path = app_config.registry_index_path.clone();
+    pub fn from_url(path: PathBuf, repository_url: String) -> Result<Self> {
+        let url = repository_url.clone();
+        let diff = crates_index_diff::Index::from_path_or_cloned_with_options(
+            &path,
+            crates_index_diff::CloneOptions { repository_url },
+        )
+        .context("initialising registry index repository")?;
+
+        let config = load_config(diff.repository()).context("loading registry config")?;
+        let api = Api::new(config.api).context("initialising registry api client")?;
+        Ok(Self {
+            path,
+            api,
+            repository_url: Some(url),
+        })
+    }
+
+    pub fn new(path: PathBuf) -> Result<Self> {
         // This initializes the repository, then closes it afterwards to avoid leaking file descriptors.
         // See https://github.com/rust-lang/docs.rs/pull/847
         let diff = crates_index_diff::Index::from_path_or_cloned(&path)
             .context("initialising registry index repository")?;
         let config = load_config(diff.repository()).context("loading registry config")?;
         let api = Api::new(config.api).context("initialising registry api client")?;
-        Ok(Self { path, api })
+        Ok(Self {
+            path,
+            api,
+            repository_url: None,
+        })
     }
 
     pub(crate) fn diff(&self) -> Result<crates_index_diff::Index> {
-        let diff = crates_index_diff::Index::from_path_or_cloned(&self.path)
+        let options = self
+            .repository_url
+            .clone()
+            .map(|repository_url| crates_index_diff::CloneOptions { repository_url })
+            .unwrap_or_default();
+        let diff = crates_index_diff::Index::from_path_or_cloned_with_options(&self.path, options)
             .context("re-opening registry index for diff")?;
         Ok(diff)
     }

src/index/mod.rs

@@ -6,6 +6,7 @@ pub use self::build_queue::BuildQueue;
 pub use self::config::Config;
 pub use self::context::Context;
 pub use self::docbuilder::DocBuilder;
+pub use self::docbuilder::PackageKind;
 pub use self::docbuilder::RustwideBuilder;
 pub use self::index::Index;
 pub use self::metrics::Metrics;

src/lib.rs

@@ -197,7 +197,10 @@ impl TestEnvironment {
     pub(crate) fn index(&self) -> Arc<Index> {
         self.index
             .get_or_init(|| {
-                Arc::new(Index::new(&*self.config()).expect("failed to initialize the index"))
+                Arc::new(
+                    Index::new(self.config().registry_index_path.clone())
+                        .expect("failed to initialize the index"),
+                )
             })
             .clone()
     }