Skip to main content

fjall_cli/command/keyspace_command/
iter_command.rs

1use crate::{OutputAffixes, OutputKind, OutputKindWriteError, PrefixKind, Suffix};
2use errgonomic::{handle, handle_bool, map_err};
3use fjall::{Database, Guard, Keyspace, KeyspaceCreateOptions};
4use std::io;
5use std::io::Write;
6use std::process::ExitCode;
7use thiserror::Error;
8
9#[derive(clap::Parser, Clone, Debug)]
10pub struct IterCommand {
11    #[arg(long, value_enum, default_value_t = OutputKind::KeyValue)]
12    kind: OutputKind,
13
14    #[arg(long, value_enum)]
15    item_prefix: Option<PrefixKind>,
16
17    #[arg(long)]
18    item_suffix: Option<Suffix>,
19
20    #[arg(long, value_enum)]
21    key_prefix: Option<PrefixKind>,
22
23    #[arg(long)]
24    key_suffix: Option<Suffix>,
25
26    #[arg(long, value_enum)]
27    value_prefix: Option<PrefixKind>,
28
29    #[arg(long)]
30    value_suffix: Option<Suffix>,
31
32    #[arg(long, default_value_t = 0, help = "Number of items to skip before writing output.")]
33    offset: usize,
34
35    #[arg(long, help = "Maximum number of items to write.")]
36    limit: Option<usize>,
37}
38
39impl IterCommand {
40    pub async fn run(self, db: &Database, keyspace: impl Into<String>) -> Result<ExitCode, IterCommandRunError> {
41        use IterCommandRunError::*;
42        let keyspace = keyspace.into();
43        let Self {
44            kind,
45            item_prefix,
46            item_suffix,
47            key_prefix,
48            key_suffix,
49            value_prefix,
50            value_suffix,
51            offset,
52            limit,
53        } = self;
54        let affixes = OutputAffixes {
55            item_prefix,
56            item_suffix,
57            key_prefix,
58            key_suffix,
59            value_prefix,
60            value_suffix,
61        };
62        handle_bool!(!db.keyspace_exists(&keyspace), KeyspaceNotFound, keyspace);
63        let keyspace_handle = handle!(db.keyspace(&keyspace, KeyspaceCreateOptions::default), KeyspaceFailed, keyspace);
64        let mut stdout = io::stdout().lock();
65        handle!(Self::write_items(&mut stdout, &keyspace_handle, &kind, &affixes, offset, limit,), WriteItemsFailed, keyspace);
66        Ok(ExitCode::SUCCESS)
67    }
68
69    pub fn write_items(writer: &mut impl Write, keyspace: &Keyspace, kind: &OutputKind, affixes: &OutputAffixes, offset: usize, limit: Option<usize>) -> Result<(), IterCommandWriteItemsError> {
70        use IterCommandWriteItemsError::*;
71        let result = match limit {
72            Some(limit) => keyspace
73                .iter()
74                .skip(offset)
75                .take(limit)
76                .try_for_each(|guard| Self::write_item(writer, kind, affixes, guard)),
77            None => keyspace
78                .iter()
79                .skip(offset)
80                .try_for_each(|guard| Self::write_item(writer, kind, affixes, guard)),
81        };
82        map_err!(result, WriteItemFailed)
83    }
84
85    pub fn write_item(writer: &mut impl Write, kind: &OutputKind, affixes: &OutputAffixes, guard: Guard) -> Result<(), IterCommandWriteItemError> {
86        use IterCommandWriteItemError::*;
87        let (key, value) = handle!(guard.into_inner(), IntoInnerFailed);
88        handle!(kind.write(writer, &key, &value, affixes), WriteFailed);
89        Ok(())
90    }
91}
92
93#[derive(Error, Debug)]
94pub enum IterCommandRunError {
95    #[error("keyspace '{keyspace}' not found")]
96    KeyspaceNotFound { keyspace: String },
97
98    #[error("failed to open keyspace '{keyspace}'")]
99    KeyspaceFailed { source: fjall::Error, keyspace: String },
100
101    #[error("failed to write items for keyspace '{keyspace}'")]
102    WriteItemsFailed { source: IterCommandWriteItemsError, keyspace: String },
103}
104
105#[derive(Error, Debug)]
106pub enum IterCommandWriteItemsError {
107    #[error("failed to write item")]
108    WriteItemFailed { source: IterCommandWriteItemError },
109}
110
111#[derive(Error, Debug)]
112pub enum IterCommandWriteItemError {
113    #[error("failed to read key-value pair")]
114    IntoInnerFailed { source: fjall::Error },
115
116    #[error("failed to write output")]
117    WriteFailed { source: OutputKindWriteError },
118}