feature: add delete chunks flag to delete file route

Author: skeptrunedev

server/src/handlers/file_handler.rs

@@ -1,4 +1,7 @@
-use super::auth_handler::{AdminOnly, LoggedUser};
+use super::{
+    auth_handler::{AdminOnly, LoggedUser},
+    group_handler::DeleteGroupData,
+};
 use crate::{
     data::models::{
         DatasetAndOrgWithSubAndPlan, DatasetConfiguration, File, FileAndGroupId, FileWorkerMessage,
@@ -336,6 +339,7 @@ pub async fn get_dataset_files_handler(
     params(
         ("TR-Dataset" = String, Header, description = "The dataset id or tracking_id to use for the request. We assume you intend to use an id if the value is a valid uuid."),
         ("file_id" = uuid::Uuid, description = "The id of the file to delete"),
+        ("delete_chunks" = bool, Query, description = "Delete the chunks within the group"),
     ),
     security(
         ("ApiKey" = ["admin"]),
@@ -344,14 +348,17 @@ pub async fn get_dataset_files_handler(
 #[tracing::instrument(skip(pool))]
 pub async fn delete_file_handler(
     file_id: web::Path<uuid::Uuid>,
+    query: web::Query<DeleteGroupData>,
     pool: web::Data<Pool>,
     _user: AdminOnly,
     dataset_org_plan_sub: DatasetAndOrgWithSubAndPlan,
 ) -> Result<HttpResponse, actix_web::Error> {
     let dataset_config =
         DatasetConfiguration::from_json(dataset_org_plan_sub.dataset.server_configuration.clone());
+
     delete_file_query(
         file_id.into_inner(),
+        query.delete_chunks,
         dataset_org_plan_sub.dataset,
         pool,
         dataset_config,

server/src/handlers/group_handler.rs

@@ -604,12 +604,10 @@ pub async fn delete_chunk_group(
     )
     .await?;
 
-    let deleted_at = chrono::Utc::now().naive_utc();
-
     delete_group_by_id_query(
         group_id,
         dataset_org_plan_sub.dataset,
-        deleted_at,
+        chrono::Utc::now().naive_utc(),
         data.delete_chunks,
         delete_group_pool,
         dataset_config,

server/src/lib.rs

@@ -141,7 +141,7 @@ impl Modify for SecurityAddon {
             name = "BSL",
             url = "https://github.com/devflowinc/trieve/blob/main/LICENSE.txt",
         ),
-        version = "0.11.8",
+        version = "0.11.9",
     ),
     servers(
         (url = "https://api.trieve.ai",

server/src/operators/file_operator.rs

@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-
 use super::chunk_operator::{create_chunk_metadata, get_row_count_for_organization_id_query};
 use super::clickhouse_operator::{ClickHouseEvent, EventQueue};
 use super::group_operator::{create_group_from_file_query, create_groups_query};
@@ -9,6 +7,7 @@ use crate::data::models::FileDTO;
 use crate::data::models::{Dataset, DatasetAndOrgWithSubAndPlan, DatasetConfiguration, EventType};
 use crate::handlers::chunk_handler::ChunkReqPayload;
 use crate::handlers::file_handler::UploadFileReqPayload;
+use crate::operators::group_operator::delete_group_by_file_id_query;
 use crate::{data::models::WorkerEvent, get_env};
 use crate::{
     data::models::{File, Pool},
@@ -18,11 +17,11 @@ use actix_web::web;
 use diesel::dsl::sql;
 use diesel::prelude::*;
 use diesel::sql_types::BigInt;
-use diesel_async::scoped_futures::ScopedFutureExt;
-use diesel_async::{AsyncConnection, RunQueryDsl};
+use diesel_async::RunQueryDsl;
 use redis::aio::MultiplexedConnection;
 use regex::Regex;
 use s3::{creds::Credentials, Bucket, Region};
+use std::collections::HashMap;
 
 #[tracing::instrument]
 pub fn get_aws_bucket() -> Result<Bucket, ServiceError> {
@@ -359,6 +358,7 @@ pub async fn get_dataset_file_query(
 #[tracing::instrument(skip(pool))]
 pub async fn delete_file_query(
     file_uuid: uuid::Uuid,
+    delete_chunks: Option<bool>,
     dataset: Dataset,
     pool: web::Data<Pool>,
     dataset_config: DatasetConfiguration,
@@ -383,29 +383,28 @@ pub async fn delete_file_query(
         .await
         .map_err(|_| ServiceError::BadRequest("Could not delete file from S3".to_string()))?;
 
-    let transaction_result = conn
-        .transaction::<_, diesel::result::Error, _>(|conn| {
-            async {
-                diesel::delete(
-                    files_columns::files
-                        .filter(files_columns::id.eq(file_uuid))
-                        .filter(files_columns::dataset_id.eq(dataset.clone().id)),
-                )
-                .execute(conn)
-                .await?;
-
-                Ok(())
-            }
-            .scope_boxed()
-        })
-        .await;
-
-    match transaction_result {
-        Ok(_) => (),
-        Err(e) => {
-            log::error!("Error deleting file with transaction {:?}", e);
-            return Err(ServiceError::BadRequest("Could not delete file".to_string()).into());
-        }
+    diesel::delete(
+        files_columns::files
+            .filter(files_columns::id.eq(file_uuid))
+            .filter(files_columns::dataset_id.eq(dataset.clone().id)),
+    )
+    .execute(&mut conn)
+    .await
+    .map_err(|e| {
+        log::error!("Error deleting file {:?}", e);
+        ServiceError::BadRequest("Could not delete file".to_string())
+    })?;
+
+    if delete_chunks.is_some_and(|delete_chunks| delete_chunks) {
+        delete_group_by_file_id_query(
+            file_uuid,
+            dataset,
+            chrono::Utc::now().naive_utc(),
+            Some(true),
+            pool.clone(),
+            dataset_config,
+        )
+        .await?;
     }
 
     Ok(())

server/src/operators/group_operator.rs

@@ -354,6 +354,104 @@ pub async fn delete_group_by_id_query(
     }
 }
 
+#[tracing::instrument(skip(pool))]
+pub async fn delete_group_by_file_id_query(
+    file_id: uuid::Uuid,
+    dataset: Dataset,
+    deleted_at: chrono::NaiveDateTime,
+    delete_chunks: Option<bool>,
+    pool: web::Data<Pool>,
+    dataset_config: DatasetConfiguration,
+) -> Result<(), ServiceError> {
+    use crate::data::schema::chunk_group::dsl as chunk_group_columns;
+    use crate::data::schema::chunk_group_bookmarks::dsl as chunk_group_bookmarks_columns;
+    use crate::data::schema::chunk_metadata::dsl as chunk_metadata_columns;
+    use crate::data::schema::groups_from_files::dsl as groups_from_files_columns;
+
+    let mut conn = pool.get().await.map_err(|_e| {
+        ServiceError::InternalServerError("Failed to get postgres connection".to_string())
+    })?;
+
+    let group_id: uuid::Uuid = groups_from_files_columns::groups_from_files
+        .filter(groups_from_files_columns::file_id.eq(file_id))
+        .select(groups_from_files_columns::group_id)
+        .first::<uuid::Uuid>(&mut conn)
+        .await
+        .map_err(|_err| {
+            ServiceError::BadRequest("Error getting group id for file_id".to_string())
+        })?;
+
+    let delete_chunks = delete_chunks.unwrap_or(false);
+    let chunks = chunk_group_bookmarks_columns::chunk_group_bookmarks
+        .inner_join(chunk_metadata_columns::chunk_metadata)
+        .filter(chunk_group_bookmarks_columns::group_id.eq(group_id))
+        .select(ChunkMetadataTable::as_select())
+        .load::<ChunkMetadataTable>(&mut conn)
+        .await
+        .map_err(|_err| ServiceError::BadRequest("Error getting chunks".to_string()))?;
+
+    let transaction_result = conn
+        .transaction::<_, diesel::result::Error, _>(|conn| {
+            async move {
+                diesel::delete(
+                    groups_from_files_columns::groups_from_files
+                        .filter(groups_from_files_columns::group_id.eq(group_id))
+                        .filter(groups_from_files_columns::created_at.le(deleted_at)),
+                )
+                .execute(conn)
+                .await?;
+
+                diesel::delete(
+                    chunk_group_bookmarks_columns::chunk_group_bookmarks
+                        .filter(chunk_group_bookmarks_columns::group_id.eq(group_id))
+                        .filter(chunk_group_bookmarks_columns::created_at.le(deleted_at)),
+                )
+                .execute(conn)
+                .await?;
+
+                diesel::delete(
+                    chunk_group_columns::chunk_group
+                        .filter(chunk_group_columns::id.eq(group_id))
+                        .filter(chunk_group_columns::dataset_id.eq(dataset.id))
+                        .filter(chunk_group_columns::created_at.le(deleted_at)),
+                )
+                .execute(conn)
+                .await?;
+
+                Ok(())
+            }
+            .scope_boxed()
+        })
+        .await;
+
+    if delete_chunks {
+        let chunk_ids = chunks.iter().map(|chunk| chunk.id).collect();
+        delete_chunk_metadata_query(
+            chunk_ids,
+            deleted_at,
+            dataset.clone(),
+            pool.clone(),
+            dataset_config.clone(),
+        )
+        .await?;
+    } else {
+        let remove_chunks_from_groups_futures = chunks.iter().map(|chunk| {
+            remove_bookmark_from_qdrant_query(
+                chunk.qdrant_point_id,
+                group_id,
+                dataset_config.clone(),
+            )
+        });
+
+        futures::future::join_all(remove_chunks_from_groups_futures).await;
+    }
+
+    match transaction_result {
+        Ok(_) => Ok(()),
+        Err(_) => Err(ServiceError::BadRequest("Error deleting group".to_string())),
+    }
+}
+
 #[tracing::instrument(skip(pool))]
 pub async fn update_chunk_group_query(
     group: ChunkGroup,