Skip to content

Commit

Permalink
support timezones in calendar (#1259)
Browse files Browse the repository at this point in the history
  • Loading branch information
conradludgate committed Sep 29, 2023
1 parent a195c38 commit b4428c2
Show file tree
Hide file tree
Showing 4 changed files with 110 additions and 154 deletions.
7 changes: 4 additions & 3 deletions atuin-server-database/src/calendar.rs
Original file line number Diff line number Diff line change
@@ -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)]
Expand Down
170 changes: 56 additions & 114 deletions atuin-server-database/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ pub mod models;
use std::{
collections::HashMap,
fmt::{Debug, Display},
ops::Range,
};

use self::{
Expand All @@ -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)]
Expand Down Expand Up @@ -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<RecordIndex>;

async fn count_history_range(
&self,
user: &User,
start: PrimitiveDateTime,
end: PrimitiveDateTime,
) -> DbResult<i64>;
async fn count_history_range(&self, user: &User, range: Range<OffsetDateTime>)
-> DbResult<i64>;

async fn list_history(
&self,
Expand All @@ -94,136 +91,81 @@ pub trait Database: Sized + Clone + Send + Sync + 'static {

async fn oldest_history(&self, user: &User) -> DbResult<History>;

/// Count the history for a given year
#[instrument(skip_all)]
async fn count_history_year(&self, user: &User, year: i32) -> DbResult<i64> {
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<i64> {
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<i64> {
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<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();
let mut ret = HashMap::new();
let iter: Box<dyn Iterator<Item = DbResult<(u64, Range<Date>)>> + 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?;

ret.insert(
year as u64,
TimePeriodInfo {
count: count as u64,
hash: "".to_string(),
},
);
}
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)?;

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(),
},
);
}

Ok(ret)
}

TimePeriod::DAY => {
let mut ret = HashMap::new();
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);

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?;
Ok((month as u64, start..end))
}))
}

ret.insert(
day as u64,
TimePeriodInfo {
count: count as u64,
hash: "".to_string(),
},
);
}
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?")))?;

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)
}
}
9 changes: 5 additions & 4 deletions atuin-server-postgres/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -176,8 +178,7 @@ impl Database for Postgres {
async fn count_history_range(
&self,
user: &User,
start: PrimitiveDateTime,
end: PrimitiveDateTime,
range: Range<OffsetDateTime>,
) -> DbResult<i64> {
let res: (i64,) = sqlx::query_as(
"select count(1) from history
Expand All @@ -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)?;
Expand Down
78 changes: 45 additions & 33 deletions atuin-server/src/handlers/history.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -166,53 +166,65 @@ pub async fn add<DB: Database>(
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<DB: Database>(
Path(focus): Path<String>,
Query(params): Query<HashMap<String, u64>>,
Query(params): Query<CalendarQuery>,
UserAuth(user): UserAuth,
state: State<AppState<DB>>,
) -> 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 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))
}

1 comment on commit b4428c2

@vercel
Copy link

@vercel vercel bot commented on b4428c2 Sep 29, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Successfully deployed to the following URLs:

atuin-docs – ./

atuin-docs.vercel.app
atuin-docs-atuin.vercel.app
atuin-docs-git-main-atuin.vercel.app

Please sign in to comment.