letterbox/server/src/bin/server.rs

// Rocket generates a lot of warnings for handlers
// TODO: figure out why
#![allow(unreachable_patterns)]
#[macro_use]
extern crate rocket;
use std::{error::Error, io::Cursor, str::FromStr};

use async_graphql::{http::GraphiQLSource, EmptySubscription, Schema};
use async_graphql_rocket::{GraphQLQuery, GraphQLRequest, GraphQLResponse};
use glog::Flags;
use notmuch::{Notmuch, NotmuchError, ThreadSet};
use rocket::{
    fairing::AdHoc,
    http::{ContentType, Header},
    request::Request,
    response::{content, Debug, Responder},
    serde::json::Json,
    Response, State,
};
use rocket_cors::{AllowedHeaders, AllowedOrigins};
use server::{
    config::Config,
    error::ServerError,
    graphql::{Attachment, GraphqlSchema, Mutation, QueryRoot},
    nm::{attachment_bytes, cid_attachment_bytes},
};
use sqlx::postgres::PgPool;
use tantivy::{Index, IndexWriter};

#[get("/refresh")]
async fn refresh(nm: &State<Notmuch>) -> Result<Json<String>, Debug<NotmuchError>> {
    Ok(Json(String::from_utf8_lossy(&nm.new()?).to_string()))
}

#[get("/show/<query>/pretty")]
async fn show_pretty(
    nm: &State<Notmuch>,
    query: &str,
) -> Result<Json<ThreadSet>, Debug<ServerError>> {
    let query = urlencoding::decode(query).map_err(|e| ServerError::from(NotmuchError::from(e)))?;
    let res = nm.show(&query).map_err(ServerError::from)?;
    Ok(Json(res))
}

#[get("/show/<query>")]
async fn show(nm: &State<Notmuch>, query: &str) -> Result<Json<ThreadSet>, Debug<NotmuchError>> {
    let query = urlencoding::decode(query).map_err(NotmuchError::from)?;
    let res = nm.show(&query)?;
    Ok(Json(res))
}

struct InlineAttachmentResponder(Attachment);

impl<'r, 'o: 'r> Responder<'r, 'o> for InlineAttachmentResponder {
    fn respond_to(self, _: &'r Request<'_>) -> rocket::response::Result<'o> {
        let mut resp = Response::build();
        if let Some(filename) = self.0.filename {
            resp.header(Header::new(
                "Content-Disposition",
                format!(r#"inline; filename="{}""#, filename),
            ));
        }
        if let Some(content_type) = self.0.content_type {
            if let Some(ct) = ContentType::parse_flexible(&content_type) {
                resp.header(ct);
            }
        }
        resp.sized_body(self.0.bytes.len(), Cursor::new(self.0.bytes))
            .ok()
    }
}

struct DownloadAttachmentResponder(Attachment);

impl<'r, 'o: 'r> Responder<'r, 'o> for DownloadAttachmentResponder {
    fn respond_to(self, _: &'r Request<'_>) -> rocket::response::Result<'o> {
        let mut resp = Response::build();
        if let Some(filename) = self.0.filename {
            resp.header(Header::new(
                "Content-Disposition",
                format!(r#"attachment; filename="{}""#, filename),
            ));
        }
        if let Some(content_type) = self.0.content_type {
            if let Some(ct) = ContentType::parse_flexible(&content_type) {
                resp.header(ct);
            }
        }
        resp.sized_body(self.0.bytes.len(), Cursor::new(self.0.bytes))
            .ok()
    }
}

#[get("/cid/<id>/<cid>")]
async fn view_cid(
    nm: &State<Notmuch>,
    id: &str,
    cid: &str,
) -> Result<InlineAttachmentResponder, Debug<ServerError>> {
    let mid = if id.starts_with("id:") {
        id.to_string()
    } else {
        format!("id:{}", id)
    };
    info!("view cid attachment {mid} {cid}");
    let attachment = cid_attachment_bytes(nm, &mid, &cid)?;
    Ok(InlineAttachmentResponder(attachment))
}

#[get("/view/attachment/<id>/<idx>/<_>")]
async fn view_attachment(
    nm: &State<Notmuch>,
    id: &str,
    idx: &str,
) -> Result<InlineAttachmentResponder, Debug<ServerError>> {
    let mid = if id.starts_with("id:") {
        id.to_string()
    } else {
        format!("id:{}", id)
    };
    info!("view attachment {mid} {idx}");
    let idx: Vec<_> = idx
        .split('.')
        .map(|s| s.parse().expect("not a usize"))
        .collect();
    let attachment = attachment_bytes(nm, &mid, &idx)?;
    Ok(InlineAttachmentResponder(attachment))
}

