Storage refactorings

src/bin/cratesfyi.rs

@@ -5,7 +5,9 @@ use std::sync::Arc;
 
 use cratesfyi::db::{self, add_path_into_database, Pool};
 use cratesfyi::utils::{remove_crate_priority, set_crate_priority};
-use cratesfyi::{BuildQueue, Config, DocBuilder, DocBuilderOptions, RustwideBuilder, Server};
+use cratesfyi::{
+    BuildQueue, Config, DocBuilder, DocBuilderOptions, RustwideBuilder, Server, Storage,
+};
 use failure::{err_msg, Error, ResultExt};
 use once_cell::sync::OnceCell;
 use structopt::StructOpt;
@@ -123,6 +125,7 @@ impl CommandLine {
                     ctx.pool()?,
                     ctx.config()?,
                     ctx.build_queue()?,
+                    ctx.storage()?,
                 )?;
             }
             Self::Daemon {
@@ -137,6 +140,7 @@ impl CommandLine {
                     ctx.config()?,
                     ctx.pool()?,
                     ctx.build_queue()?,
+                    ctx.storage()?,
                     registry_watcher == Toggle::Enabled,
                 )?;
             }
@@ -346,7 +350,7 @@ impl BuildSubcommand {
             Self::World => {
                 docbuilder.load_cache().context("Failed to load cache")?;
 
-                let mut builder = RustwideBuilder::init(ctx.pool()?)?;
+                let mut builder = RustwideBuilder::init(ctx.pool()?, ctx.storage()?)?;
                 builder
                     .build_world(&mut docbuilder)
                     .context("Failed to build world")?;
@@ -360,8 +364,8 @@ impl BuildSubcommand {
                 local,
             } => {
                 docbuilder.load_cache().context("Failed to load cache")?;
-                let mut builder =
-                    RustwideBuilder::init(ctx.pool()?).context("failed to initialize rustwide")?;
+                let mut builder = RustwideBuilder::init(ctx.pool()?, ctx.storage()?)
+                    .context("failed to initialize rustwide")?;
 
                 if let Some(path) = local {
                     builder
@@ -397,14 +401,14 @@ impl BuildSubcommand {
                     }
                 }
 
-                let mut builder = RustwideBuilder::init(ctx.pool()?)?;
+                let mut builder = RustwideBuilder::init(ctx.pool()?, ctx.storage()?)?;
                 builder
                     .update_toolchain()
                     .context("failed to update toolchain")?;
             }
 
             Self::AddEssentialFiles => {
-                let mut builder = RustwideBuilder::init(ctx.pool()?)?;
+                let mut builder = RustwideBuilder::init(ctx.pool()?, ctx.storage()?)?;
                 builder
                     .add_essential_files()
                     .context("failed to add essential files")?;
@@ -469,7 +473,7 @@ impl DatabaseSubcommand {
             }
 
             Self::AddDirectory { directory, prefix } => {
-                add_path_into_database(&*ctx.conn()?, &prefix, directory)
+                add_path_into_database(&*ctx.storage()?, &prefix, directory)
                     .context("Failed to add directory into database")?;
             }
 
@@ -553,6 +557,7 @@ enum DeleteSubcommand {
 
 struct Context {
     build_queue: OnceCell<Arc<BuildQueue>>,
+    storage: OnceCell<Arc<Storage>>,
     config: OnceCell<Arc<Config>>,
     pool: OnceCell<Pool>,
 }
@@ -561,6 +566,7 @@ impl Context {
     fn new() -> Self {
         Self {
             build_queue: OnceCell::new(),
+            storage: OnceCell::new(),
             config: OnceCell::new(),
             pool: OnceCell::new(),
         }
@@ -575,6 +581,13 @@ impl Context {
             .clone())
     }
 
+    fn storage(&self) -> Result<Arc<Storage>, Error> {
+        Ok(self
+            .storage
+            .get_or_try_init::<_, Error>(|| Ok(Arc::new(Storage::new(self.pool()?))))?
+            .clone())
+    }
+
     fn config(&self) -> Result<Arc<Config>, Error> {
         Ok(self
             .config

src/db/delete.rs

@@ -190,17 +190,17 @@ mod tests {
             let db = env.db();
 
             // Create fake packages in the database
-            let pkg1_v1_id = db
+            let pkg1_v1_id = env
                 .fake_release()
                 .name("package-1")
                 .version("1.0.0")
                 .create()?;
-            let pkg1_v2_id = db
+            let pkg1_v2_id = env
                 .fake_release()
                 .name("package-1")
                 .version("2.0.0")
                 .create()?;
-            let pkg2_id = db.fake_release().name("package-2").create()?;
+            let pkg2_id = env.fake_release().name("package-2").create()?;
 
             assert!(crate_exists(&db.conn(), "package-1")?);
             assert!(crate_exists(&db.conn(), "package-2")?);
@@ -244,13 +244,13 @@ mod tests {
             }
 
             let db = env.db();
-            let v1 = db
+            let v1 = env
                 .fake_release()
                 .name("a")
                 .version("1.0.0")
                 .author("malicious actor")
                 .create()?;
-            let v2 = db
+            let v2 = env
                 .fake_release()
                 .name("a")
                 .version("2.0.0")

src/db/file.rs

@@ -6,17 +6,10 @@
 
 use crate::error::Result;
 use crate::storage::{CompressionAlgorithms, Storage};
-use postgres::Connection;
 
 use serde_json::Value;
 use std::path::{Path, PathBuf};
 
-pub(crate) use crate::storage::Blob;
-
-pub(crate) fn get_path(conn: &Connection, path: &str, max_size: usize) -> Result<Blob> {
-    Storage::new(conn).get(path, max_size)
-}
-
 /// Store all files in a directory and return [[mimetype, filename]] as Json
 ///
 /// If there is an S3 Client configured, store files into an S3 bucket;
@@ -27,12 +20,11 @@ pub(crate) fn get_path(conn: &Connection, path: &str, max_size: usize) -> Result
 /// Note that this function is used for uploading both sources
 /// and files generated by rustdoc.
 pub fn add_path_into_database<P: AsRef<Path>>(
-    conn: &Connection,
+    storage: &Storage,
     prefix: &str,
     path: P,
 ) -> Result<(Value, CompressionAlgorithms)> {
-    let mut backend = Storage::new(conn);
-    let (file_list, algorithms) = backend.store_all(conn, prefix, path.as_ref())?;
+    let (file_list, algorithms) = storage.store_all(prefix, path.as_ref())?;
     Ok((
         file_list_to_json(file_list.into_iter().collect())?,
         algorithms,

src/db/mod.rs

@@ -5,10 +5,8 @@ pub(crate) use self::add_package::add_package_into_database;
 pub use self::delete::{delete_crate, delete_version};
 pub use self::file::add_path_into_database;
 pub use self::migrate::migrate;
-pub use self::pool::{Pool, PoolError};
-
-#[cfg(test)]
 pub(crate) use self::pool::PoolConnection;
+pub use self::pool::{Pool, PoolError};
 
 mod add_package;
 pub mod blacklist;

src/docbuilder/rustwide_builder.rs

@@ -6,10 +6,10 @@ use crate::db::{add_build_into_database, add_package_into_database, Pool};
 use crate::docbuilder::{crates::crates_from_path, Limits};
 use crate::error::Result;
 use crate::storage::CompressionAlgorithms;
+use crate::storage::Storage;
 use crate::utils::{copy_doc_dir, parse_rustc_version, CargoMetadata};
 use failure::ResultExt;
 use log::{debug, info, warn, LevelFilter};
-use postgres::Connection;
 use rustwide::cmd::{Command, SandboxBuilder};
 use rustwide::logging::{self, LogStorage};
 use rustwide::toolchain::ToolchainError;
@@ -18,6 +18,7 @@ use serde_json::Value;
 use std::borrow::Cow;
 use std::collections::HashSet;
 use std::path::Path;
+use std::sync::Arc;
 
 const USER_AGENT: &str = "docs.rs builder (https://github.com/rust-lang/docs.rs)";
 const DEFAULT_RUSTWIDE_WORKSPACE: &str = ".rustwide";
@@ -68,12 +69,13 @@ pub struct RustwideBuilder {
     workspace: Workspace,
     toolchain: Toolchain,
     db: Pool,
+    storage: Arc<Storage>,
     rustc_version: String,
     cpu_limit: Option<u32>,
 }
 
 impl RustwideBuilder {
-    pub fn init(db: Pool) -> Result<Self> {
+    pub fn init(db: Pool, storage: Arc<Storage>) -> Result<Self> {
         use rustwide::cmd::SandboxImage;
         let env_workspace_path = ::std::env::var("CRATESFYI_RUSTWIDE_WORKSPACE");
         let workspace_path = env_workspace_path
@@ -108,6 +110,7 @@ impl RustwideBuilder {
             workspace,
             toolchain,
             db,
+            storage,
             rustc_version: String::new(),
             cpu_limit,
         })
@@ -245,7 +248,7 @@ impl RustwideBuilder {
                     })?;
                 }
 
-                add_path_into_database(&conn, "", &dest)?;
+                add_path_into_database(&self.storage, "", &dest)?;
                 conn.query(
                     "INSERT INTO config (name, value) VALUES ('rustc_version', $1) \
                      ON CONFLICT (name) DO UPDATE SET value = $1;",
@@ -350,7 +353,7 @@ impl RustwideBuilder {
                     debug!("adding sources into database");
                     let prefix = format!("sources/{}/{}", name, version);
                     let (files, new_algs) =
-                        add_path_into_database(&conn, &prefix, build.host_source_dir())?;
+                        add_path_into_database(&self.storage, &prefix, build.host_source_dir())?;
                     files_list = Some(files);
                     algs.extend(new_algs);
 
@@ -379,7 +382,7 @@ impl RustwideBuilder {
                             &metadata,
                         )?;
                     }
-                    let new_algs = self.upload_docs(&conn, name, version, local_storage.path())?;
+                    let new_algs = self.upload_docs(name, version, local_storage.path())?;
                     algs.extend(new_algs);
                 };
 
@@ -576,14 +579,13 @@ impl RustwideBuilder {
 
     fn upload_docs(
         &self,
-        conn: &Connection,
         name: &str,
         version: &str,
         local_storage: &Path,
     ) -> Result<CompressionAlgorithms> {
         debug!("Adding documentation into database");
         add_path_into_database(
-            conn,
+            &self.storage,
             &format!("rustdoc/{}/{}", name, version),
             local_storage,
         )

src/lib.rs

@@ -7,6 +7,7 @@ pub use self::config::Config;
 pub use self::docbuilder::options::DocBuilderOptions;
 pub use self::docbuilder::DocBuilder;
 pub use self::docbuilder::RustwideBuilder;
+pub use self::storage::Storage;
 pub use self::web::Server;
 
 mod build_queue;