nu_plugin_mongo/
lib.rs

1mod cmds;
2use cmds::*;
3use mongodb::sync::Database;
4use nu_plugin::{Plugin, PluginCommand};
5use nu_protocol::{LabeledError, Span};
6use std::collections::HashMap;
7use std::sync::RwLock;
8
9struct Handle {
10    pub(crate) inner: HashMap<u8, (Database, String)>,
11    pub(crate) current: u8,
12}
13
14impl Handle {
15    fn new() -> Self {
16        Self {
17            inner: HashMap::default(),
18            current: 0,
19        }
20    }
21}
22pub struct MongoPlugin {
23    handlers: RwLock<Handle>,
24}
25
26impl MongoPlugin {
27    pub fn new() -> Self {
28        Self {
29            handlers: RwLock::new(Handle::new()),
30        }
31    }
32
33    pub fn connect(&self, conn_str: &str) -> Result<u8, LabeledError> {
34        let conn = mongodb::sync::Client::with_uri_str(conn_str)
35            .map_err(|err| LabeledError::new(format!("{err}")))
36            .map(|c| c.default_database());
37        match conn {
38            Err(e) => return Err(e),
39            Ok(conn) => match conn {
40                None => return Err(LabeledError::new("No default database in connection url")),
41                Some(db) => {
42                    let mut write_guard = self.handlers.write().expect("write lock should success");
43                    let id = write_guard.inner.len() as u8;
44                    write_guard.inner.insert(id, (db, conn_str.to_string()));
45                    write_guard.current = id;
46                    Ok(id)
47                }
48            },
49        }
50    }
51
52    pub fn list_handles(&self) -> Vec<(u8, String)> {
53        let read_guard = self.handlers.read().expect("read lock should success");
54        let mut result = vec![];
55        for (id, (_, conn_str)) in read_guard.inner.iter() {
56            result.push((*id, conn_str.to_string()));
57        }
58        result
59    }
60
61    pub fn get_handle(&self, id: u8, span: Span) -> Result<Database, LabeledError> {
62        let read_guard = self.handlers.read().expect("read lock should success");
63        let result = read_guard.inner.get(&id).ok_or_else(|| {
64            let err = LabeledError::new("database handle doesn't exist")
65                .with_label("not existed database handle", span)
66                .with_help("You can run `mongoc list` to list all available handles, or `mongoc open` to open a new handle");
67            err
68        })?;
69        Ok(result.0.clone())
70    }
71
72    pub fn remove_handle(&self, id: u8, span: Span) -> Result<(), LabeledError> {
73        let mut write_guard = self.handlers.write().expect("write lock should success");
74        write_guard.inner.remove(&id).ok_or_else(|| {
75            let err = LabeledError::new("database handle doesn't exist")
76                .with_label("not existed database handle", span)
77                .with_help("You can run `mongoc list` to list all available handles, or `mongoc open` to open a new handle");
78            err
79        })?;
80        // if remove current handle, reset the id.
81        if write_guard.current == id {
82            let max_id = write_guard.inner.keys().max().unwrap_or(&0);
83            write_guard.current = *max_id;
84        }
85        Ok(())
86    }
87
88    pub fn select_handle(&self, id: u8, span: Span) -> Result<(), LabeledError> {
89        let mut write_guard = self.handlers.write().expect("read lock should success");
90        if !write_guard.inner.contains_key(&id) {
91            let err = LabeledError::new("database handle doesn't exist")
92                .with_label("not existed database handle", span)
93                .with_help("You can run `mongoc list` to list all available handles, or `mongoc open` to open a new handle");
94            return Err(err);
95        }
96        write_guard.current = id;
97        Ok(())
98    }
99
100    pub fn get_current(&self) -> Result<u8, LabeledError> {
101        let read_guard = self.handlers.read().expect("read lock should success");
102        if read_guard.inner.is_empty() {
103            return Err(LabeledError::new("no database handles available")
104                .with_help("You can run `mongoc open` to open a new handle"));
105        }
106        Ok(read_guard.current)
107    }
108}
109
110impl Plugin for MongoPlugin {
111    fn version(&self) -> String {
112        env!("CARGO_PKG_VERSION").into()
113    }
114
115    fn commands(&self) -> Vec<Box<dyn PluginCommand<Plugin = Self>>> {
116        vec![
117            Box::new(Open),
118            Box::new(List),
119            Box::new(Find),
120            Box::new(FindOne),
121            Box::new(Drop),
122            Box::new(DeleteOne),
123            Box::new(DeleteMany),
124            Box::new(Remove),
125            Box::new(MongoCmd),
126            Box::new(Select),
127            Box::new(ListCollectionNames),
128            Box::new(ListIndexes),
129            Box::new(Count),
130            Box::new(Estimated),
131        ]
132    }
133}
nu_plugin_mongo/lib.rs

nu_plugin_mongo/
lib.rs