#[get("/download/attachment/<id>/<idx>/<_>")]
async fn download_attachment(
    nm: &State<Notmuch>,
    id: &str,
    idx: &str,
) -> Result<DownloadAttachmentResponder, Debug<ServerError>> {
    let mid = if id.starts_with("id:") {
        id.to_string()
    } else {
        format!("id:{}", id)
    };
    info!("download attachment {mid} {idx}");
    let idx: Vec<_> = idx
        .split('.')
        .map(|s| s.parse().expect("not a usize"))
        .collect();
    let attachment = attachment_bytes(nm, &mid, &idx)?;
    Ok(DownloadAttachmentResponder(attachment))
}

#[get("/original/<id>")]
async fn original(
    nm: &State<Notmuch>,
    id: &str,
) -> Result<(ContentType, Vec<u8>), Debug<NotmuchError>> {
    let mid = if id.starts_with("id:") {
        id.to_string()
    } else {
        format!("id:{}", id)
    };
    let res = nm.show_original(&mid)?;
    Ok((ContentType::Plain, res))
}

#[rocket::get("/")]
fn graphiql() -> content::RawHtml<String> {
    content::RawHtml(GraphiQLSource::build().endpoint("/api/graphql").finish())
}

#[rocket::post("/create-news-db")]
fn create_news_db(config: &State<Config>) -> Result<String, Debug<ServerError>> {
    create_news_db_impl(config)?;
    Ok(format!(
        "DB created in {}\n",
        config.newsreader_tantivy_db_path
    ))
}
fn create_news_db_impl(config: &Config) -> Result<(), ServerError> {
    std::fs::remove_dir_all(&config.newsreader_tantivy_db_path).map_err(ServerError::from)?;
    std::fs::create_dir_all(&config.newsreader_tantivy_db_path).map_err(ServerError::from)?;
    use tantivy::schema::*;
    let mut schema_builder = Schema::builder();
    schema_builder.add_text_field("site", STRING | STORED);
    schema_builder.add_text_field("title", TEXT | STORED);
    schema_builder.add_text_field("summary", TEXT);
    schema_builder.add_text_field("link", STRING | STORED);
    schema_builder.add_date_field("date", FAST);
    schema_builder.add_bool_field("is_read", FAST);
    schema_builder.add_text_field("uid", STRING | STORED);
    schema_builder.add_i64_field("id", FAST);

    let schema = schema_builder.build();
    Index::create_in_dir(&config.newsreader_tantivy_db_path, schema).map_err(ServerError::from)?;
    Ok(())
}

#[rocket::post("/reindex-news-db")]
async fn reindex_news_db(
    pool: &State<PgPool>,
    config: &State<Config>,
) -> Result<String, Debug<ServerError>> {
    use tantivy::{doc, Term};

    let start_time = std::time::Instant::now();
    let pool: &PgPool = pool;

    let index =
        Index::open_in_dir(&config.newsreader_tantivy_db_path).map_err(ServerError::from)?;
    let mut index_writer = index.writer(50_000_000).map_err(ServerError::from)?;
    let schema = index.schema();
    let site = schema.get_field("site").map_err(ServerError::from)?;
    let title = schema.get_field("title").map_err(ServerError::from)?;
    let summary = schema.get_field("summary").map_err(ServerError::from)?;
    let link = schema.get_field("link").map_err(ServerError::from)?;
    let date = schema.get_field("date").map_err(ServerError::from)?;
    let is_read = schema.get_field("is_read").map_err(ServerError::from)?;
    let uid = schema.get_field("uid").map_err(ServerError::from)?;
    let id = schema.get_field("id").map_err(ServerError::from)?;

    let rows = sqlx::query_file!("sql/all-posts.sql")
        .fetch_all(pool)
        .await
        .map_err(ServerError::from)?;

    let total = rows.len();
    for (i, r) in rows.into_iter().enumerate() {
        if i % 10_000 == 0 {
            info!(
                "{i}/{total} processed, elapsed {:.2}s",
                start_time.elapsed().as_secs_f32()
            );
        }
        let id_term = Term::from_field_text(uid, &r.uid);
        index_writer.delete_term(id_term);
        index_writer
            .add_document(doc!(
                site =>  r.site.expect("UNKOWN_SITE"),
                title =>  r.title.expect("UNKOWN_TITLE"),
                // TODO: clean and extract text from HTML
                summary => r.summary.expect("UNKNOWN_SUMMARY"),
                link => r.link.expect("link"),
                date => tantivy::DateTime::from_primitive(r.date.expect("date")),
                is_read => r.is_read.expect("is_read"),
                uid => r.uid,
                id => r.id as i64,
            ))
            .map_err(ServerError::from)?;
    }

    index_writer.commit().map_err(ServerError::from)?;

    info!("took {:.2}s to reindex", start_time.elapsed().as_secs_f32());
    Ok(format!(
        "DB openned in {}\n",
        config.newsreader_tantivy_db_path
    ))
}

