Merge pull request #514 from cloudflare/alewis/add-bucket-operations

Add bucket::upload, bucket::delete

docs/content/kv_commands.md

@@ -146,3 +146,25 @@ Deletes all specified keys within a given namespace.
 $ wrangler kv:bulk delete f7b02e7fc70443149ac906dd81ec1791 ./allthethings.json
 ```
 
+## `kv:bucket`
+
+### `upload`
+
+Walks the given directory and runs a bulk upload, using the path to an asset as its `key` and the asset as its `value`.
+
+#### Usage
+
+```sh
+$ wrangler kv:bucket upload f7b02e7fc70443149ac906dd81ec1791 ./public
+```
+
+### `delete`
+
+Walks the given directory and runs a bulk delete, using the paths to assets as the `key`s to delete.
+
+#### Usage
+
+```sh
+$ wrangler kv:bucket upload f7b02e7fc70443149ac906dd81ec1791 ./public
+```
+

src/commands/kv/bucket.rs

@@ -0,0 +1,109 @@
+extern crate base64;
+
+use crate::commands::kv::delete_bulk::delete_bulk;
+use crate::commands::kv::write_bulk::write_bulk;
+
+use cloudflare::endpoints::workerskv::write_bulk::KeyValuePair;
+
+use walkdir::WalkDir;
+
+use std::ffi::OsString;
+use std::fs::metadata;
+use std::path::Path;
+
+use crate::terminal::message;
+
+pub fn upload(namespace_id: &str, filename: &Path) -> Result<(), failure::Error> {
+    let pairs: Result<Vec<KeyValuePair>, failure::Error> = match &metadata(filename) {
+        Ok(file_type) if file_type.is_dir() => directory_keys_values(filename),
+        Ok(_file_type) => {
+            // any other file types (files, symlinks)
+            failure::bail!("wrangler kv:bucket upload takes a directory")
+        }
+        Err(e) => failure::bail!("{}", e),
+    };
+
+    write_bulk(namespace_id, pairs?)
+}
+
+pub fn delete(namespace_id: &str, filename: &Path) -> Result<(), failure::Error> {
+    let keys: Result<Vec<String>, failure::Error> = match &metadata(filename) {
+        Ok(file_type) if file_type.is_dir() => directory_keys_only(filename),
+        Ok(_) => {
+            // any other file types (namely, symlinks)
+            failure::bail!(
+                "{} should be a file or directory, but is a symlink",
+                filename.display()
+            )
+        }
+        Err(e) => failure::bail!("{}", e),
+    };
+
+    delete_bulk(namespace_id, keys?)
+}
+
+fn directory_keys_values(directory: &Path) -> Result<Vec<KeyValuePair>, failure::Error> {
+    let mut upload_vec: Vec<KeyValuePair> = Vec::new();
+    for entry in WalkDir::new(directory) {
+        let entry = entry.unwrap();
+        let path = entry.path();
+        if path.is_file() {
+            let key = generate_key(path, directory)?;
+
+            let value = std::fs::read(path)?;
+
+            // Need to base64 encode value
+            let b64_value = base64::encode(&value);
+            message::working(&format!("Parsing {}...", key.clone()));
+            upload_vec.push(KeyValuePair {
+                key: key,
+                value: b64_value,
+                expiration: None,
+                expiration_ttl: None,
+                base64: Some(true),
+            });
+        }
+    }
+    Ok(upload_vec)
+}
+
+fn directory_keys_only(directory: &Path) -> Result<Vec<String>, failure::Error> {
+    let mut upload_vec: Vec<String> = Vec::new();
+    for entry in WalkDir::new(directory) {
+        let entry = entry.unwrap();
+        let path = entry.path();
+        if path.is_file() {
+            let key = generate_key(path, directory)?;
+
+            upload_vec.push(key);
+        }
+    }
+    Ok(upload_vec)
+}
+
+// Courtesy of Steve Kalabnik's PoC :) Used for bulk operations (write, delete)
+fn generate_key(path: &Path, directory: &Path) -> Result<String, failure::Error> {
+    let path = path.strip_prefix(directory).unwrap();
+
+    // next, we have to re-build the paths: if we're on Windows, we have paths with
+    // `\` as separators. But we want to use `/` as separators. Because that's how URLs
+    // work.
+    let mut path_with_forward_slash = OsString::new();
+
+    for (i, component) in path.components().enumerate() {
+        // we don't want a leading `/`, so skip that
+        if i > 0 {
+            path_with_forward_slash.push("/");
+        }
+
+        path_with_forward_slash.push(component);
+    }
+
+    // if we have a non-utf8 path here, it will fail, but that's not realistically going to happen
+    let path = path_with_forward_slash.to_str().expect(&format!(
+        "found a non-UTF-8 path, {:?}",
+        path_with_forward_slash
+    ));
+
+    Ok(path.to_string())
+}

src/commands/kv/delete_bulk.rs

@@ -26,7 +26,7 @@ pub fn delete_json(namespace_id: &str, filename: &Path) -> Result<(), failure::E
     delete_bulk(namespace_id, keys?)
 }
 
-fn delete_bulk(namespace_id: &str, keys: Vec<String>) -> Result<(), failure::Error> {
+pub fn delete_bulk(namespace_id: &str, keys: Vec<String>) -> Result<(), failure::Error> {
     let client = super::api_client()?;
     let account_id = super::account_id()?;
 

src/commands/kv/mod.rs

@@ -7,6 +7,7 @@ use http::status::StatusCode;
 use crate::settings;
 use crate::terminal::message;
 
+pub mod bucket;
 mod create_namespace;
 mod delete_bulk;
 mod delete_key;

src/commands/kv/write_bulk.rs

@@ -27,7 +27,7 @@ pub fn write_json(namespace_id: &str, filename: &Path) -> Result<(), failure::Er
     write_bulk(namespace_id, pairs?)
 }
 
-fn write_bulk(namespace_id: &str, pairs: Vec<KeyValuePair>) -> Result<(), failure::Error> {
+pub fn write_bulk(namespace_id: &str, pairs: Vec<KeyValuePair>) -> Result<(), failure::Error> {
     let client = super::api_client()?;
     let account_id = super::account_id()?;
 

src/commands/publish/mod.rs

@@ -93,7 +93,7 @@ fn upload_bucket(project: &Project) -> Result<(), failure::Error> {
                 Ok(ref file_type) if file_type.is_dir() => {
                     println!("Publishing contents of directory {:?}", path.as_os_str());
 
-                    kv::write_bulk(&namespace.id, Path::new(&path))?;
+                    kv::bucket::upload(&namespace.id, Path::new(&path))?;
                 }
                 Ok(file_type) => {
                     // any other file types (namely, symlinks)

src/main.rs

@@ -244,6 +244,51 @@ fn run() -> Result<(), failure::Error> {
                         )
                 )
         )
+        .subcommand(
+            SubCommand::with_name("kv:bucket")
+                .about(&*format!(
+                    "{} Use KV as bucket-style storage",
+                    emoji::KV
+                ))
+                .subcommand(
+                    SubCommand::with_name("upload")
+                        .about("Upload the contents of a directory keyed on path")
+                        .arg(
+                            Arg::with_name("namespace-id")
+                                .help("The ID of the namespace this action applies to")
+                                .required(true)
+                                // .short("n")
+                                // .long("namespace-id")
+                                // .value_name("<ID>")
+                                // .takes_value(true)
+                        )
+                        .arg(
+                            Arg::with_name("path")
+                            .help("the directory to be uploaded to KV")
+                            .required(true)
+                            .index(2),
+                        )
+                )
+                .subcommand(
+                    SubCommand::with_name("delete")
+                        .about("Delete the contents of a directory keyed on path")
+                        .arg(
+                            Arg::with_name("namespace-id")
+                                .help("The ID of the namespace this action applies to")
+                                .required(true)
+                                // .short("n")
+                                // .long("namespace-id")
+                                // .value_name("<ID>")
+                                // .takes_value(true)
+                        )
+                        .arg(
+                            Arg::with_name("path")
+                            .help("the directory to be deleted from KV")
+                            .required(true)
+                            .index(2),
+                        )
+                )
+        )
         .subcommand(
             SubCommand::with_name("generate")
                 .about(&*format!(
@@ -518,6 +563,21 @@ fn run() -> Result<(), failure::Error> {
             ("", None) => message::warn("kv:bulk expects a subcommand"),
             _ => unreachable!(),
         }
+    } else if let Some(kv_matches) = matches.subcommand_matches("kv:bucket") {
+        match kv_matches.subcommand() {
+            ("upload", Some(write_bulk_matches)) => {
+                let id = write_bulk_matches.value_of("namespace-id").unwrap();
+                let path = write_bulk_matches.value_of("path").unwrap();
+                commands::kv::bucket::upload(id, Path::new(path))?;
+            }
+            ("delete", Some(delete_matches)) => {
+                let id = delete_matches.value_of("namespace-id").unwrap();
+                let path = delete_matches.value_of("path").unwrap();
+                commands::kv::bucket::delete(id, Path::new(path))?;
+            }
+            ("", None) => message::warn("kv:bucket expects a subcommand"),
+            _ => unreachable!(),
+        }
     }
     Ok(())
 }