From f999db9a9e3f9eff2f5bef42b44a615b62b00d36 Mon Sep 17 00:00:00 2001 From: Alek Westover Date: Wed, 26 Jul 2023 13:53:56 -0400 Subject: [PATCH] Joonas suggested an amazing refactor of extension_server json parsing --- compute_tools/src/extension_server.rs | 89 ++++++++++----------------- 1 file changed, 32 insertions(+), 57 deletions(-) diff --git a/compute_tools/src/extension_server.rs b/compute_tools/src/extension_server.rs index 1b257f8728..54acf0d8af 100644 --- a/compute_tools/src/extension_server.rs +++ b/compute_tools/src/extension_server.rs @@ -61,7 +61,7 @@ use anyhow::Context; use anyhow::{self, Result}; use futures::future::join_all; use remote_storage::*; -use serde_json::{self, Value}; +use serde_json; use std::collections::HashMap; use std::io::Read; use std::num::{NonZeroU32, NonZeroUsize}; @@ -123,67 +123,46 @@ pub async fn get_available_extensions( .download_stream .read_to_end(&mut ext_idx_buffer) .await?; - let ext_index_str = str::from_utf8(&ext_idx_buffer).expect("json err"); - let ext_index_full: Value = serde_json::from_str(ext_index_str)?; - let ext_index_full = ext_index_full.as_object().expect("json err"); - info!("ext_index: {:?}", &ext_index_full); - let public_extensions = ext_index_full["public_extensions"] - .as_array() - .expect("json err"); - // TODO: maybe enabled_extensions should be a HashSet? probably doesn't matter - let mut enabled_extensions = public_extensions - .iter() - .map(|x| x.as_str().expect("json err")) - .collect::>(); - for custom_extension in custom_extensions { - enabled_extensions.push(custom_extension); + #[derive(Debug, serde::Deserialize)] + struct Index { + public_extensions: Vec, + library_index: HashMap, + extension_data: HashMap, } - let library_index = ext_index_full["library_index"] - .as_object() - .context("error parsing json")?; - let mut parsed_lib_index = HashMap::new(); - for (key, val) in library_index { - let parsed_val = val.as_str().expect("json err").to_string(); - parsed_lib_index.insert(key.to_string(), parsed_val); + #[derive(Debug, serde::Deserialize)] + struct ExtensionData { + control_data: HashMap, + archive_path: String, } - let all_extension_data = ext_index_full["extension_data"] - .as_object() - .context("error parsing json")?; + let ext_index_full = serde_json::from_slice::(&ext_idx_buffer)?; + let mut enabled_extensions = ext_index_full.public_extensions; + enabled_extensions.extend_from_slice(custom_extensions); + let library_index = ext_index_full.library_index; + let all_extension_data = ext_index_full.extension_data; info!("enabled_extensions: {:?}", enabled_extensions); let mut ext_remote_paths = HashMap::new(); let mut file_create_tasks = Vec::new(); for extension in enabled_extensions { - let ext_data = all_extension_data[extension] - .as_object() - .context("error parsing json")?; - - let control_files = ext_data["control_data"] - .as_object() - .context("error parsing json")?; - for (control_file, control_contents) in control_files { + let ext_data = &all_extension_data[&extension]; + for (control_file, control_contents) in &ext_data.control_data { let control_path = local_sharedir.join(control_file); - let control_contents_str = control_contents.as_str().context("error parsing json")?; info!("writing file {:?}{:?}", control_path, control_contents); - file_create_tasks.push(tokio::fs::write(control_path, control_contents_str)); + file_create_tasks.push(tokio::fs::write(control_path, control_contents)); } - - let ext_archive_path = ext_data["archive_path"] - .as_str() - .context("error parsing json")?; ext_remote_paths.insert( extension.to_string(), - RemotePath::from_string(ext_archive_path)?, + RemotePath::from_string(&ext_data.archive_path)?, ); } let results = join_all(file_create_tasks).await; for result in results { result?; } - Ok((ext_remote_paths, parsed_lib_index)) + Ok((ext_remote_paths, library_index)) } // download the archive for a given extension, @@ -237,24 +216,20 @@ pub async fn download_extension( // This function initializes the necessary structs to use remote storage (should be fairly cheap) pub fn init_remote_storage(remote_ext_config: &str) -> anyhow::Result { - let remote_ext_config: serde_json::Value = serde_json::from_str(remote_ext_config)?; + #[derive(Debug, serde::Deserialize)] + struct RemoteExtJson { + bucket: String, + region: String, + endpoint: Option, + prefix: Option, + } + let remote_ext_json = serde_json::from_str::(remote_ext_config)?; - let remote_ext_bucket = remote_ext_config["bucket"] - .as_str() - .context("config parse error")?; - let remote_ext_region = remote_ext_config["region"] - .as_str() - .context("config parse error")?; - let remote_ext_endpoint = remote_ext_config["endpoint"].as_str(); - let remote_ext_prefix = remote_ext_config["prefix"].as_str(); - - // If needed, it is easy to allow modification of other parameters - // however, default values should be fine for now let config = S3Config { - bucket_name: remote_ext_bucket.to_string(), - bucket_region: remote_ext_region.to_string(), - prefix_in_bucket: remote_ext_prefix.map(str::to_string), - endpoint: remote_ext_endpoint.map(|x| x.to_string()), + bucket_name: remote_ext_json.bucket, + bucket_region: remote_ext_json.region, + prefix_in_bucket: remote_ext_json.prefix, + endpoint: remote_ext_json.endpoint, concurrency_limit: NonZeroUsize::new(100).expect("100 != 0"), max_keys_per_list_response: None, };