Merge pull request #377 from Mark-Simulacrum/s3-migrate

Migrate files database to S3

Author: QuietMisdreavus

Cargo.lock

@@ -34,7 +34,8 @@ schemamama_postgres = "0.2"
 rusoto_s3 = "0.40"
 rusoto_core = "0.40"
 rusoto_credential = "0.40"
-
+futures = "0.1"
+tokio = "0.1"
 
 # iron dependencies
 iron = "0.5"

Cargo.toml

@@ -109,6 +109,7 @@ pub fn main() {
         .subcommand(SubCommand::with_name("daemon").about("Starts cratesfyi daemon"))
         .subcommand(SubCommand::with_name("database")
             .about("Database operations")
+            .subcommand(SubCommand::with_name("move-to-s3"))
             .subcommand(SubCommand::with_name("migrate")
                 .about("Run database migrations")
                 .arg(Arg::with_name("VERSION")))
@@ -239,6 +240,18 @@ pub fn main() {
         } else if let Some(_) = matches.subcommand_matches("update-search-index") {
             let conn = db::connect_db().unwrap();
             db::update_search_index(&conn).expect("Failed to update search index");
+        } else if let Some(_) = matches.subcommand_matches("move-to-s3") {
+            let conn = db::connect_db().unwrap();
+            let mut count = 1;
+            let mut total = 0;
+            while count != 0 {
+                count = db::file::move_to_s3(&conn, 5_000).expect("Failed to upload batch to S3");
+                total += count;
+                eprintln!(
+                    "moved {} rows to s3 in this batch, total moved so far: {}",
+                    count, total
+                );
+            }
         }
     } else if let Some(matches) = matches.subcommand_matches("start-web-server") {
         start_web_server(Some(matches.value_of("SOCKET_ADDR").unwrap_or("0.0.0.0:3000")));

src/bin/cratesfyi.rs

@@ -160,7 +160,6 @@ pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
 
             let content: Option<Vec<u8>> = if let Some(client) = &client {
                 let s3_res = client.put_object(PutObjectRequest {
-                    acl: Some("public-read".into()),
                     bucket: "rust-docs-rs".into(),
                     key: bucket_path.clone(),
                     body: Some(content.clone().into()),
@@ -225,7 +224,53 @@ fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {
     Ok(file_list_json.to_json())
 }
 
+pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
+    let trans = try!(conn.transaction());
+    let client = s3_client().expect("configured s3");
+
+    let rows = try!(trans.query(
+            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
+            &[]));
+    let count = rows.len();
+
+    let mut rt = ::tokio::runtime::current_thread::Runtime::new().unwrap();
+    let mut futures = Vec::new();
+    for row in &rows {
+        let path: String = row.get(0);
+        let mime: String = row.get(1);
+        let content: Vec<u8> = row.get(2);
+        let path_1 = path.clone();
+        futures.push(client.put_object(PutObjectRequest {
+            bucket: "rust-docs-rs".into(),
+            key: path.clone(),
+            body: Some(content.into()),
+            content_type: Some(mime),
+            ..Default::default()
+        }).map(move |_| {
+            path_1
+        }).map_err(move |e| {
+            panic!("failed to upload to {}: {:?}", path, e)
+        }));
+    }
 
+    use ::futures::future::Future;
+    match rt.block_on(::futures::future::join_all(futures)) {
+        Ok(paths) => {
+            let statement = trans.prepare("UPDATE files SET content = E'in-s3' WHERE path = $1")
+                .unwrap();
+            for path in paths {
+                statement.execute(&[&path]).unwrap();
+            }
+        }
+        Err(e) => {
+            panic!("results err: {:?}", e);
+        }
+    }
+
+    try!(trans.commit());
+
+    Ok(count)
+}
 
 #[cfg(test)]
 mod test {

src/db/file.rs

@@ -33,6 +33,8 @@ extern crate schemamama_postgres;
 extern crate rusoto_s3;
 extern crate rusoto_core;
 extern crate rusoto_credential;
+extern crate futures;
+extern crate tokio;
 
 pub use self::docbuilder::DocBuilder;
 pub use self::docbuilder::ChrootBuilderResult;