diff options
Diffstat (limited to 'src/server.rs')
-rw-r--r-- | src/server.rs | 179 |
1 files changed, 179 insertions, 0 deletions
diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..bb7a550 --- /dev/null +++ b/src/server.rs @@ -0,0 +1,179 @@ +use std::str::FromStr; +use std::path::PathBuf; + +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 schema; + +mod fss { + tonic::include_proto!("fss"); // The string specified here must match the proto package name +} + +pub struct Server { + index: tantivy::Index, + + field_path: tantivy::schema::Field, + field_ft: tantivy::schema::Field, + field_body: tantivy::schema::Field, +} + +impl Server { + fn write_file_to_index(&self, filepath: &str) -> Result<()> { + let mut index_writer = self.index.writer(50_000_000)?; + let path = PathBuf::from(filepath); + + let filetype = path.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))?; + + let mut doc = tantivy::Document::default(); + doc.add_text(self.field_path, filepath); + doc.add_text(self.field_ft, &filetype); + + doc.add_text(self.field_body, std::fs::read_to_string(path)?); + + index_writer.add_document(doc); + index_writer.commit()?; + Ok(()) + } + + 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_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) { + 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_body = schema.get_field("body").ok_or_else(|| anyhow!("BUG"))?; + + Server { + index, + + field_path, + field_ft, + field_body, + } + }; + + tonic::transport::Server::builder() + .add_service(fss::fss_server::FssServer::new(server)) + .serve(addr) + .await + .map_err(Error::from) + .map(|_| ()) +} + |