From b4428c27c605ef89d7231096d15851ebfd9bfede Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Fri, 29 Sep 2023 17:49:38 +0100 Subject: [PATCH] support timezones in calendar (#1259) --- atuin-server-database/src/calendar.rs | 7 +- atuin-server-database/src/lib.rs | 168 +++++++++----------------- atuin-server-postgres/src/lib.rs | 9 +- atuin-server/src/handlers/history.rs | 78 +++++++----- 4 files changed, 109 insertions(+), 153 deletions(-) diff --git a/atuin-server-database/src/calendar.rs b/atuin-server-database/src/calendar.rs index 7c05dce..2229667 100644 --- a/atuin-server-database/src/calendar.rs +++ b/atuin-server-database/src/calendar.rs @@ -1,11 +1,12 @@ // Calendar data use serde::{Deserialize, Serialize}; +use time::Month; pub enum TimePeriod { - YEAR, - MONTH, - DAY, + Year, + Month { year: i32 }, + Day { year: i32, month: Month }, } #[derive(Debug, Serialize, Deserialize)] diff --git a/atuin-server-database/src/lib.rs b/atuin-server-database/src/lib.rs index 4ebd517..d529655 100644 --- a/atuin-server-database/src/lib.rs +++ b/atuin-server-database/src/lib.rs @@ -6,6 +6,7 @@ pub mod models; use std::{ collections::HashMap, fmt::{Debug, Display}, + ops::Range, }; use self::{ @@ -15,7 +16,7 @@ use self::{ use async_trait::async_trait; use atuin_common::record::{EncryptedData, HostId, Record, RecordId, RecordIndex}; use serde::{de::DeserializeOwned, Serialize}; -use time::{Date, Duration, Month, OffsetDateTime, PrimitiveDateTime, Time}; +use time::{Date, Duration, Month, OffsetDateTime, Time, UtcOffset}; use tracing::instrument; #[derive(Debug)] @@ -74,12 +75,8 @@ pub trait Database: Sized + Clone + Send + Sync + 'static { // Return the tail record ID for each store, so (HostID, Tag, TailRecordID) async fn tail_records(&self, user: &User) -> DbResult; - async fn count_history_range( - &self, - user: &User, - start: PrimitiveDateTime, - end: PrimitiveDateTime, - ) -> DbResult; + async fn count_history_range(&self, user: &User, range: Range) + -> DbResult; async fn list_history( &self, @@ -94,136 +91,81 @@ pub trait Database: Sized + Clone + Send + Sync + 'static { async fn oldest_history(&self, user: &User) -> DbResult; - /// Count the history for a given year - #[instrument(skip_all)] - async fn count_history_year(&self, user: &User, year: i32) -> DbResult { - let start = Date::from_calendar_date(year, time::Month::January, 1)?; - let end = Date::from_calendar_date(year + 1, time::Month::January, 1)?; - - let res = self - .count_history_range( - user, - start.with_time(Time::MIDNIGHT), - end.with_time(Time::MIDNIGHT), - ) - .await?; - Ok(res) - } - - /// Count the history for a given month - #[instrument(skip_all)] - async fn count_history_month(&self, user: &User, year: i32, month: Month) -> DbResult { - let start = Date::from_calendar_date(year, month, 1)?; - let days = time::util::days_in_year_month(year, month); - let end = start + Duration::days(days as i64); - - tracing::debug!("start: {}, end: {}", start, end); - - let res = self - .count_history_range( - user, - start.with_time(Time::MIDNIGHT), - end.with_time(Time::MIDNIGHT), - ) - .await?; - Ok(res) - } - - /// Count the history for a given day - #[instrument(skip_all)] - async fn count_history_day(&self, user: &User, day: Date) -> DbResult { - let end = day - .next_day() - .ok_or_else(|| DbError::Other(eyre::eyre!("no next day?")))?; - - let res = self - .count_history_range( - user, - day.with_time(Time::MIDNIGHT), - end.with_time(Time::MIDNIGHT), - ) - .await?; - Ok(res) - } - #[instrument(skip_all)] async fn calendar( &self, user: &User, period: TimePeriod, - year: u64, - month: Month, + tz: UtcOffset, ) -> DbResult> { - // 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(); + let mut ret = HashMap::new(); + let iter: Box)>> + Send> = match period { + TimePeriod::Year => { // 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 = OffsetDateTime::now_utc().year(); + let oldest = self + .oldest_history(user) + .await? + .timestamp + .to_offset(tz) + .year(); + let current_year = OffsetDateTime::now_utc().to_offset(tz).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?; + Box::new(years.map(|year| { + let start = Date::from_calendar_date(year, time::Month::January, 1)?; + let end = Date::from_calendar_date(year + 1, time::Month::January, 1)?; - ret.insert( - year as u64, - TimePeriodInfo { - count: count as u64, - hash: "".to_string(), - }, - ); - } - - Ok(ret) + Ok((year as u64, start..end)) + })) } - TimePeriod::MONTH => { - let mut ret = HashMap::new(); - + TimePeriod::Month { year } => { let months = std::iter::successors(Some(Month::January), |m| Some(m.next())).take(12); - for month in months { - let count = self.count_history_month(user, year as i32, month).await?; - ret.insert( - month as u64, - TimePeriodInfo { - count: count as u64, - hash: "".to_string(), - }, - ); - } + Box::new(months.map(move |month| { + let start = Date::from_calendar_date(year, month, 1)?; + let days = time::util::days_in_year_month(year, month); + let end = start + Duration::days(days as i64); - Ok(ret) + Ok((month as u64, start..end)) + })) } - TimePeriod::DAY => { - let mut ret = HashMap::new(); + TimePeriod::Day { year, month } => { + let days = 1..time::util::days_in_year_month(year, month); + Box::new(days.map(move |day| { + let start = Date::from_calendar_date(year, month, day)?; + let end = start + .next_day() + .ok_or_else(|| DbError::Other(eyre::eyre!("no next day?")))?; - for day in 1..time::util::days_in_year_month(year as i32, month) { - let count = self - .count_history_day(user, Date::from_calendar_date(year as i32, month, day)?) - .await?; - - ret.insert( - day as u64, - TimePeriodInfo { - count: count as u64, - hash: "".to_string(), - }, - ); - } - - Ok(ret) + Ok((day as u64, start..end)) + })) } + }; + + for x in iter { + let (index, range) = x?; + + let start = range.start.with_time(Time::MIDNIGHT).assume_offset(tz); + let end = range.end.with_time(Time::MIDNIGHT).assume_offset(tz); + + let count = self.count_history_range(user, start..end).await?; + + ret.insert( + index, + TimePeriodInfo { + count: count as u64, + hash: "".to_string(), + }, + ); } + + Ok(ret) } } diff --git a/atuin-server-postgres/src/lib.rs b/atuin-server-postgres/src/lib.rs index c71d03a..f22e6be 100644 --- a/atuin-server-postgres/src/lib.rs +++ b/atuin-server-postgres/src/lib.rs @@ -1,3 +1,5 @@ +use std::ops::Range; + use async_trait::async_trait; use atuin_common::record::{EncryptedData, HostId, Record, RecordId, RecordIndex}; use atuin_server_database::models::{History, NewHistory, NewSession, NewUser, Session, User}; @@ -176,8 +178,7 @@ impl Database for Postgres { async fn count_history_range( &self, user: &User, - start: PrimitiveDateTime, - end: PrimitiveDateTime, + range: Range, ) -> DbResult { let res: (i64,) = sqlx::query_as( "select count(1) from history @@ -186,8 +187,8 @@ impl Database for Postgres { and timestamp < $3::date", ) .bind(user.id) - .bind(start) - .bind(end) + .bind(into_utc(range.start)) + .bind(into_utc(range.end)) .fetch_one(&self.pool) .await .map_err(fix_error)?; diff --git a/atuin-server/src/handlers/history.rs b/atuin-server/src/handlers/history.rs index 7d6b273..508eed6 100644 --- a/atuin-server/src/handlers/history.rs +++ b/atuin-server/src/handlers/history.rs @@ -6,7 +6,7 @@ use axum::{ Json, }; use http::StatusCode; -use time::Month; +use time::{Month, UtcOffset}; use tracing::{debug, error, instrument}; use super::{ErrorResponse, ErrorResponseStatus, RespExt}; @@ -166,53 +166,65 @@ pub async fn add( Ok(()) } +#[derive(serde::Deserialize, Debug)] +pub struct CalendarQuery { + #[serde(default = "serde_calendar::zero")] + year: i32, + #[serde(default = "serde_calendar::one")] + month: u8, + + #[serde(default = "serde_calendar::utc")] + tz: UtcOffset, +} + +mod serde_calendar { + use time::UtcOffset; + + pub fn zero() -> i32 { + 0 + } + + pub fn one() -> u8 { + 1 + } + + pub fn utc() -> UtcOffset { + UtcOffset::UTC + } +} + #[instrument(skip_all, fields(user.id = user.id))] pub async fn calendar( Path(focus): Path, - Query(params): Query>, + Query(params): Query, UserAuth(user): UserAuth, state: State>, ) -> 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 month = Month::try_from(*month as u8).map_err(|e| ErrorResponseStatus { + let year = params.year; + let month = Month::try_from(params.month).map_err(|e| ErrorResponseStatus { error: ErrorResponse { reason: e.to_string().into(), }, status: http::StatusCode::BAD_REQUEST, })?; + let period = match focus { + "year" => TimePeriod::Year, + "month" => TimePeriod::Month { year }, + "day" => TimePeriod::Day { year, month }, + _ => { + return Err(ErrorResponse::reply("invalid focus: use year/month/day") + .with_status(StatusCode::BAD_REQUEST)) + } + }; + let db = &state.0.database; - 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)), - }?; + let focus = db.calendar(&user, period, params.tz).await.map_err(|_| { + ErrorResponse::reply("failed to query calendar") + .with_status(StatusCode::INTERNAL_SERVER_ERROR) + })?; Ok(Json(focus)) }