path: root/src/server.rs

use std::str::FromStr;
use std::path::PathBuf;
use std::path::Path;

use anyhow::anyhow;
use anyhow::Error;
use anyhow::Context;
use anyhow::Result;
use tonic::Request;
use tonic::Response;
use tonic::Status;
use resiter::AndThen;
use resiter::Map;

use fss::VersionReply;
use fss::VersionRequest;
use fss::IndexFileRequest;
use fss::IndexFileReply;
use fss::SearchQueryRequest;
use fss::SearchResponse;

mod config;
mod cli;
mod ft;
mod schema;

mod fss {
    tonic::include_proto!("fss"); // The string specified here must match the proto package name
}

#[derive(getset::CopyGetters)]
pub struct Server {
    index: tantivy::Index,

    #[getset(get_copy = "pub")]
    field_path: tantivy::schema::Field,

    #[getset(get_copy = "pub")]
    field_ft: tantivy::schema::Field,

    #[getset(get_copy = "pub")]
    field_size: tantivy::schema::Field,

    #[getset(get_copy = "pub")]
    field_created: tantivy::schema::Field,

    #[getset(get_copy = "pub")]
    field_modified: tantivy::schema::Field,

    #[getset(get_copy = "pub")]
    field_indexed: tantivy::schema::Field,

    #[getset(get_copy = "pub")]
    field_body: tantivy::schema::Field,
}

impl Server {
    fn write_file_to_index(&self, filepath: &Path) -> Result<()> {
        let ext = filepath
            .extension()
            .map(ToOwned::to_owned)
            .and_then(|osstr| osstr.to_str().map(|s| s.to_string()))
            .ok_or_else(|| anyhow!("Path {} is not UTF8", filepath.display()))?;

        let mut doc = tantivy::Document::default();
        doc.add_text(self.field_path(), filepath.display().to_string());
        doc.add_text(self.field_ft(), &ext);
        {
            use std::os::linux::fs::MetadataExt;
            let meta = filepath.metadata()?;
            doc.add_i64(self.field_created(), meta.st_atime());
            doc.add_i64(self.field_modified(), meta.st_mtime());
            if let Ok(now) = std::time::SystemTime::now().duration_since(std::time::SystemTime::UNIX_EPOCH).map(|d| d.as_secs()) {
                doc.add_u64(self.field_indexed(), now)
            }
        }

        let doc = self.parse_for_file(doc, filepath, &ext)?;
        let mut index_writer = self.index.writer(50_000_000)?;
        index_writer.add_document(doc);
        index_writer.commit()?;
        Ok(())
    }

    fn parse_for_file(&self, doc: tantivy::Document, filepath: &Path, ext: &str) -> Result<tantivy::Document> {
        use crate::ft::FileTypeParser;
        use crate::ft::TextFileParser;

        match ext {
            "txt" => Ok(TextFileParser {}),
            _ => Err(anyhow!("No parser available for {}", ext))
        }?
        .parse(&self, filepath, ext, doc)
    }

    fn search(&self, query_str: &str) -> Result<Vec<PathBuf>> {
        let reader = self.index
            .reader_builder()
            .reload_policy(tantivy::ReloadPolicy::OnCommit)
            .try_into()?;

         let searcher = reader.searcher();
         let query_parser = tantivy::query::QueryParser::for_index(&self.index, vec![
             self.field_path,
             self.field_ft,
             self.field_size,
             self.field_created,
             self.field_modified,
             self.field_indexed,
             self.field_body,
         ]);
         let query = query_parser.parse_query(query_str)?;

         let top_docs = searcher.search(&query, &tantivy::collector::TopDocs::with_limit(10))?;

         let results = top_docs
             .into_iter()
             .map(|(_score, adr)| searcher.doc(adr).map_err(Error::from))
             .and_then_ok(|retr| {
                retr.get_all(self.field_path)
                    .map(|value| {
                        value.text().ok_or_else(|| anyhow!("Not a text value.."))
                    })
                    .map_ok(PathBuf::from)
                    .collect::<Result<Vec<_>>>()
             })
            .collect::<Result<Vec<Vec<PathBuf>>>>()?
            .into_iter()
            .map(Vec::into_iter)
            .flatten()
            .collect();
        Ok(results)
    }
}


