From f4240aa62b47850020aa8c3e164d6d3544626f53 Mon Sep 17 00:00:00 2001 From: Ellie Huxtable Date: Wed, 13 Apr 2022 18:29:18 +0100 Subject: [PATCH] Initial implementation of calendar API (#298) This can be used in the future for sync so that we can be more intelligent with what we're doing, and only sync up what's needed I'd like to eventually replace this with something more like a merkle tree, hence the hash field I've exposed, but that can come later Although this does include a much larger number of count queries, it should also be significantly more cache-able. I'll follow up with that later, and also follow up with using this for sync :) --- Cargo.lock | 10 ++ atuin-common/src/calendar.rs | 15 ++ atuin-common/src/utils.rs | 17 +++ atuin-server/Cargo.toml | 1 + atuin-server/src/calendar.rs | 15 ++ atuin-server/src/database.rs | 209 ++++++++++++++++++++++++++- atuin-server/src/handlers/history.rs | 48 +++++- atuin-server/src/lib.rs | 1 + atuin-server/src/router.rs | 2 +- 9 files changed, 313 insertions(+), 5 deletions(-) create mode 100644 atuin-common/src/calendar.rs create mode 100644 atuin-server/src/calendar.rs diff --git a/Cargo.lock b/Cargo.lock index e7fe019..526d6bc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -144,6 +144,7 @@ dependencies = [ "axum", "base64", "chrono", + "chronoutil", "config", "eyre", "fs-err", @@ -332,6 +333,15 @@ dependencies = [ "scanlex", ] +[[package]] +name = "chronoutil" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43a58c924bb772aa201da3acf5308c46b60275c64e6d3bc89c23dd63d71e83fd" +dependencies = [ + "chrono", +] + [[package]] name = "clap" version = "3.1.8" diff --git a/atuin-common/src/calendar.rs b/atuin-common/src/calendar.rs new file mode 100644 index 0000000..d576f1a --- /dev/null +++ b/atuin-common/src/calendar.rs @@ -0,0 +1,15 @@ +// Calendar data + +pub enum TimePeriod { + YEAR, + MONTH, + DAY, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct TimePeriodInfo { + pub count: u64, + + // TODO: Use this for merkle tree magic + pub hash: String, +} diff --git a/atuin-common/src/utils.rs b/atuin-common/src/utils.rs index 7fe0c30..35647bd 100644 --- a/atuin-common/src/utils.rs +++ b/atuin-common/src/utils.rs @@ -1,5 +1,6 @@ use std::path::PathBuf; +use chrono::NaiveDate; use crypto::digest::Digest; use crypto::sha2::Sha256; use sodiumoxide::crypto::pwhash::argon2id13; @@ -51,6 +52,22 @@ pub fn data_dir() -> PathBuf { data_dir.join("atuin") } +pub fn get_days_from_month(year: i32, month: u32) -> i64 { + NaiveDate::from_ymd( + match month { + 12 => year + 1, + _ => year, + }, + match month { + 12 => 1, + _ => month + 1, + }, + 1, + ) + .signed_duration_since(NaiveDate::from_ymd(year, month, 1)) + .num_days() +} + #[cfg(test)] mod tests { use super::*; diff --git a/atuin-server/Cargo.toml b/atuin-server/Cargo.toml index 17e5d72..871a31f 100644 --- a/atuin-server/Cargo.toml +++ b/atuin-server/Cargo.toml @@ -30,3 +30,4 @@ async-trait = "0.1.49" axum = "0.5" http = "0.2" fs-err = "2.7" +chronoutil = "0.2.3" diff --git a/atuin-server/src/calendar.rs b/atuin-server/src/calendar.rs new file mode 100644 index 0000000..d576f1a --- /dev/null +++ b/atuin-server/src/calendar.rs @@ -0,0 +1,15 @@ +// Calendar data + +pub enum TimePeriod { + YEAR, + MONTH, + DAY, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct TimePeriodInfo { + pub count: u64, + + // TODO: Use this for merkle tree magic + pub hash: String, +} diff --git a/atuin-server/src/database.rs b/atuin-server/src/database.rs index 7de2a6f..9043c2d 100644 --- a/atuin-server/src/database.rs +++ b/atuin-server/src/database.rs @@ -1,12 +1,19 @@ use async_trait::async_trait; +use std::collections::HashMap; use eyre::{eyre, Result}; use sqlx::postgres::PgPoolOptions; use crate::settings::HISTORY_PAGE_SIZE; +use super::calendar::{TimePeriod, TimePeriodInfo}; use super::models::{History, NewHistory, NewSession, NewUser, Session, User}; +use chrono::{Datelike, TimeZone}; +use chronoutil::RelativeDuration; + +use atuin_common::utils::get_days_from_month; + #[async_trait] pub trait Database { async fn get_session(&self, token: &str) -> Result; @@ -18,14 +25,36 @@ pub trait Database { async fn add_user(&self, user: &NewUser) -> Result; async fn count_history(&self, user: &User) -> Result; + + async fn count_history_range( + &self, + user: &User, + start: chrono::NaiveDateTime, + end: chrono::NaiveDateTime, + ) -> Result; + async fn count_history_day(&self, user: &User, date: chrono::NaiveDate) -> Result; + async fn count_history_month(&self, user: &User, date: chrono::NaiveDate) -> Result; + async fn count_history_year(&self, user: &User, year: i32) -> Result; + async fn list_history( &self, user: &User, - created_since: chrono::NaiveDateTime, + created_after: chrono::NaiveDateTime, since: chrono::NaiveDateTime, host: &str, ) -> Result>; + async fn add_history(&self, history: &[NewHistory]) -> Result<()>; + + async fn oldest_history(&self, user: &User) -> Result; + + async fn calendar( + &self, + user: &User, + period: TimePeriod, + year: u64, + month: u64, + ) -> Result>; } #[derive(Clone)] @@ -106,10 +135,82 @@ impl Database for Postgres { Ok(res.0) } + async fn count_history_range( + &self, + user: &User, + start: chrono::NaiveDateTime, + end: chrono::NaiveDateTime, + ) -> Result { + let res: (i64,) = sqlx::query_as( + "select count(1) from history + where user_id = $1 + and timestamp >= $2::date + and timestamp < $3::date", + ) + .bind(user.id) + .bind(start) + .bind(end) + .fetch_one(&self.pool) + .await?; + + Ok(res.0) + } + + // Count the history for a given year + async fn count_history_year(&self, user: &User, year: i32) -> Result { + let start = chrono::Utc.ymd(year, 1, 1).and_hms_nano(0, 0, 0, 0); + let end = start + RelativeDuration::years(1); + + let res = self + .count_history_range(user, start.naive_utc(), end.naive_utc()) + .await?; + Ok(res) + } + + // Count the history for a given month + async fn count_history_month(&self, user: &User, month: chrono::NaiveDate) -> Result { + let start = chrono::Utc + .ymd(month.year(), month.month(), 1) + .and_hms_nano(0, 0, 0, 0); + + // ofc... + let end = if month.month() < 12 { + chrono::Utc + .ymd(month.year(), month.month() + 1, 1) + .and_hms_nano(0, 0, 0, 0) + } else { + chrono::Utc + .ymd(month.year() + 1, 1, 1) + .and_hms_nano(0, 0, 0, 0) + }; + + debug!("start: {}, end: {}", start, end); + + let res = self + .count_history_range(user, start.naive_utc(), end.naive_utc()) + .await?; + Ok(res) + } + + // Count the history for a given day + async fn count_history_day(&self, user: &User, day: chrono::NaiveDate) -> Result { + let start = chrono::Utc + .ymd(day.year(), day.month(), day.day()) + .and_hms_nano(0, 0, 0, 0); + let end = chrono::Utc + .ymd(day.year(), day.month(), day.day() + 1) + .and_hms_nano(0, 0, 0, 0); + + let res = self + .count_history_range(user, start.naive_utc(), end.naive_utc()) + .await?; + Ok(res) + } + async fn list_history( &self, user: &User, - created_since: chrono::NaiveDateTime, + created_after: chrono::NaiveDateTime, since: chrono::NaiveDateTime, host: &str, ) -> Result> { @@ -124,7 +225,7 @@ impl Database for Postgres { ) .bind(user.id) .bind(host) - .bind(created_since) + .bind(created_after) .bind(since) .bind(HISTORY_PAGE_SIZE) .fetch_all(&self.pool) @@ -211,4 +312,106 @@ impl Database for Postgres { Err(eyre!("could not find session")) } } + + async fn oldest_history(&self, user: &User) -> Result { + let res = sqlx::query_as::<_, History>( + "select * from history + where user_id = $1 + order by timestamp asc + limit 1", + ) + .bind(user.id) + .fetch_one(&self.pool) + .await?; + + Ok(res) + } + + async fn calendar( + &self, + user: &User, + period: TimePeriod, + year: u64, + month: u64, + ) -> Result> { + // TODO: Support different timezones. Right now we assume UTC and + // everything is stored as such. But it _should_ be possible to + // interpret the stored date with a different TZ + + match period { + TimePeriod::YEAR => { + let mut ret = HashMap::new(); + // First we need to work out how far back to calculate. Get the + // oldest history item + let oldest = self.oldest_history(user).await?.timestamp.year(); + let current_year = chrono::Utc::now().year(); + + // All the years we need to get data for + // The upper bound is exclusive, so include current +1 + let years = oldest..current_year + 1; + + for year in years { + let count = self.count_history_year(user, year).await?; + + ret.insert( + year as u64, + TimePeriodInfo { + count: count as u64, + hash: "".to_string(), + }, + ); + } + + Ok(ret) + } + + TimePeriod::MONTH => { + let mut ret = HashMap::new(); + + for month in 1..13 { + let count = self + .count_history_month( + user, + chrono::Utc.ymd(year as i32, month, 1).naive_utc(), + ) + .await?; + + ret.insert( + month as u64, + TimePeriodInfo { + count: count as u64, + hash: "".to_string(), + }, + ); + } + + Ok(ret) + } + + TimePeriod::DAY => { + let mut ret = HashMap::new(); + + for day in 1..get_days_from_month(year as i32, month as u32) { + let count = self + .count_history_day( + user, + chrono::Utc + .ymd(year as i32, month as u32, day as u32) + .naive_utc(), + ) + .await?; + + ret.insert( + day as u64, + TimePeriodInfo { + count: count as u64, + hash: "".to_string(), + }, + ); + } + + Ok(ret) + } + } + } } diff --git a/atuin-server/src/handlers/history.rs b/atuin-server/src/handlers/history.rs index 546e5a2..fde7cf2 100644 --- a/atuin-server/src/handlers/history.rs +++ b/atuin-server/src/handlers/history.rs @@ -1,11 +1,14 @@ use axum::extract::Query; -use axum::{Extension, Json}; +use axum::{extract::Path, Extension, Json}; use http::StatusCode; +use std::collections::HashMap; use crate::database::{Database, Postgres}; use crate::models::{NewHistory, User}; use atuin_common::api::*; +use crate::calendar::{TimePeriod, TimePeriodInfo}; + pub async fn count( user: User, db: Extension, @@ -79,3 +82,46 @@ pub async fn add( Ok(()) } + +pub async fn calendar( + Path(focus): Path, + Query(params): Query>, + user: User, + db: Extension, +) -> Result>, ErrorResponseStatus<'static>> { + let focus = focus.as_str(); + + let year = params.get("year").unwrap_or(&0); + let month = params.get("month").unwrap_or(&1); + + let focus = match focus { + "year" => db + .calendar(&user, TimePeriod::YEAR, *year, *month) + .await + .map_err(|_| { + ErrorResponse::reply("failed to query calendar") + .with_status(StatusCode::INTERNAL_SERVER_ERROR) + }), + + "month" => db + .calendar(&user, TimePeriod::MONTH, *year, *month) + .await + .map_err(|_| { + ErrorResponse::reply("failed to query calendar") + .with_status(StatusCode::INTERNAL_SERVER_ERROR) + }), + + "day" => db + .calendar(&user, TimePeriod::DAY, *year, *month) + .await + .map_err(|_| { + ErrorResponse::reply("failed to query calendar") + .with_status(StatusCode::INTERNAL_SERVER_ERROR) + }), + + _ => Err(ErrorResponse::reply("invalid focus: use year/month/day") + .with_status(StatusCode::BAD_REQUEST)), + }?; + + Ok(Json(focus)) +} diff --git a/atuin-server/src/lib.rs b/atuin-server/src/lib.rs index ca0aa11..02b3fda 100644 --- a/atuin-server/src/lib.rs +++ b/atuin-server/src/lib.rs @@ -15,6 +15,7 @@ extern crate log; extern crate serde_derive; pub mod auth; +pub mod calendar; pub mod database; pub mod handlers; pub mod models; diff --git a/atuin-server/src/router.rs b/atuin-server/src/router.rs index 6ca4722..146cc99 100644 --- a/atuin-server/src/router.rs +++ b/atuin-server/src/router.rs @@ -54,12 +54,12 @@ where async fn teapot() -> impl IntoResponse { (http::StatusCode::IM_A_TEAPOT, "☕") } - pub fn router(postgres: Postgres, settings: Settings) -> Router { Router::new() .route("/", get(handlers::index)) .route("/sync/count", get(handlers::history::count)) .route("/sync/history", get(handlers::history::list)) + .route("/sync/calendar/:focus", get(handlers::history::calendar)) .route("/history", post(handlers::history::add)) .route("/user/:username", get(handlers::user::get)) .route("/register", post(handlers::user::register))