fix: use extractor for db connections

Author: ImShyMike

rustytime/src/handlers/admin/mod.rs

@@ -1,28 +1,23 @@
 use aide::NoApi;
 use axum::extract::Path;
 use axum::{
-    extract::State,
     http::StatusCode,
     response::{IntoResponse, Response},
 };
 
+use crate::db_query;
 use crate::models::user::User;
-use crate::state::AppState;
-use crate::utils::auth::AuthenticatedUser;
-use crate::{db_query, get_db_conn};
+use crate::utils::extractors::{AuthenticatedUser, DbConnection};
 
 pub async fn change_user_admin_level(
-    State(app_state): State<AppState>,
     Path((user_id, new_level)): Path<(i32, i16)>,
     NoApi(AuthenticatedUser(current_user)): NoApi<AuthenticatedUser>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
 ) -> Result<StatusCode, Response> {
-
     if !current_user.is_owner() {
         return Err((StatusCode::FORBIDDEN, "No permission").into_response());
     }
 
-    let mut conn = get_db_conn!(app_state);
-
     let Some(target_user) = db_query!(
         User::get_by_id(&mut conn, user_id),
         "Failed to fetch target user"

rustytime/src/handlers/api/user.rs

@@ -1,3 +1,4 @@
+use aide::NoApi;
 use axum::extract::ConnectInfo;
 use axum::extract::{Json, Path, State};
 use axum::http::StatusCode;
@@ -10,13 +11,13 @@ use std::net::IpAddr;
 use std::net::SocketAddr;
 
 use crate::db::connection::DbPool;
-use crate::get_db_conn;
 use crate::models::heartbeat::Heartbeat;
 use crate::models::heartbeat::*;
 use crate::models::project::get_or_create_project_id;
 use crate::schema::heartbeats;
 use crate::state::AppState;
 use crate::utils::auth::{get_user_id_from_api_key, get_valid_api_key};
+use crate::utils::extractors::DbConnection;
 use crate::utils::http::extract_client_ip_from_headers;
 use crate::utils::time::{TimeFormat, human_readable_duration};
 use std::collections::{HashMap, hash_map};
@@ -132,6 +133,7 @@ pub async fn create_heartbeats(
 pub async fn get_statusbar_today(
     State(app_state): State<AppState>,
     Path(id): Path<String>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
     headers: axum::http::HeaderMap,
     uri: axum::http::Uri,
 ) -> Result<Json<serde_json::Value>, Response> {
@@ -164,8 +166,6 @@ pub async fn get_statusbar_today(
         .unwrap()
         .and_utc();
 
-    let mut conn = get_db_conn!(app_state);
-
     match Heartbeat::get_user_duration_seconds(
         &mut conn,
         DurationInput {

rustytime/src/handlers/data/import.rs

@@ -1,10 +1,10 @@
 use aide::NoApi;
+use axum::Json;
 use axum::extract::Query;
 use axum::extract::State;
 use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::response::Response;
-use axum::Json;
 use schemars::JsonSchema;
 use serde::Deserialize;
 use serde::Serialize;
@@ -15,8 +15,8 @@ use crate::db_query;
 use crate::jobs::import::enqueue_import;
 use crate::models::import_job::{ImportJob, ImportJobStatus};
 use crate::state::AppState;
+use crate::utils::extractors::AuthenticatedUser;
 use crate::utils::session::SessionManager;
-use crate::utils::auth::AuthenticatedUser;
 
 #[derive(Deserialize, JsonSchema)]
 pub struct ImportQuery {
@@ -67,7 +67,6 @@ pub async fn import_heartbeats(
     cookies: NoApi<Cookies>,
     NoApi(AuthenticatedUser(current_user)): NoApi<AuthenticatedUser>,
 ) -> Result<Json<ImportStartResponse>, Response> {
-
     let Some(session_id) = SessionManager::get_session_from_cookies(&cookies) else {
         return Err((StatusCode::UNAUTHORIZED, "User session is invalid").into_response());
     };

rustytime/src/handlers/data/project_aliases.rs

@@ -1,12 +1,10 @@
+use crate::db_query;
 use crate::models::project_alias::{NewProjectAlias, ProjectAlias as ProjectAliasModel};
-use crate::state::AppState;
-use crate::utils::auth::AuthenticatedUser;
-use crate::{db_query, get_db_conn};
+use crate::utils::extractors::{AuthenticatedUser, DbConnection};
 use aide::NoApi;
 use axum::Json;
 use axum::extract::Path;
 use axum::{
-    extract::State,
     http::StatusCode,
     response::{IntoResponse, Response},
 };
@@ -33,13 +31,9 @@ pub struct ProjectAliasesResponse {
 
 /// Handler for the project aliases
 pub async fn project_aliases(
-    State(app_state): State<AppState>,
     NoApi(AuthenticatedUser(current_user)): NoApi<AuthenticatedUser>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
 ) -> Result<Json<ProjectAliasesResponse>, Response> {
-
-    // get database connection
-    let mut conn = get_db_conn!(app_state);
-
     // get all project aliases
     let alias_records = db_query!(
         ProjectAliasModel::list_user_project_aliases(&mut conn, current_user.id),
@@ -71,13 +65,10 @@ pub async fn project_aliases(
 }
 
 pub async fn add_project_alias(
-    State(app_state): State<AppState>,
     NoApi(AuthenticatedUser(current_user)): NoApi<AuthenticatedUser>,
     Path((id, alias_id)): Path<(i32, i32)>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
 ) -> Result<StatusCode, Response> {
-
-    let mut conn = get_db_conn!(app_state);
-
     let new_alias = NewProjectAlias {
         user_id: current_user.id,
         project_id: alias_id,
@@ -93,13 +84,10 @@ pub async fn add_project_alias(
 }
 
 pub async fn delete_project_alias(
-    State(app_state): State<AppState>,
     NoApi(AuthenticatedUser(current_user)): NoApi<AuthenticatedUser>,
     Path(id): Path<i32>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
 ) -> Result<StatusCode, Response> {
-
-    let mut conn = get_db_conn!(app_state);
-
     db_query!(
         ProjectAliasModel::delete_project_alias(&mut conn, current_user.id, id),
         "Failed to delete project alias"

rustytime/src/handlers/data/projects.rs

@@ -1,12 +1,10 @@
+use crate::db_query;
 use crate::models::project::Project as ProjectModel;
-use crate::state::AppState;
-use crate::utils::auth::AuthenticatedUser;
-use crate::{db_query, get_db_conn};
+use crate::utils::extractors::{AuthenticatedUser, DbConnection};
 use aide::NoApi;
 use axum::Json;
 use axum::extract::Path;
 use axum::{
-    extract::State,
     http::StatusCode,
     response::{IntoResponse, Response},
 };
@@ -31,13 +29,9 @@ pub struct RepoUrlRequest {
 
 /// Handler for the projects list
 pub async fn projects_list(
-    State(app_state): State<AppState>,
     NoApi(AuthenticatedUser(current_user)): NoApi<AuthenticatedUser>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
 ) -> Result<Json<ProjectsListResponse>, Response> {
-
-    // get database connection
-    let mut conn = get_db_conn!(app_state);
-
     // get all projects
     let projects = db_query!(
         ProjectModel::list_user_projects(&mut conn, current_user.id),
@@ -54,15 +48,11 @@ pub async fn projects_list(
 }
 
 pub async fn set_project_repo(
-    State(app_state): State<AppState>,
     NoApi(AuthenticatedUser(current_user)): NoApi<AuthenticatedUser>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
     Path(project_id): Path<i32>,
     repo_url: Json<RepoUrlRequest>,
 ) -> Result<Response, Response> {
-
-    // get database connection
-    let mut conn = get_db_conn!(app_state);
-
     // validate the url if provided
     if let Some(url) = &repo_url.repo_url {
         let parsed = url::Url::parse(url).map_err(|_| {

rustytime/src/handlers/github.rs

@@ -15,10 +15,10 @@ use serde::{Deserialize, Serialize};
 use std::env;
 use tower_cookies::Cookies;
 
-use crate::state::AppState;
+use crate::models::session::Session;
 use crate::utils::session::SessionManager;
-use crate::{get_db_conn, models::session::Session};
 use crate::{models::user::User, utils::env::is_production_env};
+use crate::{state::AppState, utils::extractors::DbConnection};
 use axum::Json;
 use tracing::error;
 
@@ -137,6 +137,7 @@ pub async fn login(
 pub async fn callback(
     State(app_state): State<AppState>,
     cookies: NoApi<Cookies>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
     Query(params): Query<AuthRequest>,
 ) -> Result<Redirect, Response> {
     let cookies = cookies.0;
@@ -202,9 +203,6 @@ pub async fn callback(
         }
     };
 
-    // get database connection
-    let mut conn = get_db_conn!(app_state);
-
     let (user, session) = match conn.transaction::<_, diesel::result::Error, _>(|conn| {
         // create or update user in database
         let user = User::create_or_update(
@@ -253,6 +251,7 @@ pub async fn callback(
 pub async fn verify_session(
     State(app_state): State<AppState>,
     Query(params): Query<serde_json::Value>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
     cookies: NoApi<Cookies>,
 ) -> Result<Json<VerifySessionResponse>, Response> {
     let cookies = cookies.0;
@@ -272,7 +271,6 @@ pub async fn verify_session(
     match SessionManager::validate_session(&app_state.db_pool, session_id).await {
         Ok(Some(session_data)) => {
             // get user details
-            let mut conn = get_db_conn!(app_state);
             let user = crate::schema::users::table
                 .find(session_data.user_id)
                 .first::<User>(&mut conn)
@@ -333,15 +331,12 @@ pub async fn verify_session(
 
 /// Handler to log out the user
 pub async fn logout(
-    State(app_state): State<AppState>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
     cookies: NoApi<Cookies>,
 ) -> Result<Response, Response> {
     let cookies = cookies.0;
     // get session from cookie
     if let Some(session_id) = SessionManager::get_session_from_cookies(&cookies) {
-        // delete session from database
-        let mut conn = get_db_conn!(app_state);
-
         diesel::delete(
             crate::schema::sessions::table.filter(crate::schema::sessions::id.eq(session_id)),
         )

rustytime/src/handlers/homepage.rs

@@ -1,7 +1,7 @@
-use crate::state::AppState;
-use axum::{extract::State, response::Redirect};
+use axum::response::Redirect;
+
 /// Handler for the homepage
-pub async fn home_page(State(_app_state): State<AppState>) -> Result<Redirect, Redirect> {
+pub async fn home_page() -> Result<Redirect, Redirect> {
     let frontend_url =
         std::env::var("FRONTEND_URL").unwrap_or_else(|_| "http://localhost:5173".to_string());
     Ok(Redirect::to(&frontend_url))

rustytime/src/handlers/page/admin.rs

@@ -12,14 +12,14 @@ use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use tower_cookies::Cookies;
 
+use crate::db_query;
 use crate::models::heartbeat::Heartbeat;
 use crate::models::session::Session;
 use crate::models::user::{PartialUser, User};
 use crate::state::AppState;
 use crate::utils::cache::CachedAdminStats;
+use crate::utils::extractors::{AuthenticatedUser, DbConnection};
 use crate::utils::session::{ImpersonationContext, SessionManager};
-use crate::utils::auth::AuthenticatedUser;
-use crate::{db_query, get_db_conn};
 
 #[derive(Deserialize, JsonSchema)]
 pub struct AdminQuery {
@@ -56,8 +56,8 @@ pub async fn admin_dashboard(
     State(app_state): State<AppState>,
     Query(query): Query<AdminQuery>,
     NoApi(AuthenticatedUser(current_user)): NoApi<AuthenticatedUser>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
 ) -> Result<Json<AdminDashboardResponse>, Response> {
-
     if !current_user.is_admin() {
         return Err((StatusCode::FORBIDDEN, "No permission").into_response());
     }
@@ -70,8 +70,6 @@ pub async fn admin_dashboard(
     let daily_activity = if let Some(cached) = cached {
         cached.daily_activity
     } else {
-        let mut conn = get_db_conn!(app_state);
-
         let raw_daily_activity = db_query!(
             Heartbeat::get_daily_activity_last_week(&mut conn),
             "Failed to fetch daily activity"
@@ -87,8 +85,6 @@ pub async fn admin_dashboard(
         raw_daily_activity
     };
 
-    let mut conn = get_db_conn!(app_state);
-
     let total_heartbeats = db_query!(Heartbeat::total_heartbeat_count_estimate(&mut conn));
     let total_users = db_query!(User::count_total_users(&mut conn, false));
     let heartbeats_last_hour = db_query!(Heartbeat::count_heartbeats_last_hour(&mut conn));
@@ -142,6 +138,7 @@ pub async fn impersonate_user(
     cookies: NoApi<Cookies>,
     impersonation: NoApi<Option<Extension<ImpersonationContext>>>,
     NoApi(AuthenticatedUser(session_user)): NoApi<AuthenticatedUser>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
 ) -> Result<StatusCode, Response> {
     let cookies = cookies.0;
     let impersonation = impersonation.0;
@@ -150,8 +147,6 @@ pub async fn impersonate_user(
         return Err((StatusCode::UNAUTHORIZED, "Session missing").into_response());
     };
 
-    let mut conn = get_db_conn!(app_state);
-
     let Some(session_data) = db_query!(
         SessionManager::validate_session(&app_state.db_pool, session_id).await,
         "Failed to validate session"

rustytime/src/handlers/page/dashboard.rs

@@ -1,10 +1,11 @@
+use crate::db_query;
+use crate::models::heartbeat::Heartbeat;
 use crate::models::heartbeat::{TimeRange, UsageStat};
 use crate::state::AppState;
 use crate::utils::cache::{CachedDashboardStats, DashboardCacheKey};
+use crate::utils::extractors::{AuthenticatedUser, DbConnection};
 use crate::utils::session::SessionManager;
 use crate::utils::time::{TimeFormat, human_readable_duration};
-use crate::utils::auth::AuthenticatedUser;
-use crate::{db_query, get_db_conn, models::heartbeat::Heartbeat};
 use aide::NoApi;
 use axum::{
     extract::{Query, State},
@@ -45,6 +46,7 @@ pub async fn dashboard(
     State(app_state): State<AppState>,
     cookies: NoApi<Cookies>,
     NoApi(AuthenticatedUser(user)): NoApi<AuthenticatedUser>,
+    NoApi(DbConnection(mut conn)): NoApi<DbConnection>,
     Query(query): Query<DashboardQuery>,
 ) -> Result<Json<DashboardResponse>, Response> {
     let cookies = cookies.0;
@@ -72,8 +74,6 @@ pub async fn dashboard(
     let (total_heartbeats, dashboard_stats) = if let Some(cached) = cached {
         (cached.heartbeat_count, cached.stats)
     } else {
-        let mut conn = get_db_conn!(app_state);
-
         let total_heartbeats = db_query!(
             Heartbeat::get_user_heartbeat_count_by_range(
                 &mut conn,