1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

use anyhow::Result;
use serde_json::value::Value as JsonValue;

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 crate::settings::toml::Target;

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(
        target: &Target,
        client: HttpApiClient,
        namespace_id: &str,
        prefix: Option<&str>,
    ) -> Result<KeyList> {
        let iter = KeyList {
            keys_result: None,
            prefix: prefix.map(str::to_string),
            client,
            account_id: target.account_id.load()?.to_owned(),
            namespace_id: namespace_id.to_string(),
            cursor: None,
            init_fetch: false,
        };
        Ok(iter)
    }

    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,
        }
    }

    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);
                    key.map(Ok)
                }
                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)
    }
}