cleanup: create file object in handler

devflowinc · Jan 8, 2025 · 00a64ac · 00a64ac
1 parent cff6b51
commit 00a64ac
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 19 deletions.
diff --git a/server/src/bin/file-worker.rs b/server/src/bin/file-worker.rs
@@ -14,9 +14,7 @@ use trieve_server::{
     operators::{
         clickhouse_operator::{ClickHouseEvent, EventQueue},
         dataset_operator::get_dataset_and_organization_from_dataset_id_query,
-        file_operator::{
-            create_file_chunks, create_file_query, get_aws_bucket, preprocess_file_to_chunks,
-        },
+        file_operator::{create_file_chunks, get_aws_bucket, preprocess_file_to_chunks},
         group_operator::{create_group_from_file_query, create_groups_query},
     },
 };
@@ -313,17 +311,6 @@ async fn upload_file(
     )
     .await?;
 
-    let file_size_mb = (file_data.len() as f64 / 1024.0 / 1024.0).round() as i64;
-
-    let created_file = create_file_query(
-        file_id,
-        file_size_mb,
-        file_worker_message.upload_file_data.clone(),
-        file_worker_message.dataset_id,
-        web_pool.clone(),
-    )
-    .await?;
-
     let group_id = if !file_worker_message
         .upload_file_data
         .pdf2md_options
@@ -365,7 +352,7 @@ async fn upload_file(
 
         let group_id = chunk_group.id;
 
-        create_group_from_file_query(group_id, created_file.id, web_pool.clone())
+        create_group_from_file_query(group_id, file_worker_message.file_id, web_pool.clone())
             .await
             .map_err(|e| {
                 log::error!("Could not create group from file {:?}", e);
@@ -568,7 +555,7 @@ async fn upload_file(
 
                 if !new_chunks.is_empty() {
                     create_file_chunks(
-                        created_file.id,
+                        file_worker_message.file_id,
                         file_worker_message.upload_file_data.clone(),
                         new_chunks.clone(),
                         dataset_org_plan_sub.clone(),
@@ -724,7 +711,7 @@ async fn upload_file(
         .collect::<Vec<_>>();
 
     create_file_chunks(
-        created_file.id,
+        file_worker_message.file_id,
         file_worker_message.upload_file_data,
         chunks,
         dataset_org_plan_sub,

diff --git a/server/src/handlers/file_handler.rs b/server/src/handlers/file_handler.rs
@@ -12,8 +12,8 @@ use crate::{
     operators::{
         crawl_operator::{process_crawl_doc, Document},
         file_operator::{
-            delete_file_query, get_aws_bucket, get_csvjsonl_aws_bucket, get_dataset_file_query,
-            get_file_query,
+            create_file_query, delete_file_query, get_aws_bucket, get_csvjsonl_aws_bucket,
+            get_dataset_file_query, get_file_query,
         },
         organization_operator::{get_file_size_sum_org, hash_function},
     },
@@ -189,6 +189,17 @@ pub async fn upload_file_handler(
             ServiceError::BadRequest("Could not upload file to S3".to_string())
         })?;
 
+    let file_size_mb = (decoded_file_data.len() as f64 / 1024.0 / 1024.0).round() as i64;
+
+    create_file_query(
+        file_id,
+        file_size_mb,
+        upload_file_data.clone(),
+        dataset_org_plan_sub.dataset.id,
+        pool.clone(),
+    )
+    .await?;
+
     let message = FileWorkerMessage {
         file_id,
         dataset_id: dataset_org_plan_sub.dataset.id,