queen_search 0.2.0

use std::collections::HashSet;

use async_trait::async_trait;
use elasticsearch::{http::transport::Transport, Elasticsearch, IndexParts, SearchParts};
use serde_json::{json, Value};

mod model;
pub use model::*;

#[async_trait]
pub trait SearchEngine {
    async fn add_permission(
        &self,
        perm: &InfoLibPermissions,
    ) -> Result<(), Box<dyn std::error::Error>>;
    async fn index(&self, msg: &Message) -> Result<(), Box<dyn std::error::Error>>;
    async fn search_lib(
        &self,
        txt: &str,
        lib: &InfoLibId,
    ) -> Result<Vec<Message>, Box<dyn std::error::Error>>;
    async fn search_user(
        &self,
        txt: &str,
        user: &UserId,
    ) -> Result<Vec<Message>, Box<dyn std::error::Error>>;
}

pub struct ElasticAdapter {
    url: String,
    info_index: String,
    perm_index: String,
    client: Elasticsearch,
}

#[async_trait]
impl SearchEngine for ElasticAdapter {
    async fn add_permission(
        &self,
        perm: &InfoLibPermissions,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let response = self.client
            .index(IndexParts::Index(&self.perm_index))
            .body(json! ({
                "user": perm.user_id().0,
                "lib": perm.info_lib().0,
            }))
            .send()
            .await?;
        let json: Value = response.json().await?;

        log::info!("saved {json}");
        Ok(())
    }
    async fn index(&self, msg: &Message) -> Result<(), Box<dyn std::error::Error>> {
        let response = self.client
            .index(IndexParts::IndexId(&self.info_index, &msg.id()))
            .body(json! ({
                "lib": msg.lib_id().0,
                "content": msg.content(),
            }))
            .send()
            .await?;
        let json: Value = response.json().await?;

        log::info!("saved {json}");
        Ok(())
    }

    async fn search_lib(
        &self,
        txt: &str,
        lib: &InfoLibId,
    ) -> Result<Vec<Message>, Box<dyn std::error::Error>> {
        let search_json = json!({
            "query": {
                "bool" : {
                    "must": [
                        {
                            "match": {
                                "lib": lib.0,
                            }
                        },
                        {
                            "match": {
                                "content": txt
                            }
                        }
                    ]
                }
            }
        });
        let result_json = self.search_json(&self.info_index, &search_json).await?;

        let mut results = vec![];
        for res in result_json["hits"]["hits"]
            .as_array()
            .expect("parse error!")
            .iter()
        {
            let message_text = res["_source"]["content"].as_str().unwrap();
            let msg_id = res["_id"].as_str().unwrap();
            results.push(Message::new(
                lib.clone(),
                message_text.to_owned(),
                msg_id.to_owned(),
            ));
        }
        Ok(results)
    }

    async fn search_user(
        &self,
        txt: &str,
        user: &UserId,
    ) -> Result<Vec<Message>, Box<dyn std::error::Error>> {
        let search_json = json!({
            "query": {
                "bool" : {
                    "must": [
                        {
                            "match": {
                                "user": user.0
                            }
                        }
                    ]
                }
            }
        });
        let result_json = self.search_json(&self.perm_index, &search_json).await?;
        let mut libs = HashSet::new();
        for res in result_json["hits"]["hits"]
            .as_array()
            .expect("parse error!")
            .iter()
        {
            let lib_id = res["_source"]["lib"].as_str().unwrap();
            libs.insert(lib_id);
        }

        let mut msgs = vec![];

        for lib_id in libs {
            let result = self.search_lib(txt, &InfoLibId(lib_id.to_owned())).await;
            if let Ok(result) = result {
                msgs.extend(result);
            }
        }
        Ok(msgs)
    }
}

impl ElasticAdapter {
    fn new(url: &str, info_index: &str, perm_index: &str) -> Result<Self, Box<dyn std::error::Error>> {
        let transport = Transport::single_node(url)?;
        Ok(Self {
            url: url.to_owned(),
            info_index: info_index.to_owned(),
            perm_index: perm_index.to_owned(),
            client: Elasticsearch::new(transport)
        })
    }

    pub async fn search_json(
        &self,
        index_name: &str,
        search_json: &Value,
    ) -> Result<Value, Box<dyn std::error::Error>> {
        log::info!("searching....");
        let response = self.client
            .search(SearchParts::Index(&[index_name]))
            .body(search_json)
            .send()
            .await?;
        let json: Value = response.json().await?;
        log::info!("search ok: {json}");
        Ok(json)
    }
}

#[cfg(test)]
mod tests {
    use log4rs::config::Deserializers;

    use crate::{ElasticAdapter, SearchEngine, InfoLibId, UserId, Message, InfoLibPermissions};

    #[tokio::test]
    async fn test() {
        log4rs::init_file("log4rs.yml", Deserializers::default()).unwrap();
        let se = ElasticAdapter::new(
            "http://8.141.144.181:9043",
            "interlinked",
            "interlinked_perm",
        ).unwrap();
        let test_lib = InfoLibId("test_lib".to_owned());
        let test_user = UserId("test_user".to_owned());
        let msg = Message::new(
            test_lib.clone(),
            "test_content".to_owned(),
            "chat_2_msg_1".to_owned(),
        );
        se.index(&msg).await.unwrap();
        se.add_permission(&InfoLibPermissions::new(test_user.clone(), test_lib)).await.unwrap();

        se.search_user("test", &test_user).await.unwrap();
    }
}