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

@@ -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"

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,
+}

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::*;

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"

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,
+}

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<Session>;
@@ -18,14 +25,36 @@ pub trait Database {
     async fn add_user(&self, user: &NewUser) -> Result<i64>;
 
     async fn count_history(&self, user: &User) -> Result<i64>;
+
+    async fn count_history_range(
+        &self,
+        user: &User,
+        start: chrono::NaiveDateTime,
+        end: chrono::NaiveDateTime,
+    ) -> Result<i64>;
+    async fn count_history_day(&self, user: &User, date: chrono::NaiveDate) -> Result<i64>;
+    async fn count_history_month(&self, user: &User, date: chrono::NaiveDate) -> Result<i64>;
+    async fn count_history_year(&self, user: &User, year: i32) -> Result<i64>;
+
     async fn list_history(
         &self,
         user: &User,
-        created_since: chrono::NaiveDateTime,
+        created_after: chrono::NaiveDateTime,
         since: chrono::NaiveDateTime,
         host: &str,
     ) -> Result<Vec<History>>;
+
     async fn add_history(&self, history: &[NewHistory]) -> Result<()>;
+
+    async fn oldest_history(&self, user: &User) -> Result<History>;
+
+    async fn calendar(
+        &self,
+        user: &User,
+        period: TimePeriod,
+        year: u64,
+        month: u64,
+    ) -> Result<HashMap<u64, TimePeriodInfo>>;
 }
 
 #[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<i64> {
+        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<i64> {
+        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<i64> {
+        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<i64> {
+        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<Vec<History>> {
@@ -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<History> {
+        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<HashMap<u64, TimePeriodInfo>> {
+        // 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)
+            }
+        }
+    }
 }

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<Postgres>,
@@ -79,3 +82,46 @@ pub async fn add(
 
     Ok(())
 }
+
+pub async fn calendar(
+    Path(focus): Path<String>,
+    Query(params): Query<HashMap<String, u64>>,
+    user: User,
+    db: Extension<Postgres>,
+) -> Result<Json<HashMap<u64, TimePeriodInfo>>, 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))
+}

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;

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))