initial commit

src/routes/document.rs

@@ -0,0 +1,266 @@
+use std::collections::{BTreeSet, HashSet};
+
+use actix_web::{delete, get, post, put};
+use actix_web::{web, HttpResponse};
+use indexmap::IndexMap;
+use meilisearch_core::{update, MainReader};
+use serde_json::Value;
+use serde::Deserialize;
+
+use crate::Data;
+use crate::error::{Error, ResponseError};
+use crate::helpers::Authentication;
+use crate::routes::{IndexParam, IndexUpdateResponse};
+
+type Document = IndexMap<String, Value>;
+
+#[derive(Deserialize)]
+struct DocumentParam {
+    index_uid: String,
+    document_id: String,
+}
+
+pub fn services(cfg: &mut web::ServiceConfig) {
+    cfg.service(get_document)
+        .service(delete_document)
+        .service(get_all_documents)
+        .service(add_documents)
+        .service(update_documents)
+        .service(delete_documents)
+        .service(clear_all_documents);
+}
+
+#[get(
+    "/indexes/{index_uid}/documents/{document_id}",
+    wrap = "Authentication::Public"
+)]
+async fn get_document(
+    data: web::Data<Data>,
+    path: web::Path<DocumentParam>,
+) -> Result<HttpResponse, ResponseError> {
+    let index = data
+        .db
+        .open_index(&path.index_uid)
+        .ok_or(Error::index_not_found(&path.index_uid))?;
+
+    let reader = data.db.main_read_txn()?;
+
+    let internal_id = index
+        .main
+        .external_to_internal_docid(&reader, &path.document_id)?
+        .ok_or(Error::document_not_found(&path.document_id))?;
+
+    let document: Document = index
+        .document(&reader, None, internal_id)?
+        .ok_or(Error::document_not_found(&path.document_id))?;
+
+    Ok(HttpResponse::Ok().json(document))
+}
+
+#[delete(
+    "/indexes/{index_uid}/documents/{document_id}",
+    wrap = "Authentication::Private"
+)]
+async fn delete_document(
+    data: web::Data<Data>,
+    path: web::Path<DocumentParam>,
+) -> Result<HttpResponse, ResponseError> {
+    let index = data
+        .db
+        .open_index(&path.index_uid)
+        .ok_or(Error::index_not_found(&path.index_uid))?;
+
+    let mut documents_deletion = index.documents_deletion();
+    documents_deletion.delete_document_by_external_docid(path.document_id.clone());
+
+    let update_id = data.db.update_write(|w| documents_deletion.finalize(w))?;
+
+    Ok(HttpResponse::Accepted().json(IndexUpdateResponse::with_id(update_id)))
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+struct BrowseQuery {
+    offset: Option<usize>,
+    limit: Option<usize>,
+    attributes_to_retrieve: Option<String>,
+}
+
+pub fn get_all_documents_sync(
+    data: &web::Data<Data>,
+    reader: &MainReader,
+    index_uid: &str,
+    offset: usize,
+    limit: usize,
+    attributes_to_retrieve: Option<&String>
+) -> Result<Vec<Document>, Error> {
+    let index = data
+        .db
+        .open_index(index_uid)
+        .ok_or(Error::index_not_found(index_uid))?;
+
+
+    let documents_ids: Result<BTreeSet<_>, _> = index
+        .documents_fields_counts
+        .documents_ids(reader)?
+        .skip(offset)
+        .take(limit)
+        .collect();
+
+    let attributes: Option<HashSet<&str>> = attributes_to_retrieve
+        .map(|a| a.split(',').collect());
+
+    let mut documents = Vec::new();
+    for document_id in documents_ids? {
+        if let Ok(Some(document)) =
+            index.document::<Document>(reader, attributes.as_ref(), document_id)
+        {
+            documents.push(document);
+        }
+    }
+
+    Ok(documents)
+}
+
+#[get("/indexes/{index_uid}/documents", wrap = "Authentication::Public")]
+async fn get_all_documents(
+    data: web::Data<Data>,
+    path: web::Path<IndexParam>,
+    params: web::Query<BrowseQuery>,
+) -> Result<HttpResponse, ResponseError> {
+    let offset = params.offset.unwrap_or(0);
+    let limit = params.limit.unwrap_or(20);
+    let index_uid = &path.index_uid;
+    let reader = data.db.main_read_txn()?;
+    
+    let documents = get_all_documents_sync(
+        &data,
+        &reader,
+        index_uid,
+        offset,
+        limit,
+        params.attributes_to_retrieve.as_ref()
+    )?;
+
+    Ok(HttpResponse::Ok().json(documents))
+}
+
+fn find_primary_key(document: &IndexMap<String, Value>) -> Option<String> {
+    for key in document.keys() {
+        if key.to_lowercase().contains("id") {
+            return Some(key.to_string());
+        }
+    }
+    None
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+struct UpdateDocumentsQuery {
+    primary_key: Option<String>,
+}
+
+async fn update_multiple_documents(
+    data: web::Data<Data>,
+    path: web::Path<IndexParam>,
+    params: web::Query<UpdateDocumentsQuery>,
+    body: web::Json<Vec<Document>>,
+    is_partial: bool,
+) -> Result<HttpResponse, ResponseError> {
+    let update_id = data.get_or_create_index(&path.index_uid, |index| {
+        let reader = data.db.main_read_txn()?;
+
+        let mut schema = index
+            .main
+            .schema(&reader)?
+            .ok_or(meilisearch_core::Error::SchemaMissing)?;
+
+        if schema.primary_key().is_none() {
+            let id = match &params.primary_key {
+                Some(id) => id.to_string(),
+                None => body
+                    .first()
+                    .and_then(find_primary_key)
+                    .ok_or(meilisearch_core::Error::MissingPrimaryKey)?,
+            };
+
+            schema.set_primary_key(&id).map_err(Error::bad_request)?;
+
+            data.db.main_write(|w| index.main.put_schema(w, &schema))?;
+        }
+
+        let mut document_addition = if is_partial {
+            index.documents_partial_addition()
+        } else {
+            index.documents_addition()
+        };
+
+        for document in body.into_inner() {
+            document_addition.update_document(document);
+        }
+
+        Ok(data.db.update_write(|w| document_addition.finalize(w))?)
+    })?;
+    return Ok(HttpResponse::Accepted().json(IndexUpdateResponse::with_id(update_id)));
+}
+
+#[post("/indexes/{index_uid}/documents", wrap = "Authentication::Private")]
+async fn add_documents(
+    data: web::Data<Data>,
+    path: web::Path<IndexParam>,
+    params: web::Query<UpdateDocumentsQuery>,
+    body: web::Json<Vec<Document>>,
+) -> Result<HttpResponse, ResponseError> {
+    update_multiple_documents(data, path, params, body, false).await
+}
+
+#[put("/indexes/{index_uid}/documents", wrap = "Authentication::Private")]
+async fn update_documents(
+    data: web::Data<Data>,
+    path: web::Path<IndexParam>,
+    params: web::Query<UpdateDocumentsQuery>,
+    body: web::Json<Vec<Document>>,
+) -> Result<HttpResponse, ResponseError> {
+    update_multiple_documents(data, path, params, body, true).await
+}
+
+#[post(
+    "/indexes/{index_uid}/documents/delete-batch",
+    wrap = "Authentication::Private"
+)]
+async fn delete_documents(
+    data: web::Data<Data>,
+    path: web::Path<IndexParam>,
+    body: web::Json<Vec<Value>>,
+) -> Result<HttpResponse, ResponseError> {
+    let index = data
+        .db
+        .open_index(&path.index_uid)
+        .ok_or(Error::index_not_found(&path.index_uid))?;
+
+    let mut documents_deletion = index.documents_deletion();
+
+    for document_id in body.into_inner() {
+        let document_id = update::value_to_string(&document_id);
+        documents_deletion.delete_document_by_external_docid(document_id);
+    }
+
+    let update_id = data.db.update_write(|w| documents_deletion.finalize(w))?;
+
+    Ok(HttpResponse::Accepted().json(IndexUpdateResponse::with_id(update_id)))
+}
+
+#[delete("/indexes/{index_uid}/documents", wrap = "Authentication::Private")]
+async fn clear_all_documents(
+    data: web::Data<Data>,
+    path: web::Path<IndexParam>,
+) -> Result<HttpResponse, ResponseError> {
+    let index = data
+        .db
+        .open_index(&path.index_uid)
+        .ok_or(Error::index_not_found(&path.index_uid))?;
+
+    let update_id = data.db.update_write(|w| index.clear_all(w))?;
+
+    Ok(HttpResponse::Accepted().json(IndexUpdateResponse::with_id(update_id)))
+}