Alewis/refactor key list (#531)

* Write KeyList iterator

src/commands/kv/key/key_list.rs

@@ -0,0 +1,116 @@
+use cloudflare::endpoints::workerskv::list_namespace_keys::ListNamespaceKeys;
+use cloudflare::endpoints::workerskv::list_namespace_keys::ListNamespaceKeysParams;
+use cloudflare::endpoints::workerskv::Key;
+use cloudflare::framework::apiclient::ApiClient;
+use cloudflare::framework::response::ApiFailure;
+use cloudflare::framework::HttpApiClient;
+
+use serde_json::value::Value as JsonValue;
+
+use crate::settings::project::Project;
+
+pub struct KeyList {
+    keys_result: Option<Vec<Key>>,
+    prefix: Option<String>,
+    client: HttpApiClient,
+    account_id: String,
+    namespace_id: String,
+    cursor: Option<String>,
+    init_fetch: bool,
+}
+
+impl KeyList {
+    pub fn new(
+        project: &Project,
+        client: HttpApiClient,
+        namespace_id: &str,
+        prefix: Option<&str>,
+    ) -> KeyList {
+        KeyList {
+            keys_result: None,
+            prefix: prefix.map(str::to_string),
+            client,
+            account_id: project.account_id.to_owned(),
+            namespace_id: namespace_id.to_string(),
+            cursor: None,
+            init_fetch: false,
+        }
+    }
+
+    fn request_params(&self) -> ListNamespaceKeys {
+        let params = ListNamespaceKeysParams {
+            limit: None, // Defaults to 1000 (the maximum)
+            cursor: self.cursor.to_owned(),
+            prefix: self.prefix.to_owned(),
+        };
+
+        ListNamespaceKeys {
+            account_identifier: &self.account_id,
+            namespace_identifier: &self.namespace_id,
+            params: params,
+        }
+    }
+
+    fn get_batch(&mut self) -> Result<Vec<Key>, ApiFailure> {
+        let response = self.client.request(&self.request_params());
+
+        match response {
+            Ok(success) => {
+                self.cursor = extract_cursor(success.result_info.clone());
+                log::info!("{:?}", self.cursor);
+                Ok(success.result)
+            }
+            Err(e) => Err(e),
+        }
+    }
+}
+
+impl Iterator for KeyList {
+    type Item = Result<Key, ApiFailure>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // Attempt to extract next key from vector of keys in KeyList.
+        // If no key vector or no keys left, go to fallback case below to
+        // attempt to fetch the next page of keys from the Workers KV API.
+        if let Some(mut keys) = self.keys_result.to_owned() {
+            let key = keys.pop();
+            self.keys_result = Some(keys);
+
+            if let Some(k) = key {
+                return Some(Ok(k));
+            }
+        }
+        // Fallback case (if no remaining keys are found)
+        if self.cursor.is_none() && self.init_fetch {
+            None // Nothing left to fetch
+        } else {
+            if !self.init_fetch {
+                // At this point, initial fetch is being performed.
+                self.init_fetch = true;
+            }
+            match self.get_batch() {
+                Ok(mut keys) => {
+                    let key = keys.pop();
+                    self.keys_result = Some(keys);
+                    match key {
+                        Some(k) => Some(Ok(k)),
+                        None => None,
+                    }
+                }
+                Err(e) => Some(Err(e)),
+            }
+        }
+    }
+}
+
+// Returns Some(cursor) if cursor is non-empty, otherwise returns None.
+fn extract_cursor(result_info: Option<JsonValue>) -> Option<String> {
+    let result_info = result_info.unwrap();
+    let returned_cursor_value = &result_info["cursor"];
+    let returned_cursor = returned_cursor_value.as_str().unwrap().to_string();
+    if returned_cursor.is_empty() {
+        None
+    } else {
+        Some(returned_cursor)
+    }
+}

src/commands/kv/key/list.rs

@@ -1,12 +1,7 @@
 extern crate serde_json;
 
-use cloudflare::endpoints::workerskv::list_namespace_keys::ListNamespaceKeys;
-use cloudflare::endpoints::workerskv::list_namespace_keys::ListNamespaceKeysParams;
-use cloudflare::endpoints::workerskv::Key;
-use cloudflare::framework::apiclient::ApiClient;
-use serde_json::value::Value as JsonValue;
-
 use crate::commands::kv;
+use crate::commands::kv::key::KeyList;
 use crate::settings::global_user::GlobalUser;
 use crate::settings::project::Project;
 
@@ -16,78 +11,33 @@ use crate::settings::project::Project;
 pub fn list(
     project: &Project,
     user: GlobalUser,
-    id: &str,
+    namespace_id: &str,
     prefix: Option<&str>,
 ) -> Result<(), failure::Error> {
     let client = kv::api_client(user)?;
 
-    let params = ListNamespaceKeysParams {
-        limit: None, // Defaults to 1000 (the maximum)
-        cursor: None,
-        prefix: prefix.map(kv::url_encode_key),
-    };
-
-    let mut request_params = ListNamespaceKeys {
-        account_identifier: &project.account_id,
-        namespace_identifier: id,
-        params: params,
-    };
-
-    let mut response = client.request(&request_params);
+    let key_list = KeyList::new(project, client, namespace_id, prefix);
 
     print!("["); // Open json list bracket
 
-    // Iterate over all pages until no pages of keys are left.
-    // This is detected when a returned cursor is an empty string.
-    loop {
-        let (result, cursor) = match response {
-            Ok(success) => (
-                success.result,
-                get_cursor_from_result_info(success.result_info.clone()),
-            ),
-            Err(e) => failure::bail!(e),
-        };
+    let mut first_key = true;
 
-        match cursor {
-            None => {
-                // Case where we are done iterating through pages (no cursor returned)
-                print_page(result, true)?;
-                print!("]"); // Close json list bracket
-                break;
-            }
-            Some(_) => {
-                // Case where we still have pages to iterate through (a cursor is returned).
-                print_page(result, false)?;
+    for key_result in key_list {
+        match key_result {
+            Ok(key) => {
+                if first_key {
+                    first_key = false;
+                } else {
+                    print!(",");
+                }
 
-                // Update cursor in request_params.params, and make another request to Workers KV API.
-                request_params.params.cursor = cursor;
-                response = client.request(&request_params);
+                print!("{}", serde_json::to_string(&key)?);
             }
+            Err(e) => kv::print_error(e),
         }
     }
 
-    Ok(())
-}
+    print!("]"); // Close json list bracket
 
-// Returns Some(cursor) if cursor is non-empty, otherwise returns None.
-fn get_cursor_from_result_info(result_info: Option<JsonValue>) -> Option<String> {
-    let result_info = result_info.unwrap();
-    let returned_cursor_value = &result_info["cursor"];
-    let returned_cursor = returned_cursor_value.as_str().unwrap().to_string();
-    if returned_cursor.is_empty() {
-        None
-    } else {
-        Some(returned_cursor)
-    }
-}
-
-fn print_page(keys: Vec<Key>, last_page: bool) -> Result<(), failure::Error> {
-    for i in 0..keys.len() {
-        print!("{}", serde_json::to_string(&keys[i])?);
-        // if last key on last page, don't print final comma.
-        if !(last_page && i == keys.len() - 1) {
-            print!(",");
-        }
-    }
     Ok(())
 }

src/commands/kv/key/mod.rs

@@ -1,9 +1,11 @@
 mod delete;
 mod get;
+mod key_list;
 mod list;
 mod put;
 
 pub use delete::delete;
 pub use get::get;
+pub use key_list::KeyList;
 pub use list::list;
 pub use put::put;