diff --git a/backend/src/auth/mod.rs b/backend/src/auth/mod.rs index c5b6270..8cb2140 100644 --- a/backend/src/auth/mod.rs +++ b/backend/src/auth/mod.rs @@ -5,6 +5,8 @@ pub mod jwt; pub mod passkeys; pub mod password; +use std::sync::{Arc, Mutex}; + use axum::{async_trait, extract::FromRequestParts, http::request::Parts}; use axum_extra::headers::{authorization::Bearer, Authorization}; use axum_extra::TypedHeader; @@ -21,6 +23,23 @@ use uuid::Uuid; use crate::auth::jwt::PrincipalKind; +#[derive(Clone)] +pub struct TenantConnectionHolder { + inner: Arc>>, +} + +impl TenantConnectionHolder { + pub fn new(conn: PgPooledConnection) -> Self { + Self { + inner: Arc::new(Mutex::new(Some(conn))), + } + } + + pub fn into_conn(self) -> Option { + self.inner.lock().ok()?.take() + } +} + #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct AuthenticatedUser { pub user_id: uuid::Uuid, @@ -78,6 +97,9 @@ impl FromRequestParts for AuthenticatedUser { }; parts.extensions.insert(user.clone()); + parts + .extensions + .insert(TenantConnectionHolder::new(tenant_conn)); Ok(user) } @@ -106,7 +128,13 @@ impl FromRequestParts for TenantScopedConn { ) -> Result { let user = AuthenticatedUser::from_request_parts(parts, state).await?; let tenant_id = user.tenant_id; - let conn = state.db_for_tenant(tenant_id)?; + let conn = if let Some(holder) = parts.extensions.remove::() { + holder + .into_conn() + .ok_or_else(|| AppError::internal("tenant connection unavailable"))? + } else { + state.db_for_tenant(tenant_id)? + }; Ok(Self { conn, diff --git a/backend/src/documents/asset.rs b/backend/src/documents/asset.rs index 546cbdb..0a8d390 100644 --- a/backend/src/documents/asset.rs +++ b/backend/src/documents/asset.rs @@ -139,28 +139,21 @@ pub fn to_asset_object_response( } } -pub fn delete_asset(state: &AppState, tenant_id: Uuid, asset_id: Uuid) -> AppResult<()> { - let mut conn = state.db_for_tenant(tenant_id)?; - +pub fn delete_asset( + conn: &mut PgPooledConnection, + tenant_id: Uuid, + asset_id: Uuid, +) -> AppResult<()> { diesel::delete( document_assets::table .filter(document_assets::id.eq(asset_id)) .filter(document_assets::tenant_id.eq(tenant_id)), ) - .execute(&mut conn)?; + .execute(conn)?; Ok(()) } -pub async fn load_asset_responses( - state: &AppState, - tenant_id: Uuid, - version_id: Uuid, -) -> AppResult> { - let mut conn = state.db_for_tenant(tenant_id)?; - load_asset_responses_with_conn(&mut conn, tenant_id, version_id) -} - pub fn load_asset_responses_with_conn( conn: &mut PgPooledConnection, tenant_id: Uuid, @@ -188,8 +181,7 @@ pub fn load_asset_responses_with_conn( } pub fn load_primary_assets( - state: &AppState, - tenant_id: Uuid, + conn: &mut PgPooledConnection, documents: &[Document], ) -> AppResult)>> { if documents.is_empty() { @@ -206,10 +198,9 @@ pub fn load_primary_assets( version_ids.sort(); version_ids.dedup(); - let mut conn = state.db_for_tenant(tenant_id)?; let versions: Vec = document_versions::table .filter(document_versions::id.eq_any(&version_ids)) - .load(&mut conn)?; + .load(conn)?; let mut version_map: HashMap = HashMap::new(); for version in versions { @@ -231,9 +222,7 @@ pub fn load_primary_assets( document_assets::all_columns, document_asset_objects::all_columns.nullable(), )) - .load(&mut conn)?; - - drop(conn); + .load(conn)?; let mut assets_by_version: HashMap> = HashMap::new(); for (asset, _object) in assets { diff --git a/backend/src/routes/auth.rs b/backend/src/routes/auth.rs index 600c66e..4473f22 100644 --- a/backend/src/routes/auth.rs +++ b/backend/src/routes/auth.rs @@ -16,7 +16,7 @@ use crate::{ AuthenticationChallengeResponse, PasskeyLoginFinishPayload, PasskeyLoginStartPayload, PasskeyRegistrationFinishPayload, PasskeySummary, RegistrationChallengeResponse, }, - AuthenticatedUser, + AuthenticatedUser, TenantScopedConn, }, error::{AppError, AppResult}, http::responders::JsonResponse, @@ -180,7 +180,7 @@ pub async fn select_tenant( )] pub async fn logout( State(state): State, - user: AuthenticatedUser, + TenantScopedConn { mut conn, user, .. }: TenantScopedConn, jar: Option>, ) -> AppResult<(HeaderMap, StatusCode)> { let refresh_cookie = jar.as_ref().and_then(|cookies| { @@ -188,7 +188,7 @@ pub async fn logout( .get(SESSION_COOKIE_NAME) .map(|value| value.to_owned()) }); - AuthService::new(&state).logout(user, refresh_cookie.as_deref()) + AuthService::new(&state).logout(&mut conn, &user, refresh_cookie.as_deref()) } #[utoipa::path( diff --git a/backend/src/routes/documents.rs b/backend/src/routes/documents.rs index 16ad92c..8912483 100644 --- a/backend/src/routes/documents.rs +++ b/backend/src/routes/documents.rs @@ -186,12 +186,16 @@ pub async fn get_document( )] pub async fn upload_document( State(state): State, - scoped_conn: TenantScopedConn, + TenantScopedConn { + mut conn, + tenant_id, + user_id, + .. + }: TenantScopedConn, mut multipart: Multipart, ) -> AppResult> { - let tenant_id = scoped_conn.tenant_id; - let user_id = scoped_conn.user_id; - drop(scoped_conn); + let tenant_id = tenant_id; + let user_id = user_id; let mut file_bytes: Option> = None; let mut original_name: Option = None; let mut content_type: Option = None; @@ -353,7 +357,10 @@ pub async fn upload_document( }; let service = DocumentsService::new(&state); - let outcome = match service.upload_document(tenant_id, user_id, request).await { + let outcome = match service + .upload_document(&mut conn, tenant_id, user_id, request) + .await + { Ok(outcome) => outcome, Err(err) => { error!(error = ?err, original_name = %original_name_for_log, "document upload failed"); @@ -470,7 +477,7 @@ pub async fn get_document_asset( Path(asset_id): Path, Query(query): Query, TenantScopedConn { - mut conn, + conn, tenant_id, .. }: TenantScopedConn, @@ -485,7 +492,7 @@ pub async fn get_document_asset( } let service = DocumentsService::new(&state); let detail = service - .get_document_asset(&mut conn, tenant_id, asset_id, start, limit) + .get_document_asset(conn, tenant_id, asset_id, start, limit) .await?; ok_json(detail) } diff --git a/backend/src/routes/folders.rs b/backend/src/routes/folders.rs index 2a374d2..5bbd6bc 100644 --- a/backend/src/routes/folders.rs +++ b/backend/src/routes/folders.rs @@ -150,7 +150,7 @@ pub async fn list_folder_contents( )?; let documents = if include_documents { - service.hydrate_documents(&mut conn, tenant_id, user_id, documents)? + service.hydrate_documents(&mut conn, user_id, documents)? } else { Vec::new() }; diff --git a/backend/src/routes/webdav/mod.rs b/backend/src/routes/webdav/mod.rs index 8bab1bf..a4889ba 100644 --- a/backend/src/routes/webdav/mod.rs +++ b/backend/src/routes/webdav/mod.rs @@ -23,18 +23,18 @@ use crate::schema::{ document_versions::dsl as document_versions_dsl, documents::dsl as documents_dsl, folders::dsl as folders_dsl, user_memberships::dsl as memberships_dsl, users::dsl as users_dsl, }; -use crate::state::AppState; +use crate::state::{AppState, PgPooledConnection}; use crate::tenants::{apply_tenant_guc, apply_user_guc, clear_user_guc}; use crate::utils::{error::StorageResultExt, http::inline_content_disposition, time::to_http_date}; const REALM: &str = "Papercrate WebDAV"; const DOWNLOAD_URL_TTL_SECONDS: u64 = 300; -#[derive(Clone, Debug)] struct WebDavContext { tenant_id: Uuid, _user_id: Uuid, _username: String, + conn: PgPooledConnection, } pub fn create_router() -> Router { @@ -72,7 +72,7 @@ async fn handle_propfind( path: &str, headers: HeaderMap, ) -> Result { - let context = match authenticate(state, &headers)? { + let mut context = match authenticate(state, &headers)? { Some(user) => user, None => return Ok(unauthorized_response()), }; @@ -87,17 +87,18 @@ async fn handle_propfind( let tenant_id = context.tenant_id; let resources = if segments.is_empty() { - let contents = fetch_folder_contents(state, tenant_id, None)?; + let contents = fetch_folder_contents(&mut context.conn, tenant_id, None)?; build_resources_for_folder(None, &[], &contents, depth) } else { - let resolution = match resolve_path(state, tenant_id, &segments)? { + let resolution = match resolve_path(&mut context.conn, tenant_id, &segments)? { Some(resolved) => resolved, None => return Ok(not_found_response()), }; match resolution { ResolvedPath::Folder { folder, chain } => { - let contents = fetch_folder_contents(state, tenant_id, Some(folder.id))?; + let contents = + fetch_folder_contents(&mut context.conn, tenant_id, Some(folder.id))?; build_resources_for_folder(Some(&folder), &chain, &contents, depth) } ResolvedPath::Document { @@ -128,7 +129,7 @@ async fn handle_get_or_head( headers: HeaderMap, method: Method, ) -> Result { - let context = match authenticate(state, &headers)? { + let mut context = match authenticate(state, &headers)? { Some(user) => user, None => return Ok(unauthorized_response()), }; @@ -139,7 +140,7 @@ async fn handle_get_or_head( return Ok(method_not_allowed()); } - let resolution = match resolve_path(state, tenant_id, &segments)? { + let resolution = match resolve_path(&mut context.conn, tenant_id, &segments)? { Some(resolved) => resolved, None => return Ok(not_found_response()), }; @@ -233,18 +234,16 @@ fn parse_segments(path: &str) -> AppResult> { } fn fetch_folder_contents( - state: &AppState, + conn: &mut PgPooledConnection, tenant_id: Uuid, folder_id: Option, ) -> AppResult { - let mut conn = state.db_for_tenant(tenant_id)?; - let folder = match folder_id { Some(id) => Some( folders_dsl::folders .filter(folders_dsl::tenant_id.eq(tenant_id)) .find(id) - .first::(&mut conn)?, + .first::(conn)?, ), None => None, }; @@ -254,12 +253,12 @@ fn fetch_folder_contents( .filter(folders_dsl::tenant_id.eq(tenant_id)) .filter(folders_dsl::parent_id.eq(Some(id))) .order(folders_dsl::name.asc()) - .load(&mut conn)?, + .load(conn)?, None => folders_dsl::folders .filter(folders_dsl::tenant_id.eq(tenant_id)) .filter(folders_dsl::parent_id.is_null()) .order(folders_dsl::name.asc()) - .load(&mut conn)?, + .load(conn)?, }; let mut docs_query = documents_dsl::documents @@ -274,7 +273,7 @@ fn fetch_folder_contents( let documents: Vec = docs_query .order(documents_dsl::created_at.desc()) - .load(&mut conn)?; + .load(conn)?; let version_ids: Vec = documents.iter().map(|doc| doc.current_version_id).collect(); let versions: Vec = if version_ids.is_empty() { @@ -282,7 +281,7 @@ fn fetch_folder_contents( } else { document_versions_dsl::document_versions .filter(document_versions_dsl::id.eq_any(&version_ids)) - .load(&mut conn)? + .load(conn)? }; let mut version_map = versions @@ -498,6 +497,7 @@ fn authenticate(state: &AppState, headers: &HeaderMap) -> Result AppResult> { - let mut conn = state.db_for_tenant(tenant_id)?; let mut parent_id: Option = None; let mut chain: Vec = Vec::new(); let mut current_folder: Option = None; @@ -701,7 +700,7 @@ fn resolve_path( for (index, segment) in segments.iter().enumerate() { let is_last = index == segments.len() - 1; - if let Some(folder) = find_folder_by_name(&mut conn, tenant_id, parent_id, segment)? { + if let Some(folder) = find_folder_by_name(conn, tenant_id, parent_id, segment)? { chain.push(folder.name.clone()); if is_last { return Ok(Some(ResolvedPath::Folder { folder, chain })); @@ -713,7 +712,7 @@ fn resolve_path( if is_last { if let Some((document, version)) = - find_document_by_filename(&mut conn, tenant_id, parent_id, segment)? + find_document_by_filename(conn, tenant_id, parent_id, segment)? { chain.push(document.filename.clone()); return Ok(Some(ResolvedPath::Document { @@ -725,7 +724,7 @@ fn resolve_path( } if let Ok(uuid) = Uuid::parse_str(segment) { - if let Some(folder) = find_folder_by_id(&mut conn, tenant_id, uuid)? { + if let Some(folder) = find_folder_by_id(conn, tenant_id, uuid)? { if folder.parent_id != parent_id { return Ok(None); } @@ -738,7 +737,7 @@ fn resolve_path( continue; } - if let Some((document, version)) = find_document_by_id(&mut conn, tenant_id, uuid)? { + if let Some((document, version)) = find_document_by_id(conn, tenant_id, uuid)? { if document.folder_id != parent_id { return Ok(None); } diff --git a/backend/src/services/auth.rs b/backend/src/services/auth.rs index 496d68e..c8f55d0 100644 --- a/backend/src/services/auth.rs +++ b/backend/src/services/auth.rs @@ -1,7 +1,7 @@ use axum::http::{header::SET_COOKIE, HeaderMap, HeaderValue, StatusCode}; use axum::response::{IntoResponse, Response}; use chrono::{Duration as ChronoDuration, TimeZone, Utc}; -use diesel::{pg::PgConnection, prelude::*, OptionalExtension}; +use diesel::{pg::PgConnection, prelude::*, Connection, OptionalExtension}; use rand::{rngs::OsRng, RngCore}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; @@ -33,8 +33,7 @@ use crate::schema::{ user_sessions::{self, dsl as session_dsl}, users::dsl, }; -use crate::services::ServiceContext; -use crate::state::AppState; +use crate::state::{AppState, PgPooledConnection}; use crate::tenants::{ apply_tenant_guc, apply_user_guc, apply_user_session_hash, clear_user_guc, clear_user_session_hash, @@ -122,15 +121,11 @@ pub enum LoginResponseVariants { pub struct AuthService<'a> { state: &'a AppState, - ctx: ServiceContext<'a>, } impl<'a> AuthService<'a> { pub fn new(state: &'a AppState) -> Self { - Self { - state, - ctx: ServiceContext::new(state), - } + Self { state } } pub fn login(&self, payload: LoginRequest) -> AppResult { @@ -154,7 +149,7 @@ impl<'a> AuthService<'a> { let token_value = magic_token.unwrap(); - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; let username_hint = payload.username.trim(); let preferred_tenant_id = payload.preferred_tenant_id; @@ -175,7 +170,7 @@ impl<'a> AuthService<'a> { return Err(AppError::bad_request("api_token must not be empty")); } - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; let token = find_active_token_by_secret(&mut conn, None, secret, None)? .ok_or_else(AppError::unauthorized)?; @@ -242,7 +237,7 @@ impl<'a> AuthService<'a> { ) -> AppResult> { let username = normalize_username(&payload.username)?; - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; let exists: bool = dsl::users .filter(dsl::username.eq(&username)) .first::(&mut conn) @@ -279,7 +274,7 @@ impl<'a> AuthService<'a> { .verify_signup_token(&payload.signup_token) .map_err(|_| AppError::unauthorized())?; - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; let exists: bool = dsl::users .filter(dsl::username.eq(&claims.username)) @@ -303,7 +298,7 @@ impl<'a> AuthService<'a> { )?; let state_clone = self.state.clone(); - let response = self.ctx.tx(&mut conn, |conn| { + let response = conn.transaction::(|conn| { insert_user(conn, claims.sub, &claims.username)?; let tenant = state_clone.tenants.create_tenant_with_conn( @@ -333,7 +328,7 @@ impl<'a> AuthService<'a> { pub fn refresh(&self, refresh_value: &str) -> AppResult { let hashed = hash_session_token(refresh_value); - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; let now = Utc::now(); let now_naive = now.naive_utc(); @@ -379,7 +374,7 @@ impl<'a> AuthService<'a> { .map_err(|_| AppError::unauthorized())?, }; - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; apply_user_guc(&mut conn, user_id)?; let membership_exists = memberships_dsl::user_memberships @@ -405,10 +400,10 @@ impl<'a> AuthService<'a> { pub fn logout( &self, - user: AuthenticatedUser, + conn: &mut PgPooledConnection, + user: &AuthenticatedUser, refresh_cookie: Option<&str>, ) -> AppResult<(HeaderMap, StatusCode)> { - let mut conn = self.state.db_for_tenant(user.tenant_id)?; let now = Utc::now().naive_utc(); let revoked = if let Some(value) = refresh_cookie { @@ -423,7 +418,7 @@ impl<'a> AuthService<'a> { session_dsl::revoked_at.eq(now), session_dsl::updated_at.eq(now), )) - .execute(&mut conn)? + .execute(conn)? } else { 0 }; @@ -439,7 +434,7 @@ impl<'a> AuthService<'a> { session_dsl::revoked_at.eq(now), session_dsl::updated_at.eq(now), )) - .execute(&mut conn)?; + .execute(conn)?; } let mut headers = HeaderMap::new(); @@ -451,7 +446,7 @@ impl<'a> AuthService<'a> { &self, user: AuthenticatedUser, ) -> AppResult> { - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; apply_user_guc(&mut conn, user.user_id)?; let tenant_ids: Vec = memberships_dsl::user_memberships @@ -463,11 +458,10 @@ impl<'a> AuthService<'a> { let mut tenants = Vec::with_capacity(tenant_ids.len()); for tenant_id in tenant_ids { - let mut tenant_conn = self.state.db_for_tenant(tenant_id)?; let name: String = tenant_dsl::tenants .find(tenant_id) .select(tenant_dsl::name) - .first(&mut tenant_conn) + .first(&mut conn) .map_err(AppError::from)?; tenants.push(TenantSnippet { id: tenant_id, @@ -488,7 +482,7 @@ impl<'a> AuthService<'a> { .as_ref() .ok_or_else(|| AppError::bad_request("passkey support is disabled"))?; - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; let current_user: User = dsl::users.find(user.user_id).first(&mut conn)?; let challenge = service.start_registration(&mut conn, ¤t_user)?; ok_json(challenge) @@ -505,7 +499,7 @@ impl<'a> AuthService<'a> { .as_ref() .ok_or_else(|| AppError::bad_request("passkey support is disabled"))?; - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; let current_user: User = dsl::users.find(user.user_id).first(&mut conn)?; let PasskeyRegistrationFinishPayload { @@ -537,7 +531,7 @@ impl<'a> AuthService<'a> { .as_ref() .ok_or_else(|| AppError::bad_request("passkey support is disabled"))?; - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; let user: User = dsl::users .filter(dsl::username.eq(&username)) .first(&mut conn)?; @@ -553,7 +547,7 @@ impl<'a> AuthService<'a> { .as_ref() .ok_or_else(|| AppError::bad_request("passkey support is disabled"))?; - let mut conn = self.ctx.db_conn()?; + let mut conn = self.state.db_unscoped()?; let (user, _passkey, auth_result) = service.finish_authentication(&mut conn, payload.challenge_id, payload.credential)?; @@ -599,11 +593,12 @@ impl<'a> AuthService<'a> { let mut tenants = Vec::with_capacity(tenant_ids.len()); for tenant_id in tenant_ids { - let mut tenant_conn = self.state.db_for_tenant(tenant_id)?; + apply_tenant_guc(conn, tenant_id)?; + clear_user_guc(conn)?; let name: String = tenant_dsl::tenants .find(tenant_id) .select(tenant_dsl::name) - .first(&mut tenant_conn) + .first(conn) .map_err(AppError::from)?; tenants.push(TenantSnippet { id: tenant_id, @@ -634,7 +629,7 @@ impl<'a> AuthService<'a> { let now = Utc::now(); let now_naive = now.naive_utc(); - self.ctx.tx(conn, |conn| { + conn.transaction::(|conn| { let magic = magic_dsl::magic_tokens .filter(magic_dsl::token_hash.eq(&token_hash)) .filter(magic_dsl::expires_at.gt(now_naive)) diff --git a/backend/src/services/capability_sets.rs b/backend/src/services/capability_sets.rs index 33b34d4..6c1a105 100644 --- a/backend/src/services/capability_sets.rs +++ b/backend/src/services/capability_sets.rs @@ -1,6 +1,6 @@ use axum::http::StatusCode; use chrono::Utc; -use diesel::{pg::PgConnection, prelude::*, OptionalExtension}; +use diesel::{pg::PgConnection, prelude::*, Connection, OptionalExtension}; use serde::{Deserialize, Serialize}; use uuid::Uuid; @@ -19,7 +19,6 @@ use crate::{ capability_sets::{self, dsl as cs_dsl}, user_memberships, }, - services::TransactionExt, utils::text::normalize_identifier, }; @@ -156,7 +155,7 @@ impl CapabilitySetService { return ok_json(to_response(set, capabilities)); } - let set = conn.app_transaction(|conn| { + let set = conn.transaction::(|conn| { let mut working = set.clone(); if let Some(slug) = &payload.slug { diff --git a/backend/src/services/correspondents.rs b/backend/src/services/correspondents.rs index 8df0efb..5d84d52 100644 --- a/backend/src/services/correspondents.rs +++ b/backend/src/services/correspondents.rs @@ -1,5 +1,5 @@ use chrono::Utc; -use diesel::{dsl::not, prelude::*}; +use diesel::{dsl::not, prelude::*, Connection}; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; use uuid::Uuid; @@ -10,7 +10,6 @@ use crate::documents::correspondents::{ use crate::error::{AppError, AppResult}; use crate::schema::{document_correspondents, documents}; use crate::services::helpers::load_active_document; -use crate::services::TransactionExt; use crate::state::{AppState, PgPooledConnection}; use crate::utils::db::validate_bulk_ids; @@ -80,7 +79,7 @@ impl<'a> CorrespondentsService<'a> { .collect(); let correspondent_ids = normalize_correspondent_ids(&raw_ids)?; - conn.app_transaction(|conn| { + conn.transaction::<_, AppError, _>(|conn| { let document = load_active_document(conn, tenant_id, document_id)?; let mut updated = false; @@ -151,7 +150,7 @@ impl<'a> CorrespondentsService<'a> { let action = payload.action; - conn.app_transaction(|conn| { + conn.transaction::<_, AppError, _>(|conn| { let docs: Vec<(Uuid, Option)> = documents::table .filter(documents::id.eq_any(&payload.document_ids)) .filter(documents::tenant_id.eq(tenant_id)) diff --git a/backend/src/services/documents.rs b/backend/src/services/documents.rs index ef1997c..005a934 100644 --- a/backend/src/services/documents.rs +++ b/backend/src/services/documents.rs @@ -9,6 +9,7 @@ use diesel::{ prelude::*, result::DatabaseErrorKind, sql_types::Text, + Connection, }; use serde::{Deserialize, Serialize}; use serde_json::{json, Map, Value}; @@ -20,10 +21,9 @@ use uuid::Uuid; use crate::documents::{ asset::{ build_download_path, derive_document_title, filename_with_retained_extension, - load_asset_responses, load_asset_responses_with_conn, load_primary_assets, - to_asset_detail_response, to_asset_object_response, to_version_response, - DocumentAssetDetailResponse, DocumentAssetResponse, DocumentVersionDetailResponse, - DocumentVersionResponse, + load_asset_responses_with_conn, load_primary_assets, to_asset_detail_response, + to_asset_object_response, to_version_response, DocumentAssetDetailResponse, + DocumentAssetResponse, DocumentVersionDetailResponse, DocumentVersionResponse, }, correspondents::{ insert_document_correspondents, normalize_correspondent_ids, DocumentCorrespondentResponse, @@ -49,7 +49,7 @@ use crate::schema::{ }; use crate::services::{ correspondents::CorrespondentAssignmentInput, folders::gather_descendant_folder_ids, - helpers::load_active_document, TransactionExt, + helpers::load_active_document, }; use crate::state::{AppState, PgPooledConnection}; use crate::utils::{ @@ -372,7 +372,7 @@ impl<'a> DocumentsService<'a> { .cloned() .unwrap_or_else(|| (Vec::new(), Vec::new())); - let assets = load_asset_responses(self.state, tenant_id, current_version.id).await?; + let assets = load_asset_responses_with_conn(conn, tenant_id, current_version.id)?; let current_version_data = Some((to_version_response(current_version), assets)); let response = @@ -602,7 +602,7 @@ impl<'a> DocumentsService<'a> { } let docs: Vec = docs_query.load(conn)?; - let mut responses = self.hydrate_documents(conn, tenant_id, user_id, docs)?; + let mut responses = self.hydrate_documents(conn, user_id, docs)?; if let Some(order) = quickwit_order { let order_map: HashMap = order @@ -619,7 +619,6 @@ impl<'a> DocumentsService<'a> { pub fn hydrate_documents( &self, conn: &mut PgPooledConnection, - tenant_id: Uuid, user_id: Uuid, docs: Vec, ) -> AppResult> { @@ -629,7 +628,7 @@ impl<'a> DocumentsService<'a> { let doc_ids: Vec = docs.iter().map(|doc| doc.id).collect(); let mut relations = load_tags_and_correspondents(conn, &doc_ids)?; - let primary_versions = load_primary_assets(self.state, tenant_id, &docs)?; + let primary_versions = load_primary_assets(conn, &docs)?; docs.into_iter() .map(|doc| { @@ -644,6 +643,7 @@ impl<'a> DocumentsService<'a> { pub async fn upload_document( &self, + conn: &mut PgPooledConnection, tenant_id: Uuid, user_id: Uuid, request: DocumentUploadRequest, @@ -661,10 +661,8 @@ impl<'a> DocumentsService<'a> { skip_if_existing, } = request; - let mut tenant_conn = self.state.db_for_tenant(tenant_id)?; - if let Some(folder) = folder_id { - ensure_folder_exists_on_conn(&mut tenant_conn, tenant_id, folder)?; + ensure_folder_exists_on_conn(conn, tenant_id, folder)?; } let doc_id = Uuid::new_v4(); @@ -685,7 +683,7 @@ impl<'a> DocumentsService<'a> { if let Some(reused) = self .try_reuse_existing_document( - &mut tenant_conn, + conn, tenant_id, user_id, &checksum_hex, @@ -719,7 +717,7 @@ impl<'a> DocumentsService<'a> { metadata }; - let (document, version) = match tenant_conn.transaction(|conn| { + let (document, version) = match conn.transaction(|conn| { let new_document = NewDocument { id: doc_id, filename: stored_filename.clone(), @@ -767,13 +765,7 @@ impl<'a> DocumentsService<'a> { }; let detail = { - assign_tags_to_document( - &mut tenant_conn, - tenant_id, - &document, - &tag_ids, - Some(user_id), - )?; + assign_tags_to_document(conn, tenant_id, &document, &tag_ids, Some(user_id))?; if !correspondents.is_empty() { let raw_ids: Vec = correspondents @@ -782,7 +774,7 @@ impl<'a> DocumentsService<'a> { .collect(); let correspondent_ids = normalize_correspondent_ids(&raw_ids)?; insert_document_correspondents( - &mut tenant_conn, + conn, tenant_id, document.id, user_id, @@ -790,8 +782,7 @@ impl<'a> DocumentsService<'a> { )?; } - let tags_and_correspondents = - load_tags_and_correspondents(&mut tenant_conn, &[doc_id])?; + let tags_and_correspondents = load_tags_and_correspondents(conn, &[doc_id])?; let (tags, correspondents) = tags_and_correspondents .get(&doc_id) .cloned() @@ -809,7 +800,7 @@ impl<'a> DocumentsService<'a> { }; if let Err(err) = enqueue_job( - &mut tenant_conn, + conn, tenant_id, JOB_ANALYZE_DOCUMENT, json!({ @@ -907,12 +898,12 @@ impl<'a> DocumentsService<'a> { let document = load_active_document(conn, tenant_id, document_id)?; let version_id = document.current_version_id; - load_asset_responses(self.state, tenant_id, version_id).await + Ok(load_asset_responses_with_conn(conn, tenant_id, version_id)?) } pub async fn get_document_asset( &self, - conn: &mut PgPooledConnection, + mut conn: PgPooledConnection, tenant_id: Uuid, asset_id: Uuid, start: i32, @@ -921,7 +912,7 @@ impl<'a> DocumentsService<'a> { let asset: DocumentAsset = match document_assets::table .find(asset_id) .filter(document_assets::tenant_id.eq(tenant_id)) - .first(conn) + .first(&mut conn) .optional()? { Some(asset) => asset, @@ -945,7 +936,9 @@ impl<'a> DocumentsService<'a> { .filter(document_asset_objects::ordinal.ge(start)) .filter(document_asset_objects::ordinal.le(end)) .order(document_asset_objects::ordinal.asc()) - .load(conn)?; + .load(&mut conn)?; + + drop(conn); let expires_at = Utc::now() .timestamp_millis() @@ -1052,7 +1045,7 @@ impl<'a> DocumentsService<'a> { tenant_id: Uuid, document_id: Uuid, ) -> AppResult<()> { - conn.app_transaction(|conn| { + conn.transaction::<_, AppError, _>(|conn| { let document = documents::table .find(document_id) .filter(documents::tenant_id.eq(tenant_id)) @@ -1154,7 +1147,7 @@ impl<'a> DocumentsService<'a> { let tags_and_correspondents = load_tags_and_correspondents(conn, &[document_id])?; let version_id = current_version.id; - let assets = load_asset_responses(self.state, tenant_id, version_id).await?; + let assets = load_asset_responses_with_conn(conn, tenant_id, version_id)?; let version_response = to_version_response(current_version); let (tags, correspondents) = tags_and_correspondents .get(&document_id) diff --git a/backend/src/services/folders.rs b/backend/src/services/folders.rs index 27228a0..efbaf37 100644 --- a/backend/src/services/folders.rs +++ b/backend/src/services/folders.rs @@ -5,6 +5,7 @@ use diesel::{ dsl::{exists, sql}, prelude::*, sql_types::Text, + Connection, }; use serde::{Deserialize, Serialize}; use utoipa::{IntoParams, ToSchema}; @@ -15,10 +16,7 @@ use crate::error::{AppError, AppResult}; use crate::http::responders::{IntoAppResult, RowsAffectedExt}; use crate::models::{Document, Folder, NewFolder}; use crate::schema::{documents, folders}; -use crate::services::{ - documents::{DocumentResponse, DocumentsService}, - TransactionExt, -}; +use crate::services::documents::{DocumentResponse, DocumentsService}; use crate::state::{AppState, PgPooledConnection}; use crate::utils::{json::deserialize_patch_field, text::normalize_identifier, time::to_iso}; @@ -122,7 +120,7 @@ impl<'a> FolderService<'a> { return Err(AppError::bad_request("segments must not be empty")); } - let folder = conn.app_transaction(|conn| { + let folder = conn.transaction::(|conn| { let mut current_parent = payload.parent_id; let mut last_folder: Option = None; @@ -398,7 +396,7 @@ impl<'a> FolderService<'a> { tenant_id: Uuid, folder_id: Uuid, ) -> AppResult<()> { - conn.app_transaction(|conn| { + conn.transaction::<_, AppError, _>(|conn| { folders::table .find(folder_id) .filter(folders::tenant_id.eq(tenant_id)) @@ -451,7 +449,7 @@ impl<'a> FolderService<'a> { folder_id: Uuid, payload: UpdateFolderRequest, ) -> AppResult<()> { - conn.app_transaction(|conn| { + conn.transaction::<_, AppError, _>(|conn| { let folder: Folder = folders::table .find(folder_id) .filter(folders::tenant_id.eq(tenant_id)) @@ -557,11 +555,10 @@ impl<'a> FolderService<'a> { pub fn hydrate_documents( &self, conn: &mut PgPooledConnection, - tenant_id: Uuid, user_id: Uuid, docs: Vec, ) -> AppResult> { - DocumentsService::new(self.state).hydrate_documents(conn, tenant_id, user_id, docs) + DocumentsService::new(self.state).hydrate_documents(conn, user_id, docs) } } diff --git a/backend/src/services/mod.rs b/backend/src/services/mod.rs index affb9d3..bf15cd1 100644 --- a/backend/src/services/mod.rs +++ b/backend/src/services/mod.rs @@ -1,62 +1,3 @@ -use diesel::{Connection, PgConnection}; - -use crate::error::{AppError, AppResult}; -use crate::state::{AppState, PgPooledConnection}; - -pub trait TransactionExt { - fn app_transaction(&mut self, f: F) -> AppResult - where - F: FnOnce(&mut Self) -> AppResult; -} - -impl TransactionExt for PgConnection { - fn app_transaction(&mut self, f: F) -> AppResult - where - F: FnOnce(&mut Self) -> AppResult, - { - self.transaction::(|conn| f(conn)) - } -} - -impl TransactionExt for PgPooledConnection { - fn app_transaction(&mut self, f: F) -> AppResult - where - F: FnOnce(&mut Self) -> AppResult, - { - self.transaction::(|conn| f(conn)) - } -} - -pub struct ServiceContext<'a> { - state: &'a AppState, -} - -impl<'a> ServiceContext<'a> { - pub fn new(state: &'a AppState) -> Self { - Self { state } - } - - pub fn state(&self) -> &'a AppState { - self.state - } - - pub fn db_conn(&self) -> AppResult { - self.state.db_unscoped() - } - - pub fn db_conn_for_tenant(&self, tenant_id: uuid::Uuid) -> AppResult { - self.state.db_for_tenant(tenant_id) - } - - pub fn tx(&self, conn: &mut C, f: F) -> AppResult - where - C: TransactionExt, - F: FnOnce(&mut C) -> AppResult, - { - conn.app_transaction(f) - } -} - pub mod auth; pub mod capability_sets; pub mod correspondents; diff --git a/backend/src/state.rs b/backend/src/state.rs index 0539baf..a4c4b01 100644 --- a/backend/src/state.rs +++ b/backend/src/state.rs @@ -74,6 +74,8 @@ impl AppState { pub fn db_for_tenant(&self, tenant_id: Uuid) -> AppResult { debug_assert!(!tenant_id.is_nil(), "nil tenant_id passed to db_for_tenant"); let mut conn = self.db_unscoped()?; + let conn_ptr = &*conn as *const _; + tracing::trace!(target = "db_pool", ?conn_ptr, tenant_id = %tenant_id, "apply tenant context"); apply_tenant_guc(&mut conn, tenant_id)?; clear_user_guc(&mut conn)?; Ok(conn) @@ -84,6 +86,8 @@ impl AppState { tracing::error!(error = ?err, "database pool error"); AppError::internal("database pool error") })?; + let conn_ptr = &*conn as *const _; + tracing::trace!(target = "db_pool", ?conn_ptr, "acquired connection"); clear_tenant_context(&mut conn)?; Ok(conn) } diff --git a/backend/src/utils/tracing.rs b/backend/src/utils/tracing.rs index b3d2f30..2884b46 100644 --- a/backend/src/utils/tracing.rs +++ b/backend/src/utils/tracing.rs @@ -8,7 +8,7 @@ pub fn init_tracing(default_level: &str) { EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(default_level)); tracing_subscriber::fmt() .with_env_filter(filter) - .with_target(false) + .with_target(true) .compact() .init(); } diff --git a/backend/src/workers/ocr.rs b/backend/src/workers/ocr.rs index fb08d8a..5aa493f 100644 --- a/backend/src/workers/ocr.rs +++ b/backend/src/workers/ocr.rs @@ -19,6 +19,7 @@ use uuid::Uuid; use crate::{ documents::asset::delete_asset, + error::AppResult, jobs::{enqueue_job, JOB_GENERATE_OCR_TEXT, JOB_INDEX_DOCUMENT_TEXT}, models::{ Document, DocumentAsset, DocumentAssetObject, DocumentVersion, NewDocumentAsset, @@ -151,8 +152,9 @@ impl JobHandler for GenerateOcrTextJob { let tenant_id = context.document.tenant_id; let asset_id = existing_asset.id; let state_clone = state.clone(); - match task::spawn_blocking(move || { - delete_asset(state_clone.as_ref(), tenant_id, asset_id) + match task::spawn_blocking(move || -> AppResult<()> { + let mut conn = state_clone.db_for_tenant(tenant_id)?; + delete_asset(&mut conn, tenant_id, asset_id) }) .await { diff --git a/backend/src/workers/thumbnails.rs b/backend/src/workers/thumbnails.rs index 621bd46..766d546 100644 --- a/backend/src/workers/thumbnails.rs +++ b/backend/src/workers/thumbnails.rs @@ -13,6 +13,7 @@ use uuid::Uuid; use crate::{ documents::asset::delete_asset, + error::AppResult, jobs::JOB_GENERATE_THUMBNAILS, models::{ Document, DocumentAsset, DocumentAssetObject, DocumentVersion, NewDocumentAsset, @@ -184,8 +185,9 @@ impl JobHandler for GenerateThumbnailsJob { let tenant_id = initial.document.tenant_id; let asset_id = existing_preview.id; let state_clone = state.clone(); - match task::spawn_blocking(move || { - delete_asset(state_clone.as_ref(), tenant_id, asset_id) + match task::spawn_blocking(move || -> AppResult<()> { + let mut conn = state_clone.db_for_tenant(tenant_id)?; + delete_asset(&mut conn, tenant_id, asset_id) }) .await { @@ -224,8 +226,9 @@ impl JobHandler for GenerateThumbnailsJob { let tenant_id = initial.document.tenant_id; let asset_id = existing_thumbnail.id; let state_clone = state.clone(); - match task::spawn_blocking(move || { - delete_asset(state_clone.as_ref(), tenant_id, asset_id) + match task::spawn_blocking(move || -> AppResult<()> { + let mut conn = state_clone.db_for_tenant(tenant_id)?; + delete_asset(&mut conn, tenant_id, asset_id) }) .await {