#[rocket::get("/search-news-db")]
fn search_news_db(
    index: &State<tantivy::Index>,
    reader: &State<tantivy::IndexReader>,
) -> Result<String, Debug<ServerError>> {
    use tantivy::{collector::TopDocs, query::QueryParser, Document, TantivyDocument};

    let searcher = reader.searcher();
    let schema = index.schema();
    let site = schema.get_field("site").map_err(ServerError::from)?;
    let title = schema.get_field("title").map_err(ServerError::from)?;
    let summary = schema.get_field("summary").map_err(ServerError::from)?;
    let query_parser = QueryParser::for_index(&index, vec![site, title, summary]);

    let query = query_parser
        .parse_query("grapheme")
        .map_err(ServerError::from)?;
    let top_docs = searcher
        .search(&query, &TopDocs::with_limit(10))
        .map_err(ServerError::from)?;
    let mut results = vec![];
    info!("search found {} docs", top_docs.len());
    for (_score, doc_address) in top_docs {
        let retrieved_doc: TantivyDocument =
            searcher.doc(doc_address).map_err(ServerError::from)?;
        results.push(format!("{}", retrieved_doc.to_json(&schema)));
    }

    Ok(format!("{}", results.join(" ")))
}

#[rocket::get("/graphql?<query..>")]
async fn graphql_query(schema: &State<GraphqlSchema>, query: GraphQLQuery) -> GraphQLResponse {
    query.execute(schema.inner()).await
}

#[rocket::post("/graphql", data = "<request>", format = "application/json")]
async fn graphql_request(
    schema: &State<GraphqlSchema>,
    request: GraphQLRequest,
) -> GraphQLResponse {
    request.execute(schema.inner()).await
}

#[rocket::main]
async fn main() -> Result<(), Box<dyn Error>> {
    glog::new()
        .init(Flags {
            colorlogtostderr: true,
            //alsologtostderr: true, // use logtostderr to only write to stderr and not to files
            logtostderr: true,
            ..Default::default()
        })
        .unwrap();
    build_info::build_info!(fn bi);
    info!("Build Info: {}", shared::build_version(bi));
    let allowed_origins = AllowedOrigins::all();
    let cors = rocket_cors::CorsOptions {
        allowed_origins,
        allowed_methods: vec!["Get"]
            .into_iter()
            .map(|s| FromStr::from_str(s).unwrap())
            .collect(),
        allowed_headers: AllowedHeaders::some(&["Authorization", "Accept"]),
        allow_credentials: true,
        ..Default::default()
    }
    .to_cors()?;

    let rkt = rocket::build()
        .mount(
            shared::urls::MOUNT_POINT,
            routes![
                create_news_db,
                reindex_news_db,
                search_news_db,
                original,
                refresh,
                show_pretty,
                show,
                graphql_query,
                graphql_request,
                graphiql,
                view_cid,
                view_attachment,
                download_attachment,
            ],
        )
        .attach(cors)
        .attach(AdHoc::config::<Config>());

    let config: Config = rkt.figment().extract()?;
    if !std::fs::exists(&config.slurp_cache_path)? {
        info!("Creating slurp cache @ '{}'", &config.slurp_cache_path);
        std::fs::create_dir_all(&config.slurp_cache_path)?;
    }
    let pool = PgPool::connect(&config.newsreader_database_url).await?;
    let tantivy_newsreader_index = match Index::open_in_dir(&config.newsreader_tantivy_db_path) {
        Ok(idx) => idx,
        Err(_) => {
            create_news_db_impl(&config)?;
            Index::open_in_dir(&config.newsreader_tantivy_db_path)?
        }
    };
    let tantivy_newsreader_reader = tantivy_newsreader_index.reader()?;
    let schema = Schema::build(QueryRoot, Mutation, EmptySubscription)
        .data(Notmuch::default())
        .data(config)
        .data(pool.clone())
        .extension(async_graphql::extensions::Logger)
        .finish();

    let rkt = rkt
        .manage(schema)
        .manage(pool)
        .manage(Notmuch::default())
        .manage(tantivy_newsreader_index)
        .manage(tantivy_newsreader_reader);
    //.manage(Notmuch::with_config("../notmuch/testdata/notmuch.config"))

    rkt.launch().await?;
    Ok(())
}