Add post terminate hook handler

src/external/tus/hooks.rs

@@ -1,11 +1,11 @@
 use super::models::{ChangeFileInfo, PreCreateResponse};
-use crate::config::Config;
 use crate::external::tus::models::EventPayload;
 use crate::uploads::associations::db as AssociationDB;
 use crate::uploads::db as InputObjectDB;
+use crate::{config::Config, external::db};
 use anyhow::Result;
 use chrono::Utc;
-use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, Set};
+use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, ModelTrait, QueryFilter, Set};
 use uuid::Uuid;
 
 pub(super) async fn handle_pre_create(
@@ -58,9 +58,6 @@ pub(super) async fn handle_pre_create(
         submission_id: Set(submission_id),
         ..Default::default()
     };
-    println!("Creating s3 key");
-    let s3_key = format!("{}/{}", config.s3_prefix, object.last_insert_id);
-    println!("S3 key: {}", s3_key);
 
     match AssociationDB::Entity::insert(association_object)
         .exec(&db)
@@ -72,7 +69,9 @@ pub(super) async fn handle_pre_create(
 
     // Respond with a custom ID for tusd to upload to S3
     Ok(PreCreateResponse {
-        change_file_info: Some(ChangeFileInfo { id: s3_key }),
+        change_file_info: Some(ChangeFileInfo {
+            id: object.last_insert_id.to_string(),
+        }),
         status: "success".to_string(),
     })
 }
@@ -83,15 +82,18 @@ pub(super) async fn handle_post_create(
 ) -> Result<PreCreateResponse> {
     println!("Handling post-create");
     let upload_id = &payload.event.upload.id;
-    let object_id: Uuid = upload_id
-        .split('/')
-        .nth(1)
-        .unwrap()
+    // Split the upload_id on the + separator to get the object ID.
+    let object_id: Uuid = match upload_id
         .split('+')
         .next()
-        .unwrap()
-        .try_into()
-        .unwrap();
+        .and_then(|id_str| Uuid::parse_str(id_str).ok())
+    {
+        Some(id) => id,
+        None => {
+            println!("Invalid object ID in upload_id");
+            return Err(anyhow::anyhow!("Invalid object ID in upload_id"));
+        }
+    };
 
     let obj = match InputObjectDB::Entity::find()
         .filter(InputObjectDB::Column::Id.eq(object_id))
@@ -123,17 +125,18 @@ pub(super) async fn handle_post_receive(
 ) -> Result<PreCreateResponse> {
     println!("Handling post-receive");
     let upload_id = &payload.event.upload.id;
-    // Split the s3_prefix and then the UUID out of the upload_id to get the object ID.
-    // Then again with the + separator between UUID and TUSd upload ID
-    let object_id: Uuid = upload_id
-        .split('/')
-        .nth(1)
-        .unwrap()
+    // Split the upload_id on the + separator to get the object ID.
+    let object_id: Uuid = match upload_id
         .split('+')
         .next()
-        .unwrap()
-        .try_into()
-        .unwrap();
+        .and_then(|id_str| Uuid::parse_str(id_str).ok())
+    {
+        Some(id) => id,
+        None => {
+            println!("Invalid object ID in upload_id");
+            return Err(anyhow::anyhow!("Invalid object ID in upload_id"));
+        }
+    };
 
     let size_in_bytes = payload.event.upload.size;
     let offset = payload.event.upload.offset;
@@ -171,15 +174,18 @@ pub(super) async fn handle_pre_finish(
 ) -> Result<PreCreateResponse> {
     println!("Handling pre-finish");
     let upload_id = &payload.event.upload.id;
-    let object_id: Uuid = upload_id
-        .split('/')
-        .nth(1)
-        .unwrap()
+    // Split the upload_id on the + separator to get the object ID.
+    let object_id: Uuid = match upload_id
         .split('+')
         .next()
-        .unwrap()
-        .try_into()
-        .unwrap();
+        .and_then(|id_str| Uuid::parse_str(id_str).ok())
+    {
+        Some(id) => id,
+        None => {
+            println!("Invalid object ID in upload_id");
+            return Err(anyhow::anyhow!("Invalid object ID in upload_id"));
+        }
+    };
 
     let obj = match InputObjectDB::Entity::find()
         .filter(InputObjectDB::Column::Id.eq(object_id))
@@ -212,17 +218,18 @@ pub(super) async fn handle_post_finish(
     println!("Handling post-finish");
     let upload_id = &payload.event.upload.id;
 
-    // Unwrap the UUID from the upload_id. Something like this:
-    // labcaller/<uuid>+<tusd_upload_id> -> <uuid>
-    let object_id: Uuid = upload_id
-        .split('/')
-        .nth(1)
-        .unwrap()
+    // Split the upload_id on the + separator to get the object ID.
+    let object_id: Uuid = match upload_id
         .split('+')
         .next()
-        .unwrap()
-        .try_into()
-        .unwrap();
+        .and_then(|id_str| Uuid::parse_str(id_str).ok())
+    {
+        Some(id) => id,
+        None => {
+            println!("Invalid object ID in upload_id");
+            return Err(anyhow::anyhow!("Invalid object ID in upload_id"));
+        }
+    };
 
     let obj = match InputObjectDB::Entity::find()
         .filter(InputObjectDB::Column::Id.eq(object_id))
@@ -247,3 +254,50 @@ pub(super) async fn handle_post_finish(
         _ => Err(anyhow::anyhow!("Failed to update after upload completed")),
     }
 }
+
+pub(super) async fn handle_post_terminate(
+    db: DatabaseConnection,
+    payload: EventPayload,
+) -> Result<PreCreateResponse> {
+    // This hook is sent when the file should be cleaned up (del from db)
+    println!("Handling post-terminate");
+    let upload_id = &payload.event.upload.id;
+
+    // Split the upload_id on the + separator to get the object ID.
+    let object_id: Uuid = match upload_id
+        .split('+')
+        .next()
+        .and_then(|id_str| Uuid::parse_str(id_str).ok())
+    {
+        Some(id) => id,
+        None => {
+            println!("Invalid object ID in upload_id");
+            return Err(anyhow::anyhow!("Invalid object ID in upload_id"));
+        }
+    };
+
+    let obj = match InputObjectDB::Entity::find()
+        .filter(InputObjectDB::Column::Id.eq(object_id))
+        .one(&db)
+        .await
+    {
+        Ok(obj) => obj,
+        _ => return Err(anyhow::anyhow!("Failed to find object")),
+    };
+
+    // Delete all associations, then delete the object
+    AssociationDB::Entity::delete_many()
+        .filter(AssociationDB::Column::InputObjectId.eq(object_id))
+        .exec(&db)
+        .await
+        .unwrap();
+
+    let obj = obj.unwrap();
+    match obj.delete(&db).await {
+        Ok(_) => Ok(PreCreateResponse {
+            change_file_info: None,
+            status: "Upload terminated".to_string(),
+        }),
+        _ => Err(anyhow::anyhow!("Failed to delete object")),
+    }
+}

src/external/tus/models.rs

@@ -18,6 +18,8 @@ pub enum EventType {
     PreFinish,
     #[serde(rename = "post-finish")]
     PostFinish,
+    #[serde(rename = "post-terminate")]
+    PostTerminate,
     #[serde(other)]
     Unknown,
 }

src/external/tus/views.rs

@@ -1,6 +1,6 @@
 use super::hooks::{
-    handle_post_create, handle_post_finish, handle_post_receive, handle_pre_create,
-    handle_pre_finish,
+    handle_post_create, handle_post_finish, handle_post_receive, handle_post_terminate,
+    handle_pre_create, handle_pre_finish,
 };
 use crate::external::tus::models::{EventPayload, EventType};
 // use crate::objects::models::InputObject;
@@ -88,6 +88,16 @@ pub async fn handle_tus_hooks(
                 }),
             ),
         },
+        EventType::PostTerminate => match handle_post_terminate(db, payload).await {
+            Ok(response) => (StatusCode::CREATED, Json(response)),
+            Err(_) => (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                Json(PreCreateResponse {
+                    change_file_info: None,
+                    status: "error".to_string(),
+                }),
+            ),
+        },
         EventType::Unknown => (
             StatusCode::BAD_REQUEST,
             Json(PreCreateResponse {

src/uploads/views.rs

@@ -135,18 +135,6 @@ pub async fn delete_one(
         .await
         .unwrap();
 
-    // Delete from S3
-    let config = Config::from_env();
-    let s3_response = s3
-        .delete_object()
-        .bucket(config.s3_bucket)
-        .key(id.to_string())
-        .send()
-        .await
-        .unwrap();
-
-    println!("S3 response: {:?}", s3_response);
-
     match super::db::Entity::find_by_id(id).one(&db).await {
         Ok(Some(obj)) => {
             let obj: super::db::ActiveModel = obj.into();
@@ -156,6 +144,18 @@ pub async fn delete_one(
                     if res.rows_affected == 0 {
                         return StatusCode::NOT_FOUND;
                     }
+
+                    // Delete from S3
+                    let config = Config::from_env();
+                    let s3_response = s3
+                        .delete_object()
+                        .bucket(config.s3_bucket)
+                        .key(format!("{}/{}", config.s3_prefix, id))
+                        .send()
+                        .await
+                        .unwrap();
+
+                    println!("S3 response: {:?}", s3_response);
                     return StatusCode::NO_CONTENT;
                 }
                 Err(_) => {