Use the count cache (#312)

* Use the count cache

By default read from the count cache - if there is no value there, then
do a full COUNT. The cache will be filled when the user posts up some
more history

* clean up server db error handling

Co-authored-by: Conrad Ludgate <conrad.ludgate@truelayer.com>

2 files changed, 29 insertions, 9 deletions

diff --git a/atuin-server/src/handlers/history.rs b/atuin-server/src/handlers/history.rs
index fde7cf2d..4fa2a962 100644
--- a/atuin-server/src/handlers/history.rs
+++ b/atuin-server/src/handlers/history.rs
@@ -13,10 +13,17 @@ pub async fn count(
     user: User,
     db: Extension<Postgres>,
 ) -> Result<Json<CountResponse>, ErrorResponseStatus<'static>> {
-    match db.count_history(&user).await {
+    match db.count_history_cached(&user).await {
+        // By default read out the cached value
         Ok(count) => Ok(Json(CountResponse { count })),
-        Err(_) => Err(ErrorResponse::reply("failed to query history count")
-            .with_status(StatusCode::INTERNAL_SERVER_ERROR)),
+
+        // If that fails, fallback on a full COUNT. Cache is built on a POST
+        // only
+        Err(_) => match db.count_history(&user).await {
+            Ok(count) => Ok(Json(CountResponse { count })),
+            Err(_) => Err(ErrorResponse::reply("failed to query history count")
+                .with_status(StatusCode::INTERNAL_SERVER_ERROR)),
+        },
     }
 }
 
diff --git a/atuin-server/src/handlers/user.rs b/atuin-server/src/handlers/user.rs
index 1bcfce2f..42e4aa33 100644
--- a/atuin-server/src/handlers/user.rs
+++ b/atuin-server/src/handlers/user.rs
@@ -32,10 +32,15 @@ pub async fn get(
 ) -> Result<Json<UserResponse>, ErrorResponseStatus<'static>> {
     let user = match db.get_user(username.as_ref()).await {
         Ok(user) => user,
-        Err(e) => {
-            debug!("user not found: {}", e);
+        Err(sqlx::Error::RowNotFound) => {
+            debug!("user not found: {}", username);
             return Err(ErrorResponse::reply("user not found").with_status(StatusCode::NOT_FOUND));
         }
+        Err(err) => {
+            error!("database error: {}", err);
+            return Err(ErrorResponse::reply("database error")
+                .with_status(StatusCode::INTERNAL_SERVER_ERROR));
+        }
     };
 
     Ok(Json(UserResponse {
@@ -96,20 +101,28 @@ pub async fn login(
 ) -> Result<Json<LoginResponse>, ErrorResponseStatus<'static>> {
     let user = match db.get_user(login.username.borrow()).await {
         Ok(u) => u,
+        Err(sqlx::Error::RowNotFound) => {
+            return Err(ErrorResponse::reply("user not found").with_status(StatusCode::NOT_FOUND));
+        }
         Err(e) => {
             error!("failed to get user {}: {}", login.username.clone(), e);
 
-            return Err(ErrorResponse::reply("user not found").with_status(StatusCode::NOT_FOUND));
+            return Err(ErrorResponse::reply("database error")
+                .with_status(StatusCode::INTERNAL_SERVER_ERROR));
         }
     };
 
     let session = match db.get_user_session(&user).await {
         Ok(u) => u,
-        Err(e) => {
-            error!("failed to get session for {}: {}", login.username, e);
-
+        Err(sqlx::Error::RowNotFound) => {
+            debug!("user session not found for user id={}", user.id);
             return Err(ErrorResponse::reply("user not found").with_status(StatusCode::NOT_FOUND));
         }
+        Err(err) => {
+            error!("database error for user {}: {}", login.username, err);
+            return Err(ErrorResponse::reply("database error")
+                .with_status(StatusCode::INTERNAL_SERVER_ERROR));
+        }
     };
 
     let verified = verify_str(user.password.as_str(), login.password.borrow());