initial commit

src/data.rs

@@ -0,0 +1,162 @@
+use std::error::Error;
+use std::ops::Deref;
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use meilisearch_core::{Database, DatabaseOptions, Index};
+use sha2::Digest;
+
+use crate::error::{Error as MSError, ResponseError};
+use crate::index_update_callback;
+use crate::option::Opt;
+
+#[derive(Clone)]
+pub struct Data {
+    inner: Arc<DataInner>,
+}
+
+impl Deref for Data {
+    type Target = DataInner;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+#[derive(Clone)]
+pub struct DataInner {
+    pub db: Arc<Database>,
+    pub db_path: String,
+    pub dumps_dir: PathBuf,
+    pub dump_batch_size: usize,
+    pub api_keys: ApiKeys,
+    pub server_pid: u32,
+    pub http_payload_size_limit: usize,
+}
+
+#[derive(Clone)]
+pub struct ApiKeys {
+    pub public: Option<String>,
+    pub private: Option<String>,
+    pub master: Option<String>,
+}
+
+impl ApiKeys {
+    pub fn generate_missing_api_keys(&mut self) {
+        if let Some(master_key) = &self.master {
+            if self.private.is_none() {
+                let key = format!("{}-private", master_key);
+                let sha = sha2::Sha256::digest(key.as_bytes());
+                self.private = Some(format!("{:x}", sha));
+            }
+            if self.public.is_none() {
+                let key = format!("{}-public", master_key);
+                let sha = sha2::Sha256::digest(key.as_bytes());
+                self.public = Some(format!("{:x}", sha));
+            }
+        }
+    }
+}
+
+impl Data {
+    pub fn new(opt: Opt) -> Result<Data, Box<dyn Error>> {
+        let db_path = opt.db_path.clone();
+        let dumps_dir = opt.dumps_dir.clone();
+        let dump_batch_size = opt.dump_batch_size;
+        let server_pid = std::process::id();
+
+        let db_opt = DatabaseOptions {
+            main_map_size: opt.max_mdb_size,
+            update_map_size: opt.max_udb_size,
+        };
+
+        let http_payload_size_limit = opt.http_payload_size_limit;
+
+        let db = Arc::new(Database::open_or_create(opt.db_path, db_opt)?);
+
+        let mut api_keys = ApiKeys {
+            master: opt.master_key,
+            private: None,
+            public: None,
+        };
+
+        api_keys.generate_missing_api_keys();
+
+        let inner_data = DataInner {
+            db: db.clone(),
+            db_path,
+            dumps_dir,
+            dump_batch_size,
+            api_keys,
+            server_pid,
+            http_payload_size_limit,
+        };
+
+        let data = Data {
+            inner: Arc::new(inner_data),
+        };
+
+        let callback_context = data.clone();
+        db.set_update_callback(Box::new(move |index_uid, status| {
+            index_update_callback(&index_uid, &callback_context, status);
+        }));
+
+        Ok(data)
+    }
+
+    fn create_index(&self, uid: &str) -> Result<Index, ResponseError> {
+        if !uid
+            .chars()
+            .all(|x| x.is_ascii_alphanumeric() || x == '-' || x == '_')
+        {
+            return Err(MSError::InvalidIndexUid.into());
+        }
+
+        let created_index = self.db.create_index(&uid).map_err(|e| match e {
+            meilisearch_core::Error::IndexAlreadyExists => e.into(),
+            _ => ResponseError::from(MSError::create_index(e)),
+        })?;
+
+        self.db.main_write::<_, _, ResponseError>(|mut writer| {
+            created_index.main.put_name(&mut writer, uid)?;
+
+            created_index
+                .main
+                .created_at(&writer)?
+                .ok_or(MSError::internal("Impossible to read created at"))?;
+
+            created_index
+                .main
+                .updated_at(&writer)?
+                .ok_or(MSError::internal("Impossible to read updated at"))?;
+            Ok(())
+        })?;
+
+        Ok(created_index)
+    }
+
+    pub fn get_or_create_index<F, R>(&self, uid: &str, f: F) -> Result<R, ResponseError>
+    where
+        F: FnOnce(&Index) -> Result<R, ResponseError>,
+    {
+        let mut index_has_been_created = false;
+
+        let index = match self.db.open_index(&uid) {
+            Some(index) => index,
+            None => {
+                index_has_been_created = true;
+                self.create_index(&uid)?
+            }
+        };
+
+        match f(&index) {
+            Ok(r) => Ok(r),
+            Err(err) => {
+                if index_has_been_created {
+                    let _ = self.db.delete_index(&uid);
+                }
+                Err(err)
+            }
+        }
+    }
+}