diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 4f227bc8ae..374a923d78 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -58,6 +58,7 @@ pub struct ComputeNode { pub ext_remote_storage: Option, // (key: extension name, value: path to extension archive in remote storage) pub ext_remote_paths: OnceLock>, + pub library_index: OnceLock>, pub already_downloaded_extensions: Mutex>, pub build_tag: String, } @@ -744,7 +745,7 @@ LIMIT 100", let spec = &pspec.spec; let custom_ext = spec.custom_extensions.clone().unwrap_or(Vec::new()); info!("custom extensions: {:?}", &custom_ext); - let ext_remote_paths = extension_server::get_available_extensions( + let (ext_remote_paths, library_index) = extension_server::get_available_extensions( ext_remote_storage, &self.pgbin, &self.pgversion, @@ -754,38 +755,44 @@ LIMIT 100", .await?; self.ext_remote_paths .set(ext_remote_paths) - .expect("ext_remote_paths.set error"); + .expect("this is the only time we set ext_remote_paths"); + self.library_index + .set(library_index) + .expect("this is the only time we set library_index"); } Ok(()) } - pub async fn download_extension(&self, ext_name: &str) -> Result<()> { + pub async fn download_extension(&self, ext_name: &str, is_library: bool) -> Result<()> { match &self.ext_remote_storage { None => anyhow::bail!("No remote extension storage"), Some(remote_storage) => { - // TODO: eliminate useless LOAD Library calls to this function (if possible) - // not clear that we can distinguish between useful and useless - // library calls better than the below code - let ext_name = ext_name.replace(".so", ""); + let mut real_ext_name = ext_name.to_string(); + if is_library { + real_ext_name = real_ext_name.replace(".so", ""); + real_ext_name = + self.library_index.get().expect("oncelock err")[&real_ext_name].clone(); + } + { let mut already_downloaded_extensions = self.already_downloaded_extensions.lock().expect("bad lock"); - if already_downloaded_extensions.contains(&ext_name) { + if already_downloaded_extensions.contains(&real_ext_name) { info!( "extension {:?} already exists, skipping download", &ext_name ); return Ok(()); } else { - already_downloaded_extensions.insert(ext_name.clone()); + already_downloaded_extensions.insert(real_ext_name.clone()); } } extension_server::download_extension( - &ext_name, + &real_ext_name, &self .ext_remote_paths .get() - .expect("error accessing ext_remote_paths")[&ext_name], + .expect("error accessing ext_remote_paths")[&real_ext_name], remote_storage, &self.pgbin, ) @@ -838,7 +845,7 @@ LIMIT 100", info!("Downloading to shared preload libraries: {:?}", &libs_vec); let mut download_tasks = Vec::new(); for library in &libs_vec { - download_tasks.push(self.download_extension(library)); + download_tasks.push(self.download_extension(library, true)); } let results = join_all(download_tasks).await; for result in results { diff --git a/compute_tools/src/extension_server.rs b/compute_tools/src/extension_server.rs index 9129856ab6..e1e05e4ef6 100644 --- a/compute_tools/src/extension_server.rs +++ b/compute_tools/src/extension_server.rs @@ -59,7 +59,8 @@ More specifically, here is an example ext_index.json */ use anyhow::Context; use anyhow::{self, Result}; -use futures::future::join_all; +use futures::future::{join_all, Remote}; +use opentelemetry::InstrumentationLibrary; use remote_storage::*; use serde_json::{self, Value}; use std::collections::HashMap; @@ -109,9 +110,9 @@ pub async fn get_available_extensions( remote_storage: &GenericRemoteStorage, pgbin: &str, pg_version: &str, - enabled_extensions: &[String], + custom_extensions: &[String], build_tag: &str, -) -> Result> { +) -> Result<(HashMap, HashMap)> { let local_sharedir = Path::new(&get_pg_config("--sharedir", pgbin)).join("extension"); let index_path = format!("{build_tag}/{pg_version}/ext_index.json"); let index_path = RemotePath::new(Path::new(&index_path)).context("error forming path")?; @@ -123,16 +124,40 @@ 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("error parsing json"); + 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().context("error parsing json")?; + 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"); + 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); + } + + 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); + } + + let all_extension_data = ext_index_full["extension_data"] + .as_object() + .context("error parsing json")?; + 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 = ext_index_full[extension] + let ext_data = all_extension_data[extension] .as_object() .context("error parsing json")?; @@ -158,7 +183,7 @@ pub async fn get_available_extensions( for result in results { result?; } - Ok(ext_remote_paths) + Ok((ext_remote_paths, parsed_lib_index)) } // download the archive for a given extension, diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 8fd40cdfe1..af07412b52 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -125,13 +125,23 @@ async fn routes(req: Request, compute: &Arc) -> Response { info!("serving {:?} POST request", route); info!("req.uri {:?}", req.uri()); - let filename = route.split('/').last().unwrap().to_string(); - info!( - "serving /extension_server POST request, filename: {:?}", - &filename - ); - match compute.download_extension(&filename).await { + let mut is_library = false; + if let Some(params) = req.uri().query() { + info!("serving {:?} POST request with params: {}", route, params); + if params == "is_library=true" { + is_library = true; + } else { + let mut resp = Response::new(Body::from("Wrong request parameters")); + *resp.status_mut() = StatusCode::BAD_REQUEST; + return resp; + } + } + + let filename = route.split('/').last().unwrap().to_string(); + info!("serving /extension_server POST request, filename: {filename:?} is_library: {is_library}"); + + match compute.download_extension(&filename, is_library).await { Ok(_) => Response::new(Body::from("OK")), Err(e) => { error!("extension download failed: {}", e); diff --git a/test_runner/regress/data/extension_test/5648391853/v14/ext_index.json b/test_runner/regress/data/extension_test/5648391853/v14/ext_index.json index 3c8bf3ea7a..52d1d99b71 100644 --- a/test_runner/regress/data/extension_test/5648391853/v14/ext_index.json +++ b/test_runner/regress/data/extension_test/5648391853/v14/ext_index.json @@ -1 +1,23 @@ -{"anon": {"control_data": {"anon.control": "# PostgreSQL Anonymizer (anon) extension\ncomment = 'Data anonymization tools'\ndefault_version = '1.1.0'\ndirectory='extension/anon'\nrelocatable = false\nrequires = 'pgcrypto'\nsuperuser = false\nmodule_pathname = '$libdir/anon'\ntrusted = true\n"}, "archive_path": "5648391853/v14/extensions/anon.tar.zstd"}, "kq_imcx": {"control_data": {"kq_imcx.control": "# This file is generated content from add_postgresql_extension.\n# No point in modifying it, it will be overwritten anyway.\n\n# Default version, always set\ndefault_version = '0.1'\n\n# Module pathname generated from target shared library name. Use\n# MODULE_PATHNAME in script file.\nmodule_pathname = '$libdir/kq_imcx.so'\n\n# Comment for extension. Set using COMMENT option. Can be set in\n# script file as well.\ncomment = 'ketteQ In-Memory Calendar Extension (IMCX)'\n\n# Encoding for script file. Set using ENCODING option.\n#encoding = ''\n\n# Required extensions. Set using REQUIRES option (multi-valued).\n#requires = ''\ntrusted = true\n"}, "archive_path": "5648391853/v14/extensions/kq_imcx.tar.zstd"}} \ No newline at end of file +{ + "public_extensions": [ + "anon" + ], + "library_index": { + "anon": "5648391853/v14/extensions/anon.tar.zst", + "kq_imcx": "5648391853/v14/extensions/kq_imcx.tar.zst" + }, + "extension_data": { + "kq_imcx": { + "control_data": { + "kq_imcx.control": "# This file is generated content from add_postgresql_extension.\n# No point in modifying it, it will be overwritten anyway.\n\n# Default version, always set\ndefault_version = '0.1'\n\n# Module pathname generated from target shared library name. Use\n# MODULE_PATHNAME in script file.\nmodule_pathname = '$libdir/kq_imcx.so'\n\n# Comment for extension. Set using COMMENT option. Can be set in\n# script file as well.\ncomment = 'ketteQ In-Memory Calendar Extension (IMCX)'\n\n# Encoding for script file. Set using ENCODING option.\n#encoding = ''\n\n# Required extensions. Set using REQUIRES option (multi-valued).\n#requires = ''\ntrusted = true\n" + }, + "archive_path": "5648391853/v14/extensions/kq_imcx.tar.zstd" + }, + "anon": { + "control_data": { + "anon.control": "# PostgreSQL Anonymizer (anon) extension \ncomment = 'Data anonymization tools' \ndefault_version = '1.1.0' \ndirectory='extension/anon' \nrelocatable = false \nrequires = 'pgcrypto' \nsuperuser = false \nmodule_pathname = '$libdir/anon' \ntrusted = true \n" + }, + "archive_path": "5648391853/v14/extensions/anon.tar.zst" + } + } +} \ No newline at end of file diff --git a/test_runner/regress/data/extension_test/5648391853/v15/ext_index.json b/test_runner/regress/data/extension_test/5648391853/v15/ext_index.json index 2629384865..57c7a190b2 100644 --- a/test_runner/regress/data/extension_test/5648391853/v15/ext_index.json +++ b/test_runner/regress/data/extension_test/5648391853/v15/ext_index.json @@ -1 +1,23 @@ -{"kq_imcx": {"control_data": {"kq_imcx.control": "# This file is generated content from add_postgresql_extension.\n# No point in modifying it, it will be overwritten anyway.\n\n# Default version, always set\ndefault_version = '0.1'\n\n# Module pathname generated from target shared library name. Use\n# MODULE_PATHNAME in script file.\nmodule_pathname = '$libdir/kq_imcx.so'\n\n# Comment for extension. Set using COMMENT option. Can be set in\n# script file as well.\ncomment = 'ketteQ In-Memory Calendar Extension (IMCX)'\n\n# Encoding for script file. Set using ENCODING option.\n#encoding = ''\n\n# Required extensions. Set using REQUIRES option (multi-valued).\n#requires = ''\ntrusted = true\n"}, "archive_path": "5648391853/v15/extensions/kq_imcx.tar.zstd"}, "anon": {"control_data": {"anon.control": "# PostgreSQL Anonymizer (anon) extension\ncomment = 'Data anonymization tools'\ndefault_version = '1.1.0'\ndirectory='extension/anon'\nrelocatable = false\nrequires = 'pgcrypto'\nsuperuser = false\nmodule_pathname = '$libdir/anon'\ntrusted = true\n"}, "archive_path": "5648391853/v15/extensions/anon.tar.zstd"}} \ No newline at end of file +{ + "public_extensions": [ + "anon" + ], + "library_index": { + "anon": "anon", + "kq_imcx": "kq_imcx" + }, + "extension_data": { + "kq_imcx": { + "control_data": { + "kq_imcx.control": "# This file is generated content from add_postgresql_extension.\n# No point in modifying it, it will be overwritten anyway.\n\n# Default version, always set\ndefault_version = '0.1'\n\n# Module pathname generated from target shared library name. Use\n# MODULE_PATHNAME in script file.\nmodule_pathname = '$libdir/kq_imcx.so'\n\n# Comment for extension. Set using COMMENT option. Can be set in\n# script file as well.\ncomment = 'ketteQ In-Memory Calendar Extension (IMCX)'\n\n# Encoding for script file. Set using ENCODING option.\n#encoding = ''\n\n# Required extensions. Set using REQUIRES option (multi-valued).\n#requires = ''\ntrusted = true\n" + }, + "archive_path": "5648391853/v15/extensions/kq_imcx.tar.zstd" + }, + "anon": { + "control_data": { + "anon.control": "# PostgreSQL Anonymizer (anon) extension \ncomment = 'Data anonymization tools' \ndefault_version = '1.1.0' \ndirectory='extension/anon' \nrelocatable = false \nrequires = 'pgcrypto' \nsuperuser = false \nmodule_pathname = '$libdir/anon' \ntrusted = true \n" + }, + "archive_path": "5648391853/v15/extensions/anon.tar.zst" + } + } +} \ No newline at end of file