#[tonic::async_trait]
impl fss::fss_server::Fss for Server {
    async fn get_version(&self, _request: Request<VersionRequest>) -> Result<Response<VersionReply>, Status> {
        let reply = fss::VersionReply {
            version: version::version!().to_string(),
        };

        Ok(Response::new(reply))
    }

    async fn index_file(&self, request: Request<IndexFileRequest>) -> Result<Response<IndexFileReply>, Status> {
        let error = match self.write_file_to_index(request.get_ref().path.as_ref()) {
            Ok(()) => false,
            Err(e) => {
                log::error!("Error writing to index: {} -> {:?}", request.get_ref().path, e);
                true
            }
        };

        let reply = fss::IndexFileReply {
            error
        };
        Ok(Response::new(reply))
    }

    async fn search_query(&self, request: Request<SearchQueryRequest>) -> Result<Response<SearchResponse>, Status> {
        let (success, pathes) = match self.search(&request.get_ref().query) {
            Ok(pathes) => (true, pathes),
            Err(e) => {
                log::error!("Error searching with '{}' -> {:?}", request.get_ref().query, e);
                (false, vec![])
            }
        };

        let reply = fss::SearchResponse {
            success,
            pathes: pathes.into_iter().map(|pb| pb.display().to_string()).collect(),
        };
        Ok(Response::new(reply))
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = crate::cli::server_app();
    let matches = cli.get_matches();
    let _ = env_logger::try_init()?;
    let mut config = ::config::Config::default();
    {
        let xdg = xdg::BaseDirectories::with_prefix("fss")?;
        let xdg_config = xdg.find_config_file("config.toml")
            .ok_or_else(|| anyhow!("No configuration file found with XDG: {}", xdg.get_config_home().display()))?;

        log::debug!("Configuration file found with XDG: {}", xdg_config.display());
        config.merge(::config::File::from(xdg_config).required(false))
            .context("Failed to load config.toml from XDG configuration directory")?;
    }
    let config = config.try_into::<crate::config::Config>()?;
    let server_addr = matches.value_of("server_addr").unwrap_or_else(|| config.server_addr());
    let server_port = matches.value_of("server_port").map(usize::from_str).unwrap_or_else(|| Ok(*config.server_port()))?;

    let addr = format!("{}:{}", server_addr, server_port).parse()?;

    let server = {
        let index_path = tantivy::directory::MmapDirectory::open(config.database_path())?;
        let schema = crate::schema::schema();

        let index = tantivy::Index::open_or_create(index_path, schema.clone())?;
        let field_path = schema.get_field("path").ok_or_else(|| anyhow!("BUG"))?;
        let field_ft = schema.get_field("ft").ok_or_else(|| anyhow!("BUG"))?;
        let field_size = schema.get_field("size").ok_or_else(|| anyhow!("BUG"))?;
        let field_created = schema.get_field("created").ok_or_else(|| anyhow!("BUG"))?;
        let field_modified = schema.get_field("modified").ok_or_else(|| anyhow!("BUG"))?;
        let field_indexed = schema.get_field("indexed").ok_or_else(|| anyhow!("BUG"))?;
        let field_body = schema.get_field("body").ok_or_else(|| anyhow!("BUG"))?;

        Server {
            index,

            field_path,
            field_ft,
            field_size,
            field_created,
            field_modified,
            field_indexed,
            field_body,
        }
    };

    tonic::transport::Server::builder()
        .add_service(fss::fss_server::FssServer::new(server))
        .serve(addr)
        .await
        .map_err(Error::from)
        .map(|_| ())
}