[core] Add parentId parameter in core/front apis (#9261)

dust-tt · Dec 11, 2024 · 5e2bd9c · 5e2bd9c
1 parent a74674e
commit 5e2bd9c
Show file tree

Hide file tree

Showing 36 changed files with 335 additions and 28 deletions.
diff --git a/connectors/src/lib/data_sources.ts b/connectors/src/lib/data_sources.ts
@@ -153,6 +153,7 @@ async function _upsertToDatasource({
         title,
         mime_type: mimeType,
         tags: tags?.map((tag) => tag.substring(0, 512)),
+        parent_id: parents[1] ?? null,
         parents,
         light_document_output: true,
         upsert_context: upsertContext,
@@ -770,6 +771,7 @@ export async function upsertTableFromCsv({
     `/data_sources/${dataSourceConfig.dataSourceId}/tables/csv`;
   const dustRequestPayload: UpsertTableFromCsvRequestType = {
     name: tableName,
+    parentId: parents[1] ?? null,
     parents,
     description: tableDescription,
     csv: tableCsv,
@@ -1148,6 +1150,7 @@ export async function upsertFolderNode({
     folderId,
     timestampMs ? timestampMs : now.getTime(),
     title,
+    parents[1] ?? null,
     parents
   );
 

diff --git a/core/bin/core_api.rs b/core/bin/core_api.rs
@@ -1468,6 +1468,7 @@ async fn data_sources_documents_update_tags(
 
 #[derive(serde::Deserialize)]
 struct DataSourcesDocumentsUpdateParentsPayload {
+    parent_id: Option<String>,
     parents: Vec<String>,
 }
 
@@ -1489,6 +1490,17 @@ async fn data_sources_documents_update_parents(
         );
     }
 
+    if let Some(parent_id) = &payload.parent_id {
+        if payload.parents.get(1) != Some(parent_id) {
+            return error_response(
+                StatusCode::BAD_REQUEST,
+                "invalid_parent_id",
+                "Failed to update document parents - parents[1] and parent_id should be equal",
+                None,
+            );
+        }
+    }
+
     match state
         .store
         .load_data_source(&project, &data_source_id)
@@ -1617,6 +1629,7 @@ struct DataSourcesDocumentsUpsertPayload {
     document_id: String,
     timestamp: Option<u64>,
     tags: Vec<String>,
+    parent_id: Option<String>,
     parents: Vec<String>,
     source_url: Option<String>,
     section: Section,
@@ -1648,6 +1661,17 @@ async fn data_sources_documents_upsert(
         );
     }
 
+    if let Some(parent_id) = &payload.parent_id {
+        if payload.parents.get(1) != Some(parent_id) {
+            return error_response(
+                StatusCode::BAD_REQUEST,
+                "invalid_parent_id",
+                "Failed to upsert document - parents[1] and parent_id should be equal",
+                None,
+            );
+        }
+    }
+
     match state
         .store
         .load_data_source(&project, &data_source_id)
@@ -2074,6 +2098,7 @@ struct DatabasesTablesUpsertPayload {
     description: String,
     timestamp: Option<u64>,
     tags: Vec<String>,
+    parent_id: Option<String>,
     parents: Vec<String>,
 
     // Remote DB specifics
@@ -2103,6 +2128,17 @@ async fn tables_upsert(
         );
     }
 
+    if let Some(parent_id) = &payload.parent_id {
+        if payload.parents.get(1) != Some(parent_id) {
+            return error_response(
+                StatusCode::BAD_REQUEST,
+                "invalid_parent_id",
+                "Failed to upsert table - parents[1] and parent_id should be equal",
+                None,
+            );
+        }
+    }
+
     match state
         .store
         .upsert_data_source_table(
@@ -2351,6 +2387,17 @@ async fn tables_update_parents(
         );
     }
 
+    if let Some(parent_id) = &payload.parent_id {
+        if payload.parents.get(1) != Some(parent_id) {
+            return error_response(
+                StatusCode::BAD_REQUEST,
+                "invalid_parent_id",
+                "Failed to update table parents - parents[1] and parent_id should be equal",
+                None,
+            );
+        }
+    }
+
     match state
         .store
         .load_data_source_table(&project, &data_source_id, &table_id)
@@ -2725,6 +2772,7 @@ async fn tables_rows_list(
 struct FoldersUpsertPayload {
     folder_id: String,
     timestamp: Option<u64>,
+    parent_id: Option<String>,
     parents: Vec<String>,
     title: String,
 }
@@ -2747,6 +2795,17 @@ async fn folders_upsert(
         );
     }
 
+    if let Some(parent_id) = &payload.parent_id {
+        if payload.parents.get(1) != Some(parent_id) {
+            return error_response(
+                StatusCode::BAD_REQUEST,
+                "invalid_parent_id",
+                "Failed to upsert folder - parents[1] and parent_id should be equal",
+                None,
+            );
+        }
+    }
+
     match state
         .store
         .upsert_data_source_folder(

diff --git a/core/src/data_sources/data_source.rs b/core/src/data_sources/data_source.rs
@@ -142,6 +142,7 @@ pub struct Document {
     pub title: String,
     pub mime_type: String,
     pub tags: Vec<String>,
+    pub parent_id: Option<String>,
     pub parents: Vec<String>,
     pub source_url: Option<String>,
     pub hash: String,
@@ -161,6 +162,7 @@ impl Document {
         title: &str,
         mime_type: &str,
         tags: &Vec<String>,
+        parent_id: &Option<String>,
         parents: &Vec<String>,
         source_url: &Option<String>,
         hash: &str,
@@ -174,6 +176,7 @@ impl Document {
             title: title.to_string(),
             mime_type: mime_type.to_string(),
             tags: tags.clone(),
+            parent_id: parent_id.clone(),
             parents: parents.clone(),
             source_url: source_url.clone(),
             hash: hash.to_string(),
@@ -216,6 +219,7 @@ impl From<Document> for Node {
             document.timestamp,
             &document.title,
             &document.mime_type,
+            document.parent_id,
             document.parents.clone(),
         )
     }
@@ -686,7 +690,8 @@ impl DataSource {
             title.as_deref().unwrap_or(document_id),
             mime_type.as_deref().unwrap_or("application/octet-stream"),
             &tags,
-            &parents,
+            &parents.get(1).cloned(),
+            parents,
             source_url,
             &document_hash,
             full_text.len() as u64,

diff --git a/core/src/data_sources/folder.rs b/core/src/data_sources/folder.rs
@@ -6,6 +6,7 @@ pub struct Folder {
     folder_id: String,
     timestamp: u64,
     title: String,
+    parent_id: Option<String>,
     parents: Vec<String>,
 }
 
@@ -18,13 +19,15 @@ impl Folder {
         folder_id: String,
         timestamp: u64,
         title: String,
+        parent_id: Option<String>,
         parents: Vec<String>,
     ) -> Self {
         Folder {
-            data_source_id: data_source_id,
-            folder_id: folder_id,
+            data_source_id,
+            folder_id,
             timestamp,
-            title: title,
+            title,
+            parent_id,
             parents,
         }
     }
@@ -41,6 +44,9 @@ impl Folder {
     pub fn title(&self) -> &str {
         &self.title
     }
+    pub fn parent_id(&self) -> &Option<String> {
+        &self.parent_id
+    }
     pub fn parents(&self) -> &Vec<String> {
         &self.parents
     }

diff --git a/core/src/data_sources/node.rs b/core/src/data_sources/node.rs
@@ -17,6 +17,7 @@ pub struct Node {
     timestamp: u64,
     title: String,
     mime_type: String,
+    parent_id: Option<String>,
     parents: Vec<String>,
 }
 
@@ -28,6 +29,7 @@ impl Node {
         timestamp: u64,
         title: &str,
         mime_type: &str,
+        parent_id: Option<String>,
         parents: Vec<String>,
     ) -> Self {
         Node {
@@ -37,6 +39,7 @@ impl Node {
             timestamp,
             title: title.to_string(),
             mime_type: mime_type.to_string(),
+            parent_id: parent_id.clone(),
             parents,
         }
     }
@@ -70,6 +73,7 @@ impl Node {
             self.node_id,
             self.timestamp,
             self.title,
+            self.parent_id,
             self.parents,
         )
     }

diff --git a/core/src/databases/table.rs b/core/src/databases/table.rs
@@ -60,6 +60,7 @@ pub struct Table {
     tags: Vec<String>,
     title: String,
     mime_type: String,
+    parent_id: Option<String>,
     parents: Vec<String>,
 
     schema: Option<TableSchema>,
@@ -81,25 +82,27 @@ impl Table {
         title: String,
         mime_type: String,
         tags: Vec<String>,
+        parent_id: Option<String>,
         parents: Vec<String>,
         schema: Option<TableSchema>,
         schema_stale_at: Option<u64>,
         remote_database_table_id: Option<String>,
         remote_database_secret_id: Option<String>,
     ) -> Self {
         Table {
-            project: project,
-            data_source_id: data_source_id,
+            project,
+            data_source_id,
             created,
-            table_id: table_id,
-            name: name,
-            description: description,
+            table_id,
+            name,
+            description,
             timestamp,
             tags,
-            title: title,
-            mime_type: mime_type,
+            title,
+            mime_type,
+            parent_id,
             parents,
-            schema: schema,
+            schema,
             schema_stale_at,
             remote_database_table_id,
             remote_database_secret_id,
@@ -124,6 +127,9 @@ impl Table {
     pub fn mime_type(&self) -> &str {
         &self.mime_type
     }
+    pub fn parent_id(&self) -> &Option<String> {
+        &self.parent_id
+    }
     pub fn parents(&self) -> &Vec<String> {
         &self.parents
     }
@@ -228,7 +234,8 @@ impl From<Table> for Node {
             table.timestamp,
             &table.title,
             &table.mime_type,
-            table.parents.clone(),
+            table.parents.get(1).cloned(),
+            table.parents,
         )
     }
 }
@@ -574,6 +581,7 @@ mod tests {
             "test_dbml".to_string(),
             "text/plain".to_string(),
             vec![],
+            None,
             vec![],
             Some(schema),
             None,