add dev-util to import release-docs from production for local testing

Author: syphar

Cargo.lock

@@ -48,7 +48,7 @@ postcard = { version = "1.1.3", default-features = false, features = ["use-std"]
 pretty_assertions = "1.4.0"
 rand = "0.9"
 regex = "1"
-reqwest = { version = "0.13", features = ["json", "gzip"] }
+reqwest = { version = "0.13", features = ["json", "gzip", "stream"] }
 sentry = { version = "0.46.0", features = ["panic", "tracing", "tower-http", "anyhow", "backtrace"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
@@ -64,3 +64,4 @@ toml = "0.9.2"
 tracing = "0.1.37"
 url = { version = "2.1.1", features = ["serde"] }
 walkdir = "2"
+zip = { version = "7.0.0", default-features = false, features = ["bzip2"] }

Cargo.toml

@@ -0,0 +1,34 @@
+[package]
+name = "docs_rs_import_release"
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+repository.workspace = true
+edition.workspace = true
+description = "Import a successfully built release from docs.rs into a test deployment"
+
+[dependencies]
+anyhow = { workspace = true }
+async-tar = { version = "0.6.0", default-features = false, features = ["runtime-tokio", "xattr"] }
+clap = { workspace = true }
+docs_rs_cargo_metadata = { path = "../../lib/docs_rs_cargo_metadata" }
+docs_rs_context = { path = "../../lib/docs_rs_context" }
+docs_rs_database = { path = "../../lib/docs_rs_database" }
+docs_rs_logging = { path = "../../lib/docs_rs_logging" }
+docs_rs_registry_api = { path = "../../lib/docs_rs_registry_api" }
+docs_rs_repository_stats = { path = "../../lib/docs_rs_repository_stats" }
+docs_rs_rustdoc_json = { path = "../../lib/docs_rs_rustdoc_json" }
+docs_rs_storage = { path = "../../lib/docs_rs_storage" }
+docs_rs_types = { path = "../../lib/docs_rs_types" }
+docs_rs_utils = { path = "../../lib/docs_rs_utils" }
+docsrs-metadata = { path = "../../lib/metadata" }
+futures-util = { workspace = true }
+regex = { workspace = true }
+reqwest = { workspace = true }
+serde = { workspace = true }
+sqlx = { workspace = true }
+tempfile = { workspace = true }
+tokio = { workspace = true }
+tracing = { workspace = true }
+walkdir = { workspace = true }
+zip = { workspace = true }

crates/bin/docs_rs_import_release/Cargo.toml

@@ -0,0 +1,50 @@
+use anyhow::Result;
+use docs_rs_utils::{APP_USER_AGENT, spawn_blocking};
+use futures_util::StreamExt as _;
+use std::{fmt, sync::LazyLock};
+use tokio::{
+    fs,
+    io::{AsyncSeekExt as _, AsyncWriteExt as _},
+};
+use tracing::debug;
+
+pub(crate) const DOCS_RS: &str = "https://docs.rs";
+pub(crate) static CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
+    reqwest::Client::builder()
+        .user_agent(APP_USER_AGENT)
+        .build()
+        .expect("can't create request client & connection pool")
+});
+
+pub(crate) async fn download(url: impl reqwest::IntoUrl + fmt::Debug) -> Result<Vec<u8>> {
+    debug!("downloading...");
+
+    Ok(CLIENT
+        .get(url)
+        .send()
+        .await?
+        .error_for_status()?
+        .bytes()
+        .await?
+        .to_vec())
+}
+
+pub(crate) async fn download_to_temp_file(url: impl reqwest::IntoUrl) -> Result<fs::File> {
+    debug!("downloading to temp file..");
+
+    let response = CLIENT.get(url).send().await?.error_for_status()?;
+
+    // NOTE: even after being convert to a `tokio::fs::File`, this kind of temporary file
+    // will be cleaned up by the OS, when the last handle is closed.
+    let mut file = fs::File::from_std(spawn_blocking(|| Ok(tempfile::tempfile()?)).await?);
+
+    let mut stream = response.bytes_stream();
+    while let Some(chunk) = stream.next().await {
+        let chunk = chunk?;
+        file.write_all(&chunk).await?;
+    }
+
+    file.sync_all().await?;
+    file.seek(std::io::SeekFrom::Start(0)).await?;
+    Ok(file)
+}

crates/bin/docs_rs_import_release/src/common.rs

@@ -0,0 +1,55 @@
+use crate::common::download_to_temp_file;
+use anyhow::{Result, bail};
+use async_tar::Archive;
+use docs_rs_storage::compression::wrap_reader_for_decompression;
+use docs_rs_types::{CompressionAlgorithm, KrateName, Version};
+use docs_rs_utils::spawn_blocking;
+use std::path::{Path, PathBuf};
+use tokio::io;
+use tracing::debug;
+
+#[derive(Debug)]
+pub(crate) struct SourceDir {
+    _temp_dir: tempfile::TempDir,
+    pub(crate) source_path: PathBuf,
+}
+
+impl AsRef<Path> for SourceDir {
+    fn as_ref(&self) -> &Path {
+        &self.source_path
+    }
+}
+
+pub(crate) async fn download_and_extract_source(
+    name: &KrateName,
+    version: &Version,
+) -> Result<SourceDir> {
+    debug!("downloading source");
+    let crate_archive = download_to_temp_file(format!(
+        "https://static.crates.io/crates/{name}/{name}-{version}.crate"
+    ))
+    .await?;
+
+    let temp_dir = spawn_blocking(|| Ok(tempfile::tempdir()?)).await?;
+
+    debug!("unpacking source archive");
+    {
+        let mut file = io::BufReader::new(crate_archive);
+        let mut decompressed = wrap_reader_for_decompression(&mut file, CompressionAlgorithm::Gzip);
+        let archive = Archive::new(&mut decompressed);
+        archive.unpack(&temp_dir).await?;
+    }
+
+    let source_path = temp_dir.path().join(format!("{name}-{version}"));
+    if !source_path.is_dir() {
+        bail!(
+            "broken crate archive, missing source directory {:?}",
+            source_path
+        );
+    };
+
+    Ok(SourceDir {
+        source_path,
+        _temp_dir: temp_dir,
+    })
+}