POST v1/integrations-save

Signed-off-by: V4LER11 <[email protected]>
smallcloudai · Nov 5, 2024 · d7ec356 · d7ec356
1 parent 2f48bb8
commit d7ec356
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 8 deletions.
diff --git a/src/http/routers/v1.rs b/src/http/routers/v1.rs
@@ -35,7 +35,7 @@ use crate::http::routers::v1::subchat::{handle_v1_subchat, handle_v1_subchat_sin
 use crate::http::routers::v1::vecdb::{handle_v1_vecdb_search, handle_v1_vecdb_status};
 #[cfg(feature="vecdb")]
 use crate::http::routers::v1::handlers_memdb::{handle_mem_query, handle_mem_add, handle_mem_erase, handle_mem_update_used, handle_mem_block_until_vectorized, handle_mem_list, handle_ongoing_update_or_create, handle_ongoing_dump};
-use crate::http::routers::v1::integrations::handle_v1_integrations;
+use crate::http::routers::v1::integrations::{handle_v1_integrations, handle_v1_integrations_save};
 use crate::http::utils::telemetry_wrapper;
 
 pub mod code_completion;
@@ -85,6 +85,7 @@ pub fn make_v1_router() -> Router {
         .route("/tools", telemetry_get!(handle_v1_tools))
         .route("/tools-check-if-confirmation-needed", telemetry_post!(handle_v1_tools_check_if_confirmation_needed))
         .route("/integrations", telemetry_get!(handle_v1_integrations))
+        .route("/integrations-save", telemetry_post!(handle_v1_integrations_save))
 
         .route("/lsp-initialize", telemetry_post!(handle_v1_lsp_initialize))
         .route("/lsp-did-changed", telemetry_post!(handle_v1_lsp_did_change))

diff --git a/src/http/routers/v1/integrations.rs b/src/http/routers/v1/integrations.rs
@@ -3,17 +3,18 @@ use axum::Extension;
 use axum::http::{Response, StatusCode};
 use tokio::sync::RwLock as ARwLock;
 use hyper::Body;
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
+
 use crate::custom_error::ScratchError;
 use crate::global_context::GlobalContext;
-use crate::integrations::load_integration_schema_and_json;
+use crate::integrations::{integrations_paths, load_integration_schema_and_json, validate_integration_value};
 
 
-#[derive(Serialize)]
+#[derive(Serialize, Deserialize)]
 struct IntegrationItem {
     name: String,
-    schema: Value,
+    schema: Option<Value>,
     value: Value,
 }
 
@@ -28,7 +29,7 @@ pub async fn handle_v1_integrations(
     for (name, (schema, value)) in schemas_and_json_dict {
         let item = IntegrationItem {
             name,
-            schema,
+            schema: Some(schema),
             value,
         };
 
@@ -42,3 +43,43 @@ pub async fn handle_v1_integrations(
         .body(Body::from(payload))
         .unwrap())
 }
+
+
+pub async fn handle_v1_integrations_save(
+    Extension(gcx): Extension<Arc<ARwLock<GlobalContext>>>,
+    body_bytes: hyper::body::Bytes,
+) -> axum::response::Result<Response<Body>, ScratchError> {
+    let post = serde_json::from_slice::<IntegrationItem>(&body_bytes)
+        .map_err(|e| ScratchError::new(StatusCode::UNPROCESSABLE_ENTITY, format!("JSON problem: {}", e)))?;
+
+    let yaml_value: serde_yaml::Value = serde_json::to_string(&post.value).map_err(|e|e.to_string())
+        .and_then(|s|serde_yaml::from_str(&s).map_err(|e|e.to_string()))
+        .map_err(|e| ScratchError::new(StatusCode::UNPROCESSABLE_ENTITY, format!("ERROR converting JSON to YAML: {}", e)))?;
+
+    let yaml_value = validate_integration_value(&post.name, yaml_value)
+        .map_err(|e| ScratchError::new(StatusCode::UNPROCESSABLE_ENTITY, format!("ERROR validating integration value: {}", e)))?;
+
+    let integr_paths = integrations_paths(gcx.clone()).await;
+
+    let path = integr_paths.get(&post.name)
+        .ok_or(ScratchError::new(StatusCode::UNPROCESSABLE_ENTITY, format!("Integration {} not found", post.name)))?;
+
+    let mut file = tokio::fs::OpenOptions::new()
+        .write(true)
+        .truncate(true)
+        .open(path)
+        .await
+        .map_err(|e| ScratchError::new(StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to open file: {}", e)))?;
+
+    let yaml_string = serde_yaml::to_string(&yaml_value)
+        .map_err(|e| ScratchError::new(StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to convert YAML to string: {}", e)))?;
+
+    tokio::io::AsyncWriteExt::write_all(&mut file, yaml_string.as_bytes()).await
+        .map_err(|e| ScratchError::new(StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to write to file: {}", e)))?;
+
+    Ok(Response::builder()
+       .status(StatusCode::OK)
+       .header("Content-Type", "application/json")
+       .body(Body::from(format!("Integration {} saved", post.name)))
+       .unwrap())
+}
diff --git a/src/integrations/mod.rs b/src/integrations/mod.rs
@@ -4,14 +4,19 @@ use indexmap::IndexMap;
 use serde_json::json;
 use tracing::warn;
 use tokio::sync::{Mutex as AMutex, RwLock as ARwLock};
+
 use crate::global_context::GlobalContext;
 use crate::integrations::integr::Integration;
+use crate::integrations::integr_chrome::IntegrationChrome;
+use crate::integrations::integr_github::IntegrationGitHub;
+use crate::integrations::integr_gitlab::IntegrationGitLab;
+use crate::integrations::integr_pdb::IntegrationPdb;
+use crate::integrations::integr_postgres::IntegrationPostgres;
 use crate::tools::tools_description::Tool;
 use crate::yaml_configs::create_configs::read_yaml_into_value;
 
 pub mod sessions;
 pub mod process_io_utils;
-
 pub mod integr_github;
 pub mod integr_gitlab;
 pub mod integr_pdb;
@@ -20,7 +25,7 @@ pub mod integr_postgres;
 mod integr;
 
 
-// hint: when adding integration, update: DEFAULT_INTEGRATION_VALUES, integrations_paths, load_integration_tools, load_integration_schema_and_json
+// hint: when adding integration, update: DEFAULT_INTEGRATION_VALUES, integrations_paths, validate_integration_value, load_integration_tools, load_integration_schema_and_json
 
 
 pub const DEFAULT_INTEGRATION_VALUES: &[(&str, &str)] = &[
@@ -41,6 +46,32 @@ pub async fn integrations_paths(gcx: Arc<ARwLock<GlobalContext>>) -> IndexMap<St
     }).collect()
 }
 
+pub fn validate_integration_value(name: &str, value: serde_yaml::Value) -> Result<serde_yaml::Value, String> {
+    match name {
+        "github" => {
+            let integration: IntegrationGitHub = serde_yaml::from_value(value).map_err(|e| e.to_string())?;
+            serde_yaml::to_value(integration).map_err(|e| e.to_string())
+        }
+        "gitlab" => {
+            let integration: IntegrationGitLab = serde_yaml::from_value(value).map_err(|e| e.to_string())?;
+            serde_yaml::to_value(integration).map_err(|e| e.to_string())
+        }
+        "pdb" => {
+            let integration: IntegrationPdb = serde_yaml::from_value(value).map_err(|e| e.to_string())?;
+            serde_yaml::to_value(integration).map_err(|e| e.to_string())
+        }
+        "postgres" => {
+            let integration: IntegrationPostgres = serde_yaml::from_value(value).map_err(|e| e.to_string())?;
+            serde_yaml::to_value(integration).map_err(|e| e.to_string())
+        }
+        "chrome" => {
+            let integration: IntegrationChrome = serde_yaml::from_value(value).map_err(|e| e.to_string())?;
+            serde_yaml::to_value(integration).map_err(|e| e.to_string())
+        }
+        _ => Err(format!("Unknown integration type: {}", name)),
+    }
+}
+
 pub async fn load_integration_tools(
     gcx: Arc<ARwLock<GlobalContext>>,
 ) -> IndexMap<String, Arc<AMutex<Box<dyn Tool + Send>>>> {