feat(document): support crawl preset document (#2907)

* finish api Signed-off-by: xxs-wallace <[email protected]> * fix delete by id Signed-off-by: xxs-wallace <[email protected]> * fix active Signed-off-by: xxs-wallace <[email protected]> * fix name Signed-off-by: xxs-wallace <[email protected]> * fix api Signed-off-by: xxs-wallace <[email protected]> * add ut Signed-off-by: xxs-wallace <[email protected]> --------- Signed-off-by: xxs-wallace <[email protected]>
TabbyML · Aug 20, 2024 · 2c947f2 · 2c947f2
1 parent 0e1642f
commit 2c947f2
Show file tree

Hide file tree

Showing 8 changed files with 634 additions and 30 deletions.
diff --git a/ee/tabby-db/src/web_documents.rs b/ee/tabby-db/src/web_documents.rs
@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use chrono::{DateTime, Utc};
 use sqlx::{prelude::FromRow, query};
 use tabby_db_macros::query_paged_as;
@@ -22,9 +22,9 @@ impl DbConn {
         limit: Option<usize>,
         skip_id: Option<i32>,
         backwards: bool,
-        is_preset: Option<bool>,
+        is_preset: bool,
     ) -> Result<Vec<WebDocumentDAO>> {
-        let condition = is_preset.map(|is_preset| format!("is_preset={}", is_preset));
+        let condition = Some(format!("is_preset={}", is_preset));
 
         let urls = query_paged_as!(
             WebDocumentDAO,
@@ -58,6 +58,16 @@ impl DbConn {
         Ok(res.last_insert_rowid())
     }
 
+    pub async fn deactivate_preset_web_document(&self, name: String) -> Result<()> {
+        let res = query!("DELETE FROM web_documents WHERE name = ?;", name)
+            .execute(&self.pool)
+            .await?;
+        if res.rows_affected() != 1 {
+            return Err(anyhow!("No preset web document to deactivate"));
+        }
+        Ok(())
+    }
+
     pub async fn delete_web_document(&self, id: i64) -> Result<()> {
         query!("DELETE FROM web_documents WHERE id = ?;", id)
             .execute(&self.pool)

diff --git a/ee/tabby-schema/graphql/schema.graphql b/ee/tabby-schema/graphql/schema.graphql
@@ -257,6 +257,7 @@ type CustomWebDocument {
   name: String!
   id: ID!
   createdAt: DateTime!
+  updatedAt: DateTime!
   jobInfo: JobInfo!
 }
 
@@ -505,7 +506,7 @@ type Mutation {
   deleteWebCrawlerUrl(id: ID!): Boolean!
   createCustomDocument(input: CreateCustomDocumentInput!): ID!
   deleteCustomDocument(id: ID!): Boolean!
-  setPresetDocumentActive(input: SetPresetDocumentActiveInput!): ID!
+  setPresetDocumentActive(input: SetPresetDocumentActiveInput!): Boolean!
   "Delete pair of user message and bot response in a thread."
   deleteThreadMessagePair(threadId: ID!, userMessageId: ID!, assistantMessageId: ID!): Boolean!
 }
@@ -539,11 +540,11 @@ type PresetDocumentEdge {
 }
 
 type PresetWebDocument {
-  name: String!
   id: ID!
-  active: Boolean!
+  name: String!
   "`updated_at` is only filled when the preset is active."
   updatedAt: DateTime
+  "`job_info` is only filled when the preset is active."
   jobInfo: JobInfo
 }
 
@@ -619,7 +620,7 @@ type Query {
   """
   threadMessages(threadId: ID!, after: String, before: String, first: Int, last: Int): MessageConnection!
   customWebDocuments(after: String, before: String, first: Int, last: Int): CustomDocumentConnection!
-  presetWebDocuments(after: String, before: String, first: Int, last: Int, active: Boolean!): PresetDocumentConnection!
+  presetWebDocuments(after: String, before: String, first: Int, last: Int, isActive: Boolean!): PresetDocumentConnection!
 }
 
 type RefreshTokenResponse {

diff --git a/ee/tabby-schema/src/schema/constants.rs b/ee/tabby-schema/src/schema/constants.rs
@@ -6,7 +6,7 @@ lazy_static! {
     pub static ref USERNAME_REGEX: Regex =
         Regex::new(r"^[^0-9±!@£$%^&*_+§¡€#¢¶•ªº«\\/<>?:;|=.,]{2,20}$").unwrap();
     pub static ref WEB_DOCUMENT_NAME_REGEX: Regex =
-        Regex::new(r"^[A-Za-z][A-Za-z0-9\ ]*$").unwrap();
+        Regex::new(r"^[A-Za-z][A-Za-z0-9]*(\ [A-Za-z0-9]+)*$").unwrap();
 }
 
 #[cfg(test)]
@@ -53,6 +53,7 @@ mod tests {
             (" abc 123", false),
             ("abc123*", false),
             ("abc123_", false),
+            ("abc  123", false), // two space
         ];
 
         for (name, expected) in test_cases {

diff --git a/ee/tabby-schema/src/schema/mod.rs b/ee/tabby-schema/src/schema/mod.rs
@@ -596,7 +596,7 @@ impl Query {
         before: Option<String>,
         first: Option<i32>,
         last: Option<i32>,
-        active: bool,
+        is_active: bool,
     ) -> Result<Connection<PresetWebDocument>> {
         query_async(
             after,
@@ -606,7 +606,7 @@ impl Query {
             |after, before, first, last| async move {
                 ctx.locator
                     .web_documents()
-                    .list_preset_web_documents(after, before, first, last, active)
+                    .list_preset_web_documents(after, before, first, last, is_active)
                     .await
             },
         )
@@ -986,14 +986,13 @@ impl Mutation {
     async fn set_preset_document_active(
         ctx: &Context,
         input: SetPresetDocumentActiveInput,
-    ) -> Result<ID> {
+    ) -> Result<bool> {
         input.validate()?;
-        let id = ctx
-            .locator
+        ctx.locator
             .web_documents()
             .set_preset_web_documents_active(input.name, input.active)
             .await?;
-        Ok(id)
+        Ok(true)
     }
 
     /// Delete pair of user message and bot response in a thread.

diff --git a/ee/tabby-schema/src/schema/web_documents.rs b/ee/tabby-schema/src/schema/web_documents.rs
@@ -12,17 +12,19 @@ pub struct CustomWebDocument {
     pub name: String,
     pub id: ID,
     pub created_at: DateTime<Utc>,
+    pub updated_at: DateTime<Utc>,
     pub job_info: JobInfo,
 }
 
 #[derive(GraphQLObject)]
 #[graphql(context = Context)]
 pub struct PresetWebDocument {
-    pub name: String,
     pub id: ID,
-    pub active: bool,
+
+    pub name: String,
     /// `updated_at` is only filled when the preset is active.
     pub updated_at: Option<DateTime<Utc>>,
+    /// `job_info` is only filled when the preset is active.
     pub job_info: Option<JobInfo>,
 }
 
@@ -79,7 +81,7 @@ impl NodeType for PresetWebDocument {
     type Cursor = String;
 
     fn cursor(&self) -> Self::Cursor {
-        self.name.clone()
+        self.id.to_string()
     }
 
     fn connection_type_name() -> &'static str {
@@ -109,7 +111,7 @@ pub trait WebDocumentService: Send + Sync {
         before: Option<String>,
         first: Option<usize>,
         last: Option<usize>,
-        active: bool,
+        is_active: bool,
     ) -> Result<Vec<PresetWebDocument>>;
-    async fn set_preset_web_documents_active(&self, name: String, active: bool) -> Result<ID>;
+    async fn set_preset_web_documents_active(&self, name: String, active: bool) -> Result<()>;
 }
diff --git a/ee/tabby-webserver/src/service/mod.rs b/ee/tabby-webserver/src/service/mod.rs
@@ -7,6 +7,7 @@ pub mod event_logger;
 pub mod integration;
 pub mod job;
 mod license;
+mod preset_web_documents_data;
 pub mod repository;
 mod setting;
 mod thread;