storage: move deleting a prefix into Storage and add tests for it

Author: pietroalbini

src/bin/cratesfyi.rs

@@ -483,11 +483,11 @@ impl DatabaseSubcommand {
 
             Self::Delete {
                 command: DeleteSubcommand::Version { name, version },
-            } => db::delete_version(&*ctx.config()?, &*ctx.conn()?, &name, &version)
+            } => db::delete_version(&*ctx.conn()?, &*ctx.storage()?, &name, &version)
                 .context("failed to delete the crate")?,
             Self::Delete {
                 command: DeleteSubcommand::Crate { name },
-            } => db::delete_crate(&*ctx.config()?, &*ctx.conn()?, &name)
+            } => db::delete_crate(&*ctx.conn()?, &*ctx.storage()?, &name)
                 .context("failed to delete the crate")?,
             Self::Blacklist { command } => command.handle_args(ctx)?,
         }

src/db/delete.rs

@@ -1,10 +1,6 @@
-use crate::{
-    storage::s3::{s3_client, S3_RUNTIME},
-    Config,
-};
+use crate::Storage;
 use failure::{Error, Fail};
 use postgres::Connection;
-use rusoto_s3::{DeleteObjectsRequest, ListObjectsV2Request, ObjectIdentifier, S3Client, S3};
 
 /// List of directories in docs.rs's underlying storage (either the database or S3) containing a
 /// subdirectory named after the crate. Those subdirectories will be deleted.
@@ -16,35 +12,27 @@ enum CrateDeletionError {
     MissingCrate(String),
 }
 
