From 4dd91624150d332b14b2bddff5f78e8c53c421f8 Mon Sep 17 00:00:00 2001 From: aaldebs99 Date: Fri, 20 Jun 2025 01:32:32 +0000 Subject: [PATCH] fix(frontend): label writing and fetching logic --- frontend/src/components/Upload/UploadZone.tsx | 11 - frontend/src/pages/DocumentDetailsPage.tsx | 147 ++++++++++ src/db/documents.rs | 67 +++++ src/models.rs | 4 + src/routes/documents.rs | 270 +++++++++++++----- 5 files changed, 413 insertions(+), 86 deletions(-) diff --git a/frontend/src/components/Upload/UploadZone.tsx b/frontend/src/components/Upload/UploadZone.tsx index f74c37a..726d786 100644 --- a/frontend/src/components/Upload/UploadZone.tsx +++ b/frontend/src/components/Upload/UploadZone.tsx @@ -180,17 +180,6 @@ const UploadZone: React.FC = ({ onUploadComplete }) => { : f )); - // Assign labels to the uploaded document if any are selected - if (selectedLabels.length > 0) { - try { - const labelIds = selectedLabels.map(label => label.id); - await api.put(`/labels/documents/${response.data.id}`, { label_ids: labelIds }); - } catch (error) { - console.warn('Failed to assign labels to document:', error); - // Don't fail the upload if label assignment fails - } - } - if (onUploadComplete) { onUploadComplete(response.data); } diff --git a/frontend/src/pages/DocumentDetailsPage.tsx b/frontend/src/pages/DocumentDetailsPage.tsx index 2dabd5a..1a2055f 100644 --- a/frontend/src/pages/DocumentDetailsPage.tsx +++ b/frontend/src/pages/DocumentDetailsPage.tsx @@ -30,6 +30,7 @@ import { CalendarToday as DateIcon, Storage as SizeIcon, Tag as TagIcon, + Label as LabelIcon, Visibility as ViewIcon, Search as SearchIcon, Edit as EditIcon, @@ -37,6 +38,9 @@ import { } from '@mui/icons-material'; import { documentService, OcrResponse } from '../services/api'; import DocumentViewer from '../components/DocumentViewer'; +import LabelSelector from '../components/Labels/LabelSelector'; +import { type LabelData } from '../components/Labels/Label'; +import api from '../services/api'; interface Document { id: string; @@ -64,6 +68,10 @@ const DocumentDetailsPage: React.FC = () => { const [processedImageUrl, setProcessedImageUrl] = useState(null); const [processedImageLoading, setProcessedImageLoading] = useState(false); const [thumbnailUrl, setThumbnailUrl] = useState(null); + const [documentLabels, setDocumentLabels] = useState([]); + const [availableLabels, setAvailableLabels] = useState([]); + const [showLabelDialog, setShowLabelDialog] = useState(false); + const [labelsLoading, setLabelsLoading] = useState(false); useEffect(() => { if (id) { @@ -80,9 +88,14 @@ const DocumentDetailsPage: React.FC = () => { useEffect(() => { if (document) { loadThumbnail(); + fetchDocumentLabels(); } }, [document]); + useEffect(() => { + fetchAvailableLabels(); + }, []); + const fetchDocumentDetails = async (): Promise => { if (!id) { setError('No document ID provided'); @@ -204,6 +217,58 @@ const DocumentDetailsPage: React.FC = () => { }); }; + const fetchDocumentLabels = async (): Promise => { + if (!id) return; + + try { + const response = await api.get(`/labels/documents/${id}`); + if (response.status === 200 && Array.isArray(response.data)) { + setDocumentLabels(response.data); + } + } catch (error) { + console.error('Failed to fetch document labels:', error); + } + }; + + const fetchAvailableLabels = async (): Promise => { + try { + setLabelsLoading(true); + const response = await api.get('/labels?include_counts=false'); + if (response.status === 200 && Array.isArray(response.data)) { + setAvailableLabels(response.data); + } + } catch (error) { + console.error('Failed to fetch available labels:', error); + } finally { + setLabelsLoading(false); + } + }; + + const handleCreateLabel = async (labelData: Omit) => { + try { + const response = await api.post('/labels', labelData); + const newLabel = response.data; + setAvailableLabels(prev => [...prev, newLabel]); + return newLabel; + } catch (error) { + console.error('Failed to create label:', error); + throw error; + } + }; + + const handleSaveLabels = async (selectedLabels: LabelData[]): Promise => { + if (!id) return; + + try { + const labelIds = selectedLabels.map(label => label.id); + await api.put(`/labels/documents/${id}`, { label_ids: labelIds }); + setDocumentLabels(selectedLabels); + setShowLabelDialog(false); + } catch (error) { + console.error('Failed to save labels:', error); + } + }; + if (loading) { return ( @@ -435,6 +500,48 @@ const DocumentDetailsPage: React.FC = () => { )} + + {/* Labels Section */} + + + + + + + Labels + + + + + {documentLabels.length > 0 ? ( + + {documentLabels.map((label) => ( + + ))} + + ) : ( + + No labels assigned to this document + + )} + + @@ -754,6 +861,46 @@ const DocumentDetailsPage: React.FC = () => { + + {/* Label Edit Dialog */} + setShowLabelDialog(false)} + maxWidth="md" + fullWidth + > + + Edit Document Labels + + + + + Select labels to assign to this document + + + + + + + + + ); }; diff --git a/src/db/documents.rs b/src/db/documents.rs index 78ec47d..a837e80 100644 --- a/src/db/documents.rs +++ b/src/db/documents.rs @@ -3,6 +3,7 @@ use sqlx::{Row, QueryBuilder}; use uuid::Uuid; use crate::models::{Document, SearchRequest, SearchMode, SearchSnippet, HighlightRange, EnhancedDocumentResponse}; +use crate::routes::labels::Label; use super::Database; impl Database { @@ -1292,4 +1293,70 @@ impl Database { Ok((duplicates, total)) } + + pub async fn get_document_labels(&self, document_id: Uuid) -> Result> { + let labels = sqlx::query_as::<_, Label>( + r#" + SELECT + l.id, l.user_id, l.name, l.description, l.color, + l.background_color, l.icon, l.is_system, l.created_at, l.updated_at, + 0::bigint as document_count, 0::bigint as source_count + FROM labels l + INNER JOIN document_labels dl ON l.id = dl.label_id + WHERE dl.document_id = $1 + ORDER BY l.name + "# + ) + .bind(document_id) + .fetch_all(&self.pool) + .await?; + + Ok(labels) + } + + pub async fn get_labels_for_documents(&self, document_ids: &[Uuid]) -> Result>> { + if document_ids.is_empty() { + return Ok(std::collections::HashMap::new()); + } + + let rows = sqlx::query( + r#" + SELECT + dl.document_id, + l.id, l.user_id, l.name, l.description, l.color, + l.background_color, l.icon, l.is_system, l.created_at, l.updated_at + FROM labels l + INNER JOIN document_labels dl ON l.id = dl.label_id + WHERE dl.document_id = ANY($1) + ORDER BY dl.document_id, l.name + "# + ) + .bind(document_ids) + .fetch_all(&self.pool) + .await?; + + let mut labels_map: std::collections::HashMap> = std::collections::HashMap::new(); + + for row in rows { + let document_id: Uuid = row.get("document_id"); + let label = Label { + id: row.get("id"), + user_id: row.get("user_id"), + name: row.get("name"), + description: row.get("description"), + color: row.get("color"), + background_color: row.get("background_color"), + icon: row.get("icon"), + is_system: row.get("is_system"), + created_at: row.get("created_at"), + updated_at: row.get("updated_at"), + document_count: 0, + source_count: 0, + }; + + labels_map.entry(document_id).or_insert_with(Vec::new).push(label); + } + + Ok(labels_map) + } } \ No newline at end of file diff --git a/src/models.rs b/src/models.rs index ead4a20..b0c8b7c 100644 --- a/src/models.rs +++ b/src/models.rs @@ -115,6 +115,9 @@ pub struct DocumentResponse { pub mime_type: String, /// Tags associated with the document pub tags: Vec, + /// Labels associated with the document + #[serde(default)] + pub labels: Vec, /// When the document was created pub created_at: DateTime, /// Whether OCR text has been extracted @@ -260,6 +263,7 @@ impl From for DocumentResponse { file_size: doc.file_size, mime_type: doc.mime_type, tags: doc.tags, + labels: Vec::new(), // Labels will be populated separately where needed created_at: doc.created_at, has_ocr_text: doc.ocr_text.is_some(), ocr_confidence: doc.ocr_confidence, diff --git a/src/routes/documents.rs b/src/routes/documents.rs index 316b1dd..fe5fe45 100644 --- a/src/routes/documents.rs +++ b/src/routes/documents.rs @@ -7,9 +7,11 @@ use axum::{ }; use serde::Deserialize; use std::sync::Arc; +use std::collections::HashMap; use utoipa::ToSchema; use sha2::{Sha256, Digest}; use sqlx::Row; +use axum::body::Bytes; use crate::{ auth::AuthUser, @@ -69,6 +71,13 @@ async fn get_document_by_id( .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? .ok_or(StatusCode::NOT_FOUND)?; + // Get labels for this document + let labels = state + .db + .get_document_labels(document_id) + .await + .unwrap_or_else(|_| Vec::new()); + // Convert to DocumentResponse let response = DocumentResponse { id: document.id, @@ -79,6 +88,7 @@ async fn get_document_by_id( created_at: document.created_at, has_ocr_text: document.ocr_text.is_some(), tags: document.tags, + labels, ocr_confidence: document.ocr_confidence, ocr_word_count: document.ocr_word_count, ocr_processing_time_ms: document.ocr_processing_time_ms, @@ -118,92 +128,190 @@ async fn upload_document( .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? .unwrap_or_else(|| crate::models::Settings::default()); + let mut label_ids: Option> = None; + let mut file_data: Option<(String, Bytes)> = None; + + // First pass: collect all multipart fields while let Some(field) = multipart.next_field().await.map_err(|_| StatusCode::BAD_REQUEST)? { let name = field.name().unwrap_or("").to_string(); - if name == "file" { + tracing::info!("Processing multipart field: {}", name); + + if name == "label_ids" { + let label_ids_text = field.text().await.map_err(|_| StatusCode::BAD_REQUEST)?; + tracing::info!("Received label_ids field: {}", label_ids_text); + + match serde_json::from_str::>(&label_ids_text) { + Ok(ids) => { + tracing::info!("Successfully parsed {} label IDs: {:?}", ids.len(), ids); + label_ids = Some(ids); + }, + Err(e) => { + tracing::warn!("Failed to parse label_ids from upload: {} - Error: {}", label_ids_text, e); + } + } + } else if name == "file" { let filename = field .file_name() .ok_or(StatusCode::BAD_REQUEST)? .to_string(); - if !file_service.is_allowed_file_type(&filename, &settings.allowed_file_types) { - return Err(StatusCode::BAD_REQUEST); - } - let data = field.bytes().await.map_err(|_| StatusCode::BAD_REQUEST)?; - let file_size = data.len() as i64; - - // Check file size limit - let max_size_bytes = (settings.max_file_size_mb as i64) * 1024 * 1024; - if file_size > max_size_bytes { - return Err(StatusCode::PAYLOAD_TOO_LARGE); - } - - // Calculate file hash for deduplication - let file_hash = calculate_file_hash(&data); - - // Check if this exact file content already exists using efficient hash lookup - match state.db.get_document_by_user_and_hash(auth_user.user.id, &file_hash).await { - Ok(Some(existing_doc)) => { - // Return the existing document instead of creating a duplicate - return Ok(Json(existing_doc.into())); - } - Ok(None) => { - // No duplicate found, proceed with upload - } - Err(_) => { - // Continue even if duplicate check fails - } - } - - let mime_type = mime_guess::from_path(&filename) - .first_or_octet_stream() - .to_string(); - - let file_path = file_service - .save_file(&filename, &data) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - - let document = file_service.create_document( - &filename, - &filename, - &file_path, - file_size, - &mime_type, - auth_user.user.id, - Some(file_hash), - ); - - let saved_document = state - .db - .create_document(document) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - - let document_id = saved_document.id; - let enable_background_ocr = settings.enable_background_ocr; - - if enable_background_ocr { - // Use the shared queue service from AppState instead of creating a new one - // Calculate priority based on file size - let priority = match file_size { - 0..=1048576 => 10, // <= 1MB: highest priority - ..=5242880 => 8, // 1-5MB: high priority - ..=10485760 => 6, // 5-10MB: medium priority - ..=52428800 => 4, // 10-50MB: low priority - _ => 2, // > 50MB: lowest priority - }; - - state.queue_service.enqueue_document(document_id, priority, file_size).await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - } - - return Ok(Json(saved_document.into())); + let data_len = data.len(); + file_data = Some((filename.clone(), data)); + tracing::info!("Received file: {}, size: {} bytes", filename, data_len); } } + // Process the file after collecting all fields + if let Some((filename, data)) = file_data { + if !file_service.is_allowed_file_type(&filename, &settings.allowed_file_types) { + return Err(StatusCode::BAD_REQUEST); + } + + let file_size = data.len() as i64; + + // Check file size limit + let max_size_bytes = (settings.max_file_size_mb as i64) * 1024 * 1024; + if file_size > max_size_bytes { + return Err(StatusCode::PAYLOAD_TOO_LARGE); + } + + // Calculate file hash for deduplication + let file_hash = calculate_file_hash(&data); + + // Check if this exact file content already exists using efficient hash lookup + match state.db.get_document_by_user_and_hash(auth_user.user.id, &file_hash).await { + Ok(Some(existing_doc)) => { + // Return the existing document instead of creating a duplicate + let labels = state + .db + .get_document_labels(existing_doc.id) + .await + .unwrap_or_else(|_| Vec::new()); + + let mut response: DocumentResponse = existing_doc.into(); + response.labels = labels; + + return Ok(Json(response)); + } + Ok(None) => { + // No duplicate found, proceed with upload + } + Err(_) => { + // Continue even if duplicate check fails + } + } + + let mime_type = mime_guess::from_path(&filename) + .first_or_octet_stream() + .to_string(); + + let file_path = file_service + .save_file(&filename, &data) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + let document = file_service.create_document( + &filename, + &filename, + &file_path, + file_size, + &mime_type, + auth_user.user.id, + Some(file_hash), + ); + + let saved_document = state + .db + .create_document(document) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + let document_id = saved_document.id; + let enable_background_ocr = settings.enable_background_ocr; + + // Assign labels if provided + tracing::info!("Processing label assignment for document {}, label_ids: {:?}", document_id, label_ids); + + if let Some(ref label_ids_vec) = label_ids { + if !label_ids_vec.is_empty() { + tracing::info!("Attempting to assign {} labels to document {}", label_ids_vec.len(), document_id); + + // Verify all labels exist and are accessible to the user + let label_count = sqlx::query( + "SELECT COUNT(*) as count FROM labels WHERE id = ANY($1) AND (user_id = $2 OR is_system = TRUE)" + ) + .bind(label_ids_vec) + .bind(auth_user.user.id) + .fetch_one(state.db.get_pool()) + .await + .map_err(|e| { + tracing::error!("Failed to verify labels during upload: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let count: i64 = label_count.try_get("count").unwrap_or(0); + tracing::info!("Label verification: found {} valid labels out of {} requested", count, label_ids_vec.len()); + + if count as usize == label_ids_vec.len() { + // All labels are valid, assign them + for label_id in label_ids_vec { + match sqlx::query( + "INSERT INTO document_labels (document_id, label_id, assigned_by) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING" + ) + .bind(document_id) + .bind(label_id) + .bind(auth_user.user.id) + .execute(state.db.get_pool()) + .await + { + Ok(result) => { + tracing::info!("Successfully assigned label {} to document {}, rows affected: {}", label_id, document_id, result.rows_affected()); + }, + Err(e) => { + tracing::error!("Failed to assign label {} to document {}: {}", label_id, document_id, e); + } + } + } + } else { + tracing::warn!("Label verification failed: Some labels were not accessible to user {} during upload (found {}/{} labels)", auth_user.user.id, count, label_ids_vec.len()); + } + } else { + tracing::info!("No labels to assign (empty label_ids vector)"); + } + } else { + tracing::info!("No labels to assign (label_ids is None)"); + } + + if enable_background_ocr { + // Use the shared queue service from AppState instead of creating a new one + // Calculate priority based on file size + let priority = match file_size { + 0..=1048576 => 10, // <= 1MB: highest priority + ..=5242880 => 8, // 1-5MB: high priority + ..=10485760 => 6, // 5-10MB: medium priority + ..=52428800 => 4, // 10-50MB: low priority + _ => 2, // > 50MB: lowest priority + }; + + state.queue_service.enqueue_document(document_id, priority, file_size).await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + } + + // Get labels for this document (if any were assigned) + let labels = state + .db + .get_document_labels(document_id) + .await + .unwrap_or_else(|_| Vec::new()); + + let mut response: DocumentResponse = saved_document.into(); + response.labels = labels; + + return Ok(Json(response)); + } + Err(StatusCode::BAD_REQUEST) } @@ -258,7 +366,19 @@ async fn list_documents( ) ).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - let documents_response: Vec = documents.into_iter().map(|doc| doc.into()).collect(); + // Get labels for all documents efficiently + let document_ids: Vec = documents.iter().map(|doc| doc.id).collect(); + let labels_map = state + .db + .get_labels_for_documents(&document_ids) + .await + .unwrap_or_else(|_| std::collections::HashMap::new()); + + let documents_response: Vec = documents.into_iter().map(|doc| { + let mut response: DocumentResponse = doc.into(); + response.labels = labels_map.get(&response.id).cloned().unwrap_or_else(Vec::new); + response + }).collect(); let response = serde_json::json!({ "documents": documents_response,