-pub fn delete_crate(config: &Config, conn: &Connection, name: &str) -> Result<(), Error> {
+pub fn delete_crate(conn: &Connection, storage: &Storage, name: &str) -> Result<(), Error> {
     let crate_id = get_id(conn, name)?;
     delete_crate_from_database(conn, name, crate_id)?;
 
-    if let Some(client) = s3_client() {
-        for prefix in STORAGE_PATHS_TO_DELETE {
-            delete_prefix_from_s3(config, &client, &format!("{}/{}/", prefix, name))?;
-        }
+    for prefix in STORAGE_PATHS_TO_DELETE {
+        storage.delete_prefix(&format!("{}/{}/", prefix, name))?;
     }
 
     Ok(())
 }
 
 pub fn delete_version(
-    config: &Config,
     conn: &Connection,
+    storage: &Storage,
     name: &str,
     version: &str,
 ) -> Result<(), Error> {
     delete_version_from_database(conn, name, version)?;
 
-    if let Some(client) = s3_client() {
-        for prefix in STORAGE_PATHS_TO_DELETE {
-            delete_prefix_from_s3(
-                config,
-                &client,
-                &format!("{}/{}/{}/", prefix, name, version),
-            )?;
-        }
+    for prefix in STORAGE_PATHS_TO_DELETE {
+        storage.delete_prefix(&format!("{}/{}/{}/", prefix, name, version))?;
     }
 
     Ok(())
@@ -120,70 +108,12 @@ fn delete_crate_from_database(conn: &Connection, name: &str, crate_id: i32) -> R
     transaction.execute("DELETE FROM releases WHERE crate_id = $1;", &[&crate_id])?;
     transaction.execute("DELETE FROM crates WHERE id = $1;", &[&crate_id])?;
 
-    for prefix in STORAGE_PATHS_TO_DELETE {
-        transaction.execute(
-            "DELETE FROM files WHERE path LIKE $1;",
-            &[&format!("{}/{}/%", prefix, name)],
-        )?;
-    }
-
     // Transactions automatically rollback when not committing, so if any of the previous queries
     // fail the whole transaction will be aborted.
     transaction.commit()?;
     Ok(())
 }
 
-fn delete_prefix_from_s3(config: &Config, s3: &S3Client, name: &str) -> Result<(), Error> {
-    S3_RUNTIME.handle().block_on(async {
-        let mut continuation_token = None;
-        loop {
-            let list = s3
-                .list_objects_v2(ListObjectsV2Request {
-                    bucket: config.s3_bucket.clone(),
-                    prefix: Some(name.into()),
-                    continuation_token,
-                    ..ListObjectsV2Request::default()
-                })
-                .await?;
-
-            let to_delete = list
-                .contents
-                .unwrap_or_else(Vec::new)
-                .into_iter()
-                .filter_map(|o| o.key)
-                .map(|key| ObjectIdentifier {
-                    key,
-                    version_id: None,
-                })
-                .collect::<Vec<_>>();
-
-            let resp = s3
-                .delete_objects(DeleteObjectsRequest {
-                    bucket: config.s3_bucket.clone(),
-                    delete: rusoto_s3::Delete {
-                        objects: to_delete,
-                        quiet: None,
-                    },
-                    ..DeleteObjectsRequest::default()
-                })
-                .await?;
-
-            if let Some(errs) = resp.errors {
-                for err in &errs {
-                    log::error!("error deleting file from s3: {:?}", err);
-                }
-
-                failure::bail!("uploading to s3 failed");
-            }
-
-            continuation_token = list.continuation_token;
-            if continuation_token.is_none() {
-                return Ok(());
-            }
-        }
-    })
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -289,7 +219,7 @@ mod tests {
                 vec!["malicious actor".to_string(), "Peter Rabbit".to_string()]
             );
 
-            delete_version(&*env.config(), &db.conn(), "a", "1.0.0")?;
+            delete_version(&db.conn(), &*env.storage(), "a", "1.0.0")?;
             assert!(!release_exists(&db.conn(), v1)?);
             assert!(release_exists(&db.conn(), v2)?);
             assert_eq!(

src/storage/database.rs

@@ -99,6 +99,14 @@ impl<'a> StorageTransaction for DatabaseStorageTransaction<'a> {
         Ok(())
     }
 
+    fn delete_prefix(&mut self, prefix: &str) -> Result<(), Error> {
+        self.transaction.execute(
+            "DELETE FROM files WHERE path LIKE $1;",
+            &[&format!("{}%", prefix.replace('%', "\\%"))],
+        )?;
+        Ok(())
+    }
+
     fn complete(self: Box<Self>) -> Result<(), Error> {
         self.transaction.commit()?;
         Ok(())

src/storage/mod.rs

@@ -111,6 +111,24 @@ impl Storage {
         Ok(blob)
     }
 
+    fn transaction<T, F>(&self, f: F) -> Result<T, Error>
+    where
+        F: FnOnce(&mut dyn StorageTransaction) -> Result<T, Error>,
+    {
+        let conn;
+        let mut trans: Box<dyn StorageTransaction> = match &self.backend {
+            StorageBackend::Database(db) => {
+                conn = db.start_connection()?;
+                Box::new(conn.start_storage_transaction()?)
+            }
+            StorageBackend::S3(s3) => Box::new(s3.start_storage_transaction()?),
+        };
+
+        let res = f(trans.as_mut())?;
+        trans.complete()?;
+        Ok(res)
+    }
+
     // Store all files in `root_dir` into the backend under `prefix`.
     //
     // If the environment is configured with S3 credentials, this will upload to S3;
@@ -167,30 +185,23 @@ impl Storage {
         &self,
         mut blobs: impl Iterator<Item = Result<Blob, Error>>,
     ) -> Result<(), Error> {
-        let conn;
-        let mut trans: Box<dyn StorageTransaction> = match &self.backend {
-            StorageBackend::Database(db) => {
-                conn = db.start_connection()?;
-                Box::new(conn.start_storage_transaction()?)
-            }
-            StorageBackend::S3(s3) => Box::new(s3.start_storage_transaction()?),
-        };
-
-        loop {
-            let batch: Vec<_> = blobs
-                .by_ref()
-                .take(MAX_CONCURRENT_UPLOADS)
-                .collect::<Result<_, Error>>()?;
-
-            if batch.is_empty() {
-                break;
+        self.transaction(|trans| {
+            loop {
+                let batch: Vec<_> = blobs
+                    .by_ref()
+                    .take(MAX_CONCURRENT_UPLOADS)
+                    .collect::<Result<_, Error>>()?;
+                if batch.is_empty() {
+                    break;
+                }
+                trans.store_batch(batch)?;
             }
+            Ok(())
+        })
+    }
 
-            trans.store_batch(batch)?;
-        }
-
-        trans.complete()?;
-        Ok(())
+    pub(crate) fn delete_prefix(&self, prefix: &str) -> Result<(), Error> {
+        self.transaction(|trans| trans.delete_prefix(prefix))
     }
 
     // We're using `&self` instead of consuming `self` or creating a Drop impl because during tests
@@ -216,6 +227,7 @@ impl std::fmt::Debug for Storage {
 
 trait StorageTransaction {
     fn store_batch(&mut self, batch: Vec<Blob>) -> Result<(), Error>;
+    fn delete_prefix(&mut self, prefix: &str) -> Result<(), Error>;
     fn complete(self: Box<Self>) -> Result<(), Error>;
 }
 
@@ -453,6 +465,70 @@ mod backend_tests {
         Ok(())
     }
 
+    fn test_delete_prefix(storage: &Storage) -> Result<(), Error> {
+        test_deletion(
+            storage,
+            "foo/bar/",
+            &[
+                "foo.txt",
+                "foo/bar.txt",
+                "foo/bar/baz.txt",
+                "foo/bar/foobar.txt",
+                "bar.txt",
+            ],
+            &["foo.txt", "foo/bar.txt", "bar.txt"],
+            &["foo/bar/baz.txt", "foo/bar/foobar.txt"],
+        )
+    }
+
+    fn test_delete_percent(storage: &Storage) -> Result<(), Error> {
+        // PostgreSQL treats "%" as a special char when deleting a prefix. Make sure any "%" in the
+        // provided prefix is properly escaped.
+        test_deletion(
+            storage,
+            "foo/%/",
+            &["foo/bar.txt", "foo/%/bar.txt"],
+            &["foo/bar.txt"],
+            &["foo/%/bar.txt"],
+        )
+    }
+
+    fn test_deletion(
+        storage: &Storage,
+        prefix: &str,
+        start: &[&str],
+        present: &[&str],
+        missing: &[&str],
+    ) -> Result<(), Error> {
+        storage.store_blobs(
+            start
+                .iter()
+                .map(|path| Blob {
+                    path: (*path).to_string(),
+                    content: b"foo\n".to_vec(),
+                    compression: None,
+                    mime: "text/plain".into(),
+                    date_updated: Utc::now(),
+                })
+                .collect(),
+        )?;
+
+        storage.delete_prefix(prefix)?;
+
+        for existing in present {
+            assert!(storage.get(existing, std::usize::MAX).is_ok());
+        }
+        for missing in missing {
+            assert!(storage
+                .get(missing, std::usize::MAX)
+                .unwrap_err()
+                .downcast_ref::<PathNotFoundError>()
+                .is_some());
+        }
+
+        Ok(())
+    }
+
     // Remember to add the test name to the macro below when adding a new one.
 
     macro_rules! backend_tests {
@@ -495,6 +571,8 @@ mod backend_tests {
             test_get_too_big,
             test_store_blobs,
             test_store_all,
+            test_delete_prefix,
+            test_delete_percent,
         }
     }